from nmigen import *

"""
Insn.py
Class for generic RISC-V instructions
"""

class Insn(Elaboratable):
    def __init__(self, RISCV_FORMAL_ILEN, RISCV_FORMAL_XLEN, RISCV_FORMAL_CSR_MISA, RISCV_FORMAL_COMPRESSED):
        # Core-specific constants
        self.RISCV_FORMAL_ILEN = RISCV_FORMAL_ILEN
        self.RISCV_FORMAL_XLEN = RISCV_FORMAL_XLEN
        self.RISCV_FORMAL_CSR_MISA = RISCV_FORMAL_CSR_MISA
        self.RISCV_FORMAL_COMPRESSED = RISCV_FORMAL_COMPRESSED

        # RVFI input ports
        self.rvfi_valid = Signal(1)
        self.rvfi_insn = Signal(self.RISCV_FORMAL_ILEN)
        self.rvfi_pc_rdata = Signal(self.RISCV_FORMAL_XLEN)
        self.rvfi_rs1_rdata = Signal(self.RISCV_FORMAL_XLEN)
        self.rvfi_rs2_rdata = Signal(self.RISCV_FORMAL_XLEN)
        self.rvfi_mem_rdata = Signal(self.RISCV_FORMAL_XLEN)
        if self.RISCV_FORMAL_CSR_MISA:
            self.rvfi_csr_misa_rdata = Signal(self.RISCV_FORMAL_XLEN)

        # RVFI output ports
        if self.RISCV_FORMAL_CSR_MISA:
            self.spec_csr_misa_rmask = Signal(self.RISCV_FORMAL_XLEN)
        self.spec_valid = Signal(1)
        self.spec_trap = Signal(1)
        self.spec_rs1_addr = Signal(5)
        self.spec_rs2_addr = Signal(5)
        self.spec_rd_addr = Signal(5)
        self.spec_rd_wdata = Signal(self.RISCV_FORMAL_XLEN)
        self.spec_pc_wdata = Signal(self.RISCV_FORMAL_XLEN)
        self.spec_mem_addr = Signal(self.RISCV_FORMAL_XLEN)
        self.spec_mem_rmask = Signal(int(self.RISCV_FORMAL_XLEN // 8))
        self.spec_mem_wmask = Signal(int(self.RISCV_FORMAL_XLEN // 8))
        self.spec_mem_wdata = Signal(self.RISCV_FORMAL_XLEN)

        # Additional wires and registers
        self.insn_padding = Signal(self.RISCV_FORMAL_ILEN)
        self.insn_imm = Signal(self.RISCV_FORMAL_XLEN)
        self.insn_funct7 = Signal(7)
        self.insn_funct6 = Signal(6)
        self.insn_shamt = Signal(6)
        self.insn_rs2 = Signal(5)
        self.insn_rs1 = Signal(5)
        self.insn_funct3 = Signal(3)
        self.insn_rd = Signal(5)
        self.insn_opcode = Signal(7)
        self.misa_ok = Signal(1)
        self.ialign16 = Signal(1)
    def ports(self):
        input_ports = [
            self.rvfi_valid,
            self.rvfi_insn,
            self.rvfi_pc_rdata,
            self.rvfi_rs1_rdata,
            self.rvfi_rs2_rdata,
            self.rvfi_mem_rdata
        ]
        if self.RISCV_FORMAL_CSR_MISA:
            input_ports.append(self.rvfi_csr_misa_rdata)
        output_ports = [
            self.spec_valid,
            self.spec_trap,
            self.spec_rs1_addr,
            self.spec_rs2_addr,
            self.spec_rd_addr,
            self.spec_rd_wdata,
            self.spec_pc_wdata,
            self.spec_mem_addr,
            self.spec_mem_rmask,
            self.spec_mem_wmask,
            self.spec_mem_wdata
        ]
        if self.RISCV_FORMAL_CSR_MISA:
            output_ports.append(self.spec_csr_misa_rmask)
        return input_ports + output_ports
    def elaborate(self, platform):
        m = Module()

        m.d.comb += self.insn_padding.eq(self.rvfi_insn >> 32)
        m.d.comb += self.insn_funct7.eq(self.rvfi_insn[25:32])
        m.d.comb += self.insn_funct6.eq(self.rvfi_insn[26:32])
        m.d.comb += self.insn_shamt.eq(self.rvfi_insn[20:26])
        m.d.comb += self.insn_rs2.eq(self.rvfi_insn[20:25])
        m.d.comb += self.insn_rs1.eq(self.rvfi_insn[15:20])
        m.d.comb += self.insn_funct3.eq(self.rvfi_insn[12:15])
        m.d.comb += self.insn_rd.eq(self.rvfi_insn[7:12])
        m.d.comb += self.insn_opcode.eq(self.rvfi_insn[:7])

        # default assignments
        m.d.comb += self.spec_valid.eq(0)
        m.d.comb += self.spec_trap.eq(~self.misa_ok)
        m.d.comb += self.spec_rs1_addr.eq(0)
        m.d.comb += self.spec_rs2_addr.eq(0)
        m.d.comb += self.spec_rd_addr.eq(0)
        m.d.comb += self.spec_rd_wdata.eq(0)
        m.d.comb += self.spec_pc_wdata.eq(0)
        m.d.comb += self.spec_mem_addr.eq(0)
        m.d.comb += self.spec_mem_rmask.eq(0)
        m.d.comb += self.spec_mem_wmask.eq(0)
        m.d.comb += self.spec_mem_wdata.eq(0)

        return m