from nmigen import *
from nmigen.asserts import *

"""
PC Backward Check
"""

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

        # 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_pc_rdata = Signal(self.params.xlen)
        self.rvfi_pc_wdata = Signal(self.params.xlen)

    def ports(self):
        input_ports = [
            self.reset,
            self.check,
            self.rvfi_valid,
            self.rvfi_order,
            self.rvfi_pc_rdata,
            self.rvfi_pc_wdata
        ]
        return input_ports

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

        insn_order = AnyConst(64)
        expect_pc = Signal(self.params.xlen)
        expect_pc_valid = Signal(1, reset=0)

        pc_wdata = Signal(self.params.xlen)
        m.d.comb += pc_wdata.eq(self.rvfi_pc_wdata)

        with m.If(self.reset):
            m.d.sync += expect_pc_valid.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(expect_pc_valid):
                    m.d.comb += Assert(self.rvformal_addr_eq(expect_pc, pc_wdata))
            with m.Else():
                with m.If(self.rvfi_valid & (self.rvfi_order == insn_order + 1)):
                    m.d.sync += expect_pc.eq(self.rvfi_pc_rdata)
                    m.d.sync += expect_pc_valid.eq(1)

        return m