from nmigen import *
from nmigen.hdl.ast import *

# Ugly hack for now to get the insn_add import working
import sys
sys.path.append('../insns')
from insn_add import *

class rvfi_insn_check(Elaboratable):
    def __init__(self, RISCV_FORMAL_ILEN=32, RISCV_FORMAL_XLEN=32):
        self.RISCV_FORMAL_ILEN = RISCV_FORMAL_ILEN
        self.RISCV_FORMAL_XLEN = RISCV_FORMAL_XLEN
        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)
    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
        ]
        output_ports = []
        return input_ports + output_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)

        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))

        # Change this submodule accordingly to test for different instructions(?)
        m.submodules.insn_spec = insn_spec = rvfi_insn_add(self.RISCV_FORMAL_ILEN, self.RISCV_FORMAL_XLEN)

        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)

        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))

        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((~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():
                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(spec_pc_wdata == pc_wdata)

                    with m.If(spec_mem_wmask | spec_mem_rmask):
                        m.d.comb += Assert(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