from nmigen import *
from nmigen.asserts import *

"""
Instruction Check
"""

class InsnCheck(Elaboratable):
    def __init__(self, RISCV_FORMAL_ILEN, RISCV_FORMAL_XLEN, RISCV_FORMAL_CSR_MISA, RISCV_FORMAL_COMPRESSED, RISCV_FORMAL_ALIGNED_MEM, insn_model, rvformal_addr_valid):
        # 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
        self.RISCV_FORMAL_ALIGNED_MEM = RISCV_FORMAL_ALIGNED_MEM

        # Instruction under test
        self.insn_model = insn_model

        # Address validity and equality
        self.rvformal_addr_valid = rvformal_addr_valid
        self.rvformal_addr_eq = lambda a, b: (self.rvformal_addr_valid(a) == self.rvformal_addr_valid(b)) & ((~self.rvformal_addr_valid(a)) | (a == b))

        # Input ports
        self.reset = Signal(1)
        self.check = Signal(1)
        self.rvfi_valid = Signal(1)
        self.rvfi_order = Signal(64)
        self.rvfi_insn = Signal(self.RISCV_FORMAL_ILEN)
        self.rvfi_trap = Signal(1)
        self.rvfi_halt = Signal(1)
        self.rvfi_intr = Signal(1)
        self.rvfi_mode = Signal(2)
        self.rvfi_ixl = Signal(2)
        self.rvfi_rs1_addr = Signal(5)
        self.rvfi_rs2_addr = Signal(5)
        self.rvfi_rs1_rdata = Signal(self.RISCV_FORMAL_XLEN)
        self.rvfi_rs2_rdata = Signal(self.RISCV_FORMAL_XLEN)
        self.rvfi_rd_addr = Signal(5)
        self.rvfi_rd_wdata = Signal(self.RISCV_FORMAL_XLEN)
        self.rvfi_pc_rdata = Signal(self.RISCV_FORMAL_XLEN)
        self.rvfi_pc_wdata = Signal(self.RISCV_FORMAL_XLEN)
        self.rvfi_mem_addr = Signal(self.RISCV_FORMAL_XLEN)
        self.rvfi_mem_rmask = Signal(int(self.RISCV_FORMAL_XLEN // 8))
        self.rvfi_mem_wmask = Signal(int(self.RISCV_FORMAL_XLEN // 8))
        self.rvfi_mem_rdata = Signal(self.RISCV_FORMAL_XLEN)
        self.rvfi_mem_wdata = Signal(self.RISCV_FORMAL_XLEN)
        if self.RISCV_FORMAL_CSR_MISA:
            self.rvfi_csr_misa_rmask = Signal(self.RISCV_FORMAL_XLEN)
            self.rvfi_csr_misa_wmask = Signal(self.RISCV_FORMAL_XLEN)
            self.rvfi_csr_misa_rdata = Signal(self.RISCV_FORMAL_XLEN)
            self.rvfi_csr_misa_wdata = Signal(self.RISCV_FORMAL_XLEN)
    def ports(self):
        input_ports = [
            self.reset,
            self.check,
            self.rvfi_valid,
            self.rvfi_order,
            self.rvfi_insn,
            self.rvfi_trap,
            self.rvfi_halt,
            self.rvfi_intr,
            self.rvfi_mode,
            self.rvfi_ixl,
            self.rvfi_rs1_addr,
            self.rvfi_rs2_addr,
            self.rvfi_rs1_rdata,
            self.rvfi_rs2_rdata,
            self.rvfi_rd_addr,
            self.rvfi_rd_wdata,
            self.rvfi_pc_rdata,
            self.rvfi_pc_wdata,
            self.rvfi_mem_addr,
            self.rvfi_mem_rmask,
            self.rvfi_mem_wmask,
            self.rvfi_mem_rdata,
            self.rvfi_mem_wdata
        ]
        if self.RISCV_FORMAL_CSR_MISA:
            input_ports.extend([
                self.rvfi_csr_misa_rmask,
                self.rvfi_csr_misa_wmask,
                self.rvfi_csr_misa_rdata,
                self.rvfi_csr_misa_wdata
            ])
        return input_ports
    def elaborate(self, platform):
        m = Module()

        valid = Signal(1)
        m.d.comb += valid.eq((~self.reset) & self.rvfi_valid)
        insn = Signal(self.RISCV_FORMAL_ILEN)
        m.d.comb += insn.eq(self.rvfi_insn)
        trap = Signal(1)
        m.d.comb += trap.eq(self.rvfi_trap)
        halt = Signal(1)
        m.d.comb += halt.eq(self.rvfi_halt)
        intr = Signal(1)
        m.d.comb += intr.eq(self.rvfi_intr)
        rs1_addr = Signal(5)
        m.d.comb += rs1_addr.eq(self.rvfi_rs1_addr)
        rs2_addr = Signal(5)
        m.d.comb += rs2_addr.eq(self.rvfi_rs2_addr)
        rs1_rdata = Signal(self.RISCV_FORMAL_XLEN)
        m.d.comb += rs1_rdata.eq(self.rvfi_rs1_rdata)
        rs2_rdata = Signal(self.RISCV_FORMAL_XLEN)
        m.d.comb += rs2_rdata.eq(self.rvfi_rs2_rdata)
        rd_addr = Signal(5)
        m.d.comb += rd_addr.eq(self.rvfi_rd_addr)
        rd_wdata = Signal(self.RISCV_FORMAL_XLEN)
        m.d.comb += rd_wdata.eq(self.rvfi_rd_wdata)
        pc_rdata = Signal(self.RISCV_FORMAL_XLEN)
        m.d.comb += pc_rdata.eq(self.rvfi_pc_rdata)
        pc_wdata = Signal(self.RISCV_FORMAL_XLEN)
        m.d.comb += pc_wdata.eq(self.rvfi_pc_wdata)

        mem_addr = Signal(self.RISCV_FORMAL_XLEN)
        m.d.comb += mem_addr.eq(self.rvfi_mem_addr)
        mem_rmask = Signal(int(self.RISCV_FORMAL_XLEN // 8))
        m.d.comb += mem_rmask.eq(self.rvfi_mem_rmask)
        mem_wmask = Signal(int(self.RISCV_FORMAL_XLEN // 8))
        m.d.comb += mem_wmask.eq(self.rvfi_mem_wmask)
        mem_rdata = Signal(self.RISCV_FORMAL_XLEN)
        m.d.comb += mem_rdata.eq(self.rvfi_mem_rdata)
        mem_wdata = Signal(self.RISCV_FORMAL_XLEN)
        m.d.comb += mem_wdata.eq(self.rvfi_mem_wdata)

        if self.RISCV_FORMAL_CSR_MISA:
            csr_misa_rdata = Signal(self.RISCV_FORMAL_XLEN)
            m.d.comb += csr_misa_rdata.eq(self.rvfi_csr_misa_rdata)
            csr_misa_rmask = Signal(self.RISCV_FORMAL_XLEN)
            m.d.comb += csr_misa_rmask.eq(self.rvfi_csr_misa_rmask)
            spec_csr_misa_rmask = Signal(self.RISCV_FORMAL_XLEN)

        spec_valid = Signal(1)
        spec_trap = Signal(1)
        spec_rs1_addr = Signal(5)
        spec_rs2_addr = Signal(5)
        spec_rd_addr = Signal(5)
        spec_rd_wdata = Signal(self.RISCV_FORMAL_XLEN)
        spec_pc_wdata = Signal(self.RISCV_FORMAL_XLEN)
        spec_mem_addr = Signal(self.RISCV_FORMAL_XLEN)
        spec_mem_rmask = Signal(int(self.RISCV_FORMAL_XLEN // 8))
        spec_mem_wmask = Signal(int(self.RISCV_FORMAL_XLEN // 8))
        spec_mem_wdata = Signal(self.RISCV_FORMAL_XLEN)

        rs1_rdata_or_zero = Signal(self.RISCV_FORMAL_XLEN)
        m.d.comb += rs1_rdata_or_zero.eq(Mux(spec_rs1_addr != 0, rs1_rdata, 0))
        rs2_rdata_or_zero = Signal(self.RISCV_FORMAL_XLEN)
        m.d.comb += rs2_rdata_or_zero.eq(Mux(spec_rs2_addr != 0, rs2_rdata, 0))

        try:
            m.submodules.insn_spec = insn_spec = self.insn_model(RISCV_FORMAL_ILEN=self.RISCV_FORMAL_ILEN, RISCV_FORMAL_XLEN=self.RISCV_FORMAL_XLEN, RISCV_FORMAL_CSR_MISA=self.RISCV_FORMAL_CSR_MISA)
        except:
            try:
                m.submodules.insn_spec = insn_spec = self.insn_model(RISCV_FORMAL_ILEN=self.RISCV_FORMAL_ILEN, RISCV_FORMAL_XLEN=self.RISCV_FORMAL_XLEN, RISCV_FORMAL_CSR_MISA=self.RISCV_FORMAL_CSR_MISA, RISCV_FORMAL_COMPRESSED=self.RISCV_FORMAL_COMPRESSED)
            except:
                m.submodules.insn_spec = insn_spec = self.insn_model(RISCV_FORMAL_ILEN=self.RISCV_FORMAL_ILEN, RISCV_FORMAL_XLEN=self.RISCV_FORMAL_XLEN, RISCV_FORMAL_CSR_MISA=self.RISCV_FORMAL_CSR_MISA, RISCV_FORMAL_ALIGNED_MEM=self.RISCV_FORMAL_ALIGNED_MEM)

        m.d.comb += insn_spec.rvfi_valid.eq(valid)
        m.d.comb += insn_spec.rvfi_insn.eq(insn)
        m.d.comb += insn_spec.rvfi_pc_rdata.eq(pc_rdata)
        m.d.comb += insn_spec.rvfi_rs1_rdata.eq(rs1_rdata_or_zero)
        m.d.comb += insn_spec.rvfi_rs2_rdata.eq(rs2_rdata_or_zero)
        m.d.comb += insn_spec.rvfi_mem_rdata.eq(mem_rdata)

        if self.RISCV_FORMAL_CSR_MISA:
            m.d.comb += insn_spec.rvfi_csr_misa_rdata.eq(csr_misa_rdata)
            m.d.comb += spec_csr_misa_rmask.eq(insn_spec.spec_csr_misa_rmask)

        m.d.comb += spec_valid.eq(insn_spec.spec_valid)
        m.d.comb += spec_trap.eq(insn_spec.spec_trap)
        m.d.comb += spec_rs1_addr.eq(insn_spec.spec_rs1_addr)
        m.d.comb += spec_rs2_addr.eq(insn_spec.spec_rs2_addr)
        m.d.comb += spec_rd_addr.eq(insn_spec.spec_rd_addr)
        m.d.comb += spec_rd_wdata.eq(insn_spec.spec_rd_wdata)
        m.d.comb += spec_pc_wdata.eq(insn_spec.spec_pc_wdata)
        m.d.comb += spec_mem_addr.eq(insn_spec.spec_mem_addr)
        m.d.comb += spec_mem_rmask.eq(insn_spec.spec_mem_rmask)
        m.d.comb += spec_mem_wmask.eq(insn_spec.spec_mem_wmask)
        m.d.comb += spec_mem_wdata.eq(insn_spec.spec_mem_wdata)

        insn_pma_x = Signal(1)
        mem_pma_r = Signal(1)
        mem_pma_w = Signal(1)

        mem_log2len = Signal(2)
        m.d.comb += mem_log2len.eq(Mux((spec_mem_rmask | spec_mem_wmask) & 0b11110000, 3, Mux((spec_mem_rmask | spec_mem_wmask) & 0b00001100, 2, Mux((spec_mem_rmask | spec_mem_wmask) & 0b00000010, 1, 0))))

        m.d.comb += insn_pma_x.eq(1)
        m.d.comb += mem_pma_r.eq(1)
        m.d.comb += mem_pma_w.eq(1)

        mem_access_fault = Signal(1)
        m.d.comb += mem_access_fault.eq((spec_mem_rmask & ~mem_pma_r) | (spec_mem_wmask & ~mem_pma_w) | ((spec_mem_rmask | spec_mem_wmask) & ~self.rvformal_addr_valid(spec_mem_addr)))

        with m.If(~self.reset):
            m.d.comb += Cover(spec_valid)
            m.d.comb += Cover(spec_valid & ~trap)
            m.d.comb += Cover(self.check & spec_valid)
            m.d.comb += Cover(self.check & spec_valid & ~trap)
        with m.If((~self.reset) & self.check):
            m.d.comb += Assume(spec_valid)

            with m.If((~self.rvformal_addr_valid(pc_rdata)) | (~insn_pma_x) | mem_access_fault):
                m.d.comb += Assert(trap)
                m.d.comb += Assert(rd_addr == 0)
                m.d.comb += Assert(rd_wdata == 0)
                m.d.comb += Assert(mem_wmask == 0)
            with m.Else():
                if self.RISCV_FORMAL_CSR_MISA:
                    m.d.comb += Assert((spec_csr_misa_rmask & csr_misa_rmask) == spec_csr_misa_rmask)

                with m.If(rs1_addr == 0):
                    m.d.comb += Assert(rs1_rdata == 0)
                
                with m.If(rs2_addr == 0):
                    m.d.comb += Assert(rs2_rdata == 0)

                with m.If(~spec_trap):
                    with m.If(spec_rs1_addr != 0):
                        m.d.comb += Assert(spec_rs1_addr == rs1_addr)

                    with m.If(spec_rs2_addr != 0):
                        m.d.comb += Assert(spec_rs2_addr == rs2_addr)

                    m.d.comb += Assert(spec_rd_addr == rd_addr)
                    m.d.comb += Assert(spec_rd_wdata == rd_wdata)
                    m.d.comb += Assert(self.rvformal_addr_eq(spec_pc_wdata, pc_wdata))

                    with m.If(spec_mem_wmask | spec_mem_rmask):
                        m.d.comb += Assert(self.rvformal_addr_eq(spec_mem_addr, mem_addr))

                    for i in range(int(self.RISCV_FORMAL_XLEN // 8)):
                        with m.If(spec_mem_wmask[i]):
                            m.d.comb += Assert(mem_wmask[i])
                            m.d.comb += Assert(spec_mem_wdata[i*8:i*8+8] == mem_wdata[i*8:i*8+8])
                        with m.Elif(mem_wmask[i]):
                            m.d.comb += Assert(mem_rmask[i])
                            m.d.comb += Assert(mem_rdata[i*8:i*8+8] == mem_wdata[i*8:i*8+8])
                        with m.If(spec_mem_rmask[i]):
                            m.d.comb += Assert(mem_rmask[i])

                m.d.comb += Assert(spec_trap == trap)

        return m