from nmigen import *
from nmigen.asserts import *

"""
Register Check
"""

class RegCheck(Elaboratable):
    def __init__(self, params):
        # Core-specific parameters
        self.params = params

        # Input ports
        self.reset = Signal(1)
        self.check = Signal(1)
        self.rvfi_valid = Signal(1)
        self.rvfi_order = Signal(64)
        self.rvfi_rs1_addr = Signal(5)
        self.rvfi_rs1_rdata = Signal(self.params.xlen)
        self.rvfi_rs2_addr = Signal(5)
        self.rvfi_rs2_rdata = Signal(self.params.xlen)
        self.rvfi_rd_addr = Signal(5)
        self.rvfi_rd_wdata = Signal(self.params.xlen)

    def ports(self):
        input_ports = [
            self.reset,
            self.check,
            self.rvfi_valid,
            self.rvfi_order,
            self.rvfi_rs1_addr,
            self.rvfi_rs1_rdata,
            self.rvfi_rs2_addr,
            self.rvfi_rs2_rdata,
            self.rvfi_rd_addr,
            self.rvfi_rd_wdata
        ]
        return input_ports

    def elaborate(self, platform):
        m = Module()

        insn_order = AnyConst(64)
        register_index = AnyConst(5)
        register_shadow = Signal(self.params.xlen, reset=0)
        register_written = Signal(1, reset=0)

        with m.If(self.reset):
            m.d.sync += register_shadow.eq(0)
            m.d.sync += register_written.eq(0)
        with m.Else():
            with m.If(self.check):
                m.d.comb += Assume(self.rvfi_valid)
                m.d.comb += Assume(insn_order == self.rvfi_order)
                with m.If(register_written & (register_index == self.rvfi_rs1_addr)):
                    m.d.comb += Assert(register_shadow == self.rvfi_rs1_rdata)
                with m.If(register_written & (register_index == self.rvfi_rs2_addr)):
                    m.d.comb += Assert(register_shadow == self.rvfi_rs2_rdata)
            with m.Else():
                with m.If(self.rvfi_valid & (self.rvfi_order < insn_order) & (register_index == self.rvfi_rd_addr)):
                    m.d.sync += register_shadow.eq(self.rvfi_rd_wdata)
                    m.d.sync += register_written.eq(1)

        return m