from amaranth import *
from amaranth.asserts import *

from .. import PowerFVCheck
from ... import pfv, tb


__all__ = ["GPRSpec", "GPRCheck"]


class GPRSpec(Elaboratable):
    """GPR consistency specification.

    Checks that:
      * reads from a GPR return the last value that was written to it;
      * writes to multiple GPRs by a single instruction (e.g. load with update) do not target the
        same register.
    """
    def __init__(self, post):
        self.pfv  = pfv.Interface()
        self.post = tb.Trigger(cycle=post)

    def triggers(self):
        yield self.post

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

        spec_order     = AnyConst(self.pfv.order.width)
        spec_gpr_index = AnyConst(5)

        gpr_ra_read    = Signal()
        gpr_ra_write   = Signal()
        gpr_rb_read    = Signal()
        gpr_rs_read    = Signal()
        gpr_rt_read    = Signal()
        gpr_rt_write   = Signal()

        gpr_conflict   = Signal()
        gpr_written    = Signal()
        gpr_shadow     = Signal(64)

        m.d.comb += [
            gpr_ra_read .eq(self.pfv.ra.r_stb & (self.pfv.ra.index == spec_gpr_index)),
            gpr_ra_write.eq(self.pfv.ra.w_stb & (self.pfv.ra.index == spec_gpr_index)),
            gpr_rb_read .eq(self.pfv.rb.r_stb & (self.pfv.rb.index == spec_gpr_index)),
            gpr_rs_read .eq(self.pfv.rs.r_stb & (self.pfv.rs.index == spec_gpr_index)),
            gpr_rt_read .eq(self.pfv.rt.r_stb & (self.pfv.rt.index == spec_gpr_index)),
            gpr_rt_write.eq(self.pfv.rt.w_stb & (self.pfv.rt.index == spec_gpr_index)),
        ]

        with m.If(self.pfv.stb & ~self.pfv.intr & (self.pfv.order <= spec_order)):
            with m.If(gpr_ra_write & gpr_rt_write):
                m.d.sync += gpr_conflict.eq(1)
            with m.If(gpr_ra_write | gpr_rt_write):
                m.d.sync += gpr_written.eq(1)
                with m.If(gpr_ra_write):
                    m.d.sync += gpr_shadow.eq(self.pfv.ra.w_data)
                with m.Else():
                    m.d.sync += gpr_shadow.eq(self.pfv.rt.w_data)

        with m.If(self.post.stb):
            m.d.sync += [
                Assume(Past(self.pfv.stb)),
                Assume(Past(self.pfv.order) == spec_order),
            ]
            with m.If(~Past(self.pfv.intr)):
                m.d.sync += Assert(gpr_written.implies(~gpr_conflict))
                with m.If(gpr_written & Past(gpr_ra_read)):
                    m.d.sync += Assert(Past(gpr_shadow) == Past(self.pfv.ra.r_data))
                with m.If(gpr_written & Past(gpr_rb_read)):
                    m.d.sync += Assert(Past(gpr_shadow) == Past(self.pfv.rb.r_data))
                with m.If(gpr_written & Past(gpr_rs_read)):
                    m.d.sync += Assert(Past(gpr_shadow) == Past(self.pfv.rs.r_data))
                with m.If(gpr_written & Past(gpr_rt_read)):
                    m.d.sync += Assert(Past(gpr_shadow) == Past(self.pfv.rt.r_data))

        return m


class GPRCheck(PowerFVCheck, name="cons_gpr"):
    def get_testbench(self, dut, post):
        tb_spec = GPRSpec(post)
        tb_top  = tb.Testbench(tb_spec, dut)
        return tb_top