from nmigen.test.utils import *
from ...checks.insn_check import *
from ...checks.pc_fwd_check import *
from ...checks.pc_bwd_check import *
from ...checks.reg_check import *
from ...checks.causal_check import *
from ...checks.liveness_check import *
from ...checks.unique_check import *
from minerva.core import *
from ...insns.insn_lui import *
from ...insns.insn_auipc import *
from ...insns.insn_jal import *
from ...insns.insn_jalr import *
from ...insns.insn_beq import *
from ...insns.insn_bne import *
from ...insns.insn_blt import *
from ...insns.insn_bge import *
from ...insns.insn_bltu import *
from ...insns.insn_bgeu import *
from ...insns.insn_lb import *
from ...insns.insn_lh import *
from ...insns.insn_lw import *
from ...insns.insn_lbu import *
from ...insns.insn_lhu import *
from ...insns.insn_sb import *
from ...insns.insn_sh import *
from ...insns.insn_sw import *
from ...insns.insn_addi import *
from ...insns.insn_slti import *
from ...insns.insn_sltiu import *
from ...insns.insn_xori import *
from ...insns.insn_ori import *
from ...insns.insn_andi import *
from ...insns.insn_slli import *
from ...insns.insn_srli import *
from ...insns.insn_srai import *
from ...insns.insn_add import *
from ...insns.insn_sub import *
from ...insns.insn_sll import *
from ...insns.insn_slt import *
from ...insns.insn_sltu import *
from ...insns.insn_xor import *
from ...insns.insn_srl import *
from ...insns.insn_sra import *
from ...insns.insn_or import *
from ...insns.insn_and import *
from ...insns.insn_mul import *
from ...insns.insn_mulh import *
from ...insns.insn_mulhsu import *
from ...insns.insn_mulhu import *
from ...insns.insn_div import *
from ...insns.insn_divu import *
from ...insns.insn_rem import *
from ...insns.insn_remu import *
from collections import namedtuple
from multiprocessing import Process

RISCVFormalParameters = namedtuple('RISCVFormalParameters',
    ['ilen', 'xlen', 'csr_misa', 'compressed', 'aligned_mem', 'altops'])

class InsnSpec(Elaboratable):
    def __init__(self, insn_model):
        self.insn_model = insn_model

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

        m.submodules.cpu = cpu = Minerva(
            with_rvfi=True,
            with_muldiv=True,
            with_icache=True,
            icache_nways=2,
            icache_nlines=2,
            icache_nwords=4,
            icache_base=0x1000,
            icache_limit=0x4000,
            with_dcache=True,
            dcache_nways=2,
            dcache_nlines=2,
            dcache_nwords=4,
            dcache_base=0x1000,
            dcache_limit=0x4000)
        m.submodules.insn_spec = insn_spec = InsnCheck(
            params=RISCVFormalParameters(32, 32, False, False, True, True),
            insn_model=self.insn_model,
            rvformal_addr_valid=lambda x:Const(1))

        # Wire input ports to Minerva core
        m.d.comb += cpu.external_interrupt.eq(0)
        m.d.comb += cpu.timer_interrupt.eq(0)
        m.d.comb += cpu.software_interrupt.eq(0)
        m.d.comb += cpu.ibus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.ibus.ack.eq(AnySeq(1))
        m.d.comb += cpu.ibus.err.eq(0)
        m.d.comb += cpu.dbus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.dbus.ack.eq(AnySeq(1))
        m.d.comb += cpu.dbus.err.eq(0)

        m.d.comb += insn_spec.reset.eq(AnySeq(1))
        m.d.comb += insn_spec.check.eq(AnySeq(1))

        m.d.comb += insn_spec.rvfi_valid.eq(cpu.rvfi.valid)
        m.d.comb += insn_spec.rvfi_order.eq(cpu.rvfi.order)
        m.d.comb += insn_spec.rvfi_insn.eq(cpu.rvfi.insn)
        m.d.comb += insn_spec.rvfi_trap.eq(cpu.rvfi.trap)
        m.d.comb += insn_spec.rvfi_halt.eq(cpu.rvfi.halt)
        m.d.comb += insn_spec.rvfi_intr.eq(cpu.rvfi.intr)
        m.d.comb += insn_spec.rvfi_mode.eq(cpu.rvfi.mode)
        m.d.comb += insn_spec.rvfi_ixl.eq(cpu.rvfi.ixl)
        m.d.comb += insn_spec.rvfi_rs1_addr.eq(cpu.rvfi.rs1_addr)
        m.d.comb += insn_spec.rvfi_rs2_addr.eq(cpu.rvfi.rs2_addr)
        m.d.comb += insn_spec.rvfi_rs1_rdata.eq(cpu.rvfi.rs1_rdata)
        m.d.comb += insn_spec.rvfi_rs2_rdata.eq(cpu.rvfi.rs2_rdata)
        m.d.comb += insn_spec.rvfi_rd_addr.eq(cpu.rvfi.rd_addr)
        m.d.comb += insn_spec.rvfi_rd_wdata.eq(cpu.rvfi.rd_wdata)
        m.d.comb += insn_spec.rvfi_pc_rdata.eq(cpu.rvfi.pc_rdata)
        m.d.comb += insn_spec.rvfi_pc_wdata.eq(cpu.rvfi.pc_wdata)
        m.d.comb += insn_spec.rvfi_mem_addr.eq(cpu.rvfi.mem_addr)
        m.d.comb += insn_spec.rvfi_mem_rmask.eq(cpu.rvfi.mem_rmask)
        m.d.comb += insn_spec.rvfi_mem_wmask.eq(cpu.rvfi.mem_wmask)
        m.d.comb += insn_spec.rvfi_mem_rdata.eq(cpu.rvfi.mem_rdata)
        m.d.comb += insn_spec.rvfi_mem_wdata.eq(cpu.rvfi.mem_wdata)

        return m

class InsnTestCase(FHDLTestCase):
    def verify(self):
        def verify_insn(insn_spec, spec_name):
            self.assertFormal(InsnSpec(insn_spec), mode="cover", depth=20, spec_name=spec_name)
            self.assertFormal(InsnSpec(insn_spec), mode="bmc", depth=20, spec_name=spec_name)
            print("%s PASS" % spec_name)
        insns = [
            (InsnLui, "verify_lui"),
            (InsnAuipc, "verify_auipc"),
            (InsnJal, "verify_jal"),
            (InsnJalr, "verify_jalr"),
            (InsnBeq, "verify_beq"),
            (InsnBne, "verify_bne"),
            (InsnBlt, "verify_blt"),
            (InsnBge, "verify_bge"),
            (InsnBltu, "verify_bltu"),
            (InsnBgeu, "verify_bgeu"),
            (InsnLb, "verify_lb"),
            (InsnLh, "verify_lh"),
            (InsnLw, "verify_lw"),
            (InsnLbu, "verify_lbu"),
            (InsnLhu, "verify_lhu"),
            (InsnSb, "verify_sb"),
            (InsnSh, "verify_sh"),
            (InsnSw, "verify_sw"),
            (InsnAddi, "verify_addi"),
            (InsnSlti, "verify_slti"),
            (InsnSltiu, "verify_sltiu"),
            (InsnXori, "verify_xori"),
            (InsnOri, "verify_ori"),
            (InsnAndi, "verify_andi"),
            (InsnSlli, "verify_slli"),
            (InsnSrli, "verify_srli"),
            (InsnSrai, "verify_srai"),
            (InsnAdd, "verify_add"),
            (InsnSub, "verify_sub"),
            (InsnSll, "verify_sll"),
            (InsnSlt, "verify_slt"),
            (InsnSltu, "verify_sltu"),
            (InsnXor, "verify_xor"),
            (InsnSrl, "verify_srl"),
            (InsnSra, "verify_sra"),
            (InsnOr, "verify_or"),
            (InsnAnd, "verify_and"),
            (InsnMul, "verify_mul"),
            (InsnMulh, "verify_mulh"),
            (InsnMulhsu, "verify_mulhsu"),
            (InsnMulhu, "verify_mulhu"),
            (InsnDiv, "verify_div"),
            (InsnDivu, "verify_divu"),
            (InsnRem, "verify_rem"),
            (InsnRemu, "verify_remu")
        ]
        ps = []
        for insn_spec, spec_name in insns:
            p = Process(target=verify_insn, args=(insn_spec,spec_name))
            ps.append(p)
            p.start()
        for p in ps:
            p.join()

class PcFwdSpec(Elaboratable):
    def elaborate(self, platform):
        m = Module()

        m.submodules.cpu = cpu = Minerva(
            with_rvfi=True,
            with_muldiv=True,
            with_icache=True,
            icache_nways=2,
            icache_nlines=2,
            icache_nwords=4,
            icache_base=0x1000,
            icache_limit=0x4000,
            with_dcache=True,
            dcache_nways=2,
            dcache_nlines=2,
            dcache_nwords=4,
            dcache_base=0x1000,
            dcache_limit=0x4000)
        m.submodules.pc_fwd_spec = pc_fwd_spec = PcFwdCheck(
            params=RISCVFormalParameters(32, 32, False, False, True, True),
            rvformal_addr_valid=lambda x:Const(1))

        # Wire input ports to Minerva core
        m.d.comb += cpu.external_interrupt.eq(0)
        m.d.comb += cpu.timer_interrupt.eq(0)
        m.d.comb += cpu.software_interrupt.eq(0)
        m.d.comb += cpu.ibus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.ibus.ack.eq(AnySeq(1))
        m.d.comb += cpu.ibus.err.eq(0)
        m.d.comb += cpu.dbus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.dbus.ack.eq(AnySeq(1))
        m.d.comb += cpu.dbus.err.eq(0)

        m.d.comb += pc_fwd_spec.reset.eq(AnySeq(1))
        m.d.comb += pc_fwd_spec.check.eq(AnySeq(1))

        m.d.comb += pc_fwd_spec.rvfi_valid.eq(cpu.rvfi.valid)
        m.d.comb += pc_fwd_spec.rvfi_order.eq(cpu.rvfi.order)
        m.d.comb += pc_fwd_spec.rvfi_pc_rdata.eq(cpu.rvfi.pc_rdata)
        m.d.comb += pc_fwd_spec.rvfi_pc_wdata.eq(cpu.rvfi.pc_wdata)

        return m

class PcFwdTestCase(FHDLTestCase):
    def verify(self):
        self.assertFormal(PcFwdSpec(), mode="cover", depth=20, spec_name="verify_pc_fwd")
        self.assertFormal(PcFwdSpec(), mode="bmc", depth=20, spec_name="verify_pc_fwd")
        print("verify_pc_fwd PASS")

class PcBwdSpec(Elaboratable):
    def elaborate(self, platform):
        m = Module()

        m.submodules.cpu = cpu = Minerva(
            with_rvfi=True,
            with_muldiv=True,
            with_icache=True,
            icache_nways=2,
            icache_nlines=2,
            icache_nwords=4,
            icache_base=0x1000,
            icache_limit=0x4000,
            with_dcache=True,
            dcache_nways=2,
            dcache_nlines=2,
            dcache_nwords=4,
            dcache_base=0x1000,
            dcache_limit=0x4000)
        m.submodules.pc_bwd_spec = pc_bwd_spec = PcBwdCheck(
            params=RISCVFormalParameters(32, 32, False, False, True, True),
            rvformal_addr_valid=lambda x:Const(1))

        # Wire input ports to Minerva core
        m.d.comb += cpu.external_interrupt.eq(0)
        m.d.comb += cpu.timer_interrupt.eq(0)
        m.d.comb += cpu.software_interrupt.eq(0)
        m.d.comb += cpu.ibus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.ibus.ack.eq(AnySeq(1))
        m.d.comb += cpu.ibus.err.eq(0)
        m.d.comb += cpu.dbus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.dbus.ack.eq(AnySeq(1))
        m.d.comb += cpu.dbus.err.eq(0)

        m.d.comb += pc_bwd_spec.reset.eq(AnySeq(1))
        m.d.comb += pc_bwd_spec.check.eq(AnySeq(1))

        m.d.comb += pc_bwd_spec.rvfi_valid.eq(cpu.rvfi.valid)
        m.d.comb += pc_bwd_spec.rvfi_order.eq(cpu.rvfi.order)
        m.d.comb += pc_bwd_spec.rvfi_pc_rdata.eq(cpu.rvfi.pc_rdata)
        m.d.comb += pc_bwd_spec.rvfi_pc_wdata.eq(cpu.rvfi.pc_wdata)

        return m

class PcBwdTestCase(FHDLTestCase):
    def verify(self):
        self.assertFormal(PcBwdSpec(), mode="cover", depth=20, spec_name="verify_pc_bwd")
        self.assertFormal(PcBwdSpec(), mode="bmc", depth=20, spec_name="verify_pc_bwd")
        print("verify_pc_bwd PASS")

class RegSpec(Elaboratable):
    def elaborate(self, platform):
        m = Module()

        m.submodules.cpu = cpu = Minerva(
            with_rvfi=True,
            with_muldiv=True,
            with_icache=True,
            icache_nways=2,
            icache_nlines=2,
            icache_nwords=4,
            icache_base=0x1000,
            icache_limit=0x4000,
            with_dcache=True,
            dcache_nways=2,
            dcache_nlines=2,
            dcache_nwords=4,
            dcache_base=0x1000,
            dcache_limit=0x4000)
        m.submodules.reg_spec = reg_spec = RegCheck(params=RISCVFormalParameters(32, 32, False, False, True, True))

        # Wire input ports to Minerva core
        m.d.comb += cpu.external_interrupt.eq(0)
        m.d.comb += cpu.timer_interrupt.eq(0)
        m.d.comb += cpu.software_interrupt.eq(0)
        m.d.comb += cpu.ibus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.ibus.ack.eq(AnySeq(1))
        m.d.comb += cpu.ibus.err.eq(0)
        m.d.comb += cpu.dbus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.dbus.ack.eq(AnySeq(1))
        m.d.comb += cpu.dbus.err.eq(0)

        m.d.comb += reg_spec.reset.eq(AnySeq(1))
        m.d.comb += reg_spec.check.eq(AnySeq(1))

        m.d.comb += reg_spec.rvfi_valid.eq(cpu.rvfi.valid)
        m.d.comb += reg_spec.rvfi_order.eq(cpu.rvfi.order)
        m.d.comb += reg_spec.rvfi_rs1_addr.eq(cpu.rvfi.rs1_addr)
        m.d.comb += reg_spec.rvfi_rs1_rdata.eq(cpu.rvfi.rs1_rdata)
        m.d.comb += reg_spec.rvfi_rs2_addr.eq(cpu.rvfi.rs2_addr)
        m.d.comb += reg_spec.rvfi_rs2_rdata.eq(cpu.rvfi.rs2_rdata)
        m.d.comb += reg_spec.rvfi_rd_addr.eq(cpu.rvfi.rd_addr)
        m.d.comb += reg_spec.rvfi_rd_wdata.eq(cpu.rvfi.rd_wdata)

        return m

class RegTestCase(FHDLTestCase):
    def verify(self):
        self.assertFormal(RegSpec(), mode="cover", depth=10, spec_name="verify_reg")
        self.assertFormal(RegSpec(), mode="bmc", depth=10, spec_name="verify_reg")
        print("verify_reg PASS")

class CausalSpec(Elaboratable):
    def elaborate(self, platform):
        m = Module()

        m.submodules.cpu = cpu = Minerva(
            with_rvfi=True,
            with_muldiv=True,
            with_icache=True,
            icache_nways=2,
            icache_nlines=2,
            icache_nwords=4,
            icache_base=0x1000,
            icache_limit=0x4000,
            with_dcache=True,
            dcache_nways=2,
            dcache_nlines=2,
            dcache_nwords=4,
            dcache_base=0x1000,
            dcache_limit=0x4000)
        m.submodules.causal_spec = causal_spec = CausalCheck()

        # Wire input ports to Minerva core
        m.d.comb += cpu.external_interrupt.eq(0)
        m.d.comb += cpu.timer_interrupt.eq(0)
        m.d.comb += cpu.software_interrupt.eq(0)
        m.d.comb += cpu.ibus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.ibus.ack.eq(AnySeq(1))
        m.d.comb += cpu.ibus.err.eq(0)
        m.d.comb += cpu.dbus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.dbus.ack.eq(AnySeq(1))
        m.d.comb += cpu.dbus.err.eq(0)

        m.d.comb += causal_spec.reset.eq(AnySeq(1))
        m.d.comb += causal_spec.check.eq(AnySeq(1))

        m.d.comb += causal_spec.rvfi_valid.eq(cpu.rvfi.valid)
        m.d.comb += causal_spec.rvfi_rd_addr.eq(cpu.rvfi.rd_addr)
        m.d.comb += causal_spec.rvfi_order.eq(cpu.rvfi.order)
        m.d.comb += causal_spec.rvfi_rs1_addr.eq(cpu.rvfi.rs1_addr)
        m.d.comb += causal_spec.rvfi_rs2_addr.eq(cpu.rvfi.rs2_addr)

        return m

class CausalTestCase(FHDLTestCase):
    def verify(self):
        self.assertFormal(CausalSpec(), mode="cover", depth=20, spec_name="verify_causal")
        self.assertFormal(CausalSpec(), mode="bmc", depth=20, spec_name="verify_causal")
        print("verify_causal PASS")

class LivenessSpec(Elaboratable):
    def elaborate(self, platform):
        m = Module()

        m.submodules.cpu = cpu = Minerva(
            with_rvfi=True,
            with_muldiv=True,
            with_icache=True,
            icache_nways=2,
            icache_nlines=2,
            icache_nwords=4,
            icache_base=0x1000,
            icache_limit=0x4000,
            with_dcache=True,
            dcache_nways=2,
            dcache_nlines=2,
            dcache_nwords=4,
            dcache_base=0x1000,
            dcache_limit=0x4000)
        m.submodules.liveness_spec = liveness_spec = LivenessCheck()

        # Wire input ports to Minerva core
        m.d.comb += cpu.external_interrupt.eq(0)
        m.d.comb += cpu.timer_interrupt.eq(0)
        m.d.comb += cpu.software_interrupt.eq(0)
        m.d.comb += cpu.ibus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.ibus.ack.eq(AnySeq(1))
        m.d.comb += cpu.ibus.err.eq(0)
        m.d.comb += cpu.dbus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.dbus.ack.eq(AnySeq(1))
        m.d.comb += cpu.dbus.err.eq(0)

        cycle = Signal(8, reset=0)
        with m.If(cycle != 0xFF):
            m.d.sync += cycle.eq(cycle + 1)
        m.d.comb += liveness_spec.reset.eq(cycle < 1)
        m.d.comb += liveness_spec.trig.eq(cycle == 11)
        m.d.comb += liveness_spec.check.eq(cycle == 39)

        m.d.comb += liveness_spec.rvfi_valid.eq(cpu.rvfi.valid)
        m.d.comb += liveness_spec.rvfi_order.eq(cpu.rvfi.order)
        m.d.comb += liveness_spec.rvfi_halt.eq(cpu.rvfi.halt)

        # Bounded fairness constraints
        ibus_wait = Signal(2, reset=0)
        dbus_wait = Signal(2, reset=0)
        with m.If(cpu.ibus.cyc & cpu.ibus.stb & ~(cpu.ibus.ack | cpu.ibus.err)):
            m.d.sync += ibus_wait.eq(ibus_wait + 1)
        with m.Else():
            m.d.sync += ibus_wait.eq(0)
        with m.If(cpu.dbus.cyc & cpu.dbus.stb & ~(cpu.dbus.ack | cpu.dbus.err)):
            m.d.sync += dbus_wait.eq(dbus_wait + 1)
        with m.Else():
            m.d.sync += dbus_wait.eq(0)
        m.d.comb += Assume((ibus_wait < 2) & (dbus_wait < 2))
        with m.If(liveness_spec.reset):
            m.d.sync += ibus_wait.eq(0)
            m.d.sync += dbus_wait.eq(0)

        return m

class LivenessTestCase(FHDLTestCase):
    def verify(self):
        self.assertFormal(LivenessSpec(), mode="cover", depth=40, spec_name="verify_liveness")
        self.assertFormal(LivenessSpec(), mode="bmc", depth=40, spec_name="verify_liveness")
        print("verify_liveness PASS")

class UniqueSpec(Elaboratable):
    def elaborate(self, platform):
        m = Module()

        m.submodules.cpu = cpu = Minerva(
            with_rvfi=True,
            with_muldiv=True,
            with_icache=True,
            icache_nways=2,
            icache_nlines=2,
            icache_nwords=4,
            icache_base=0x1000,
            icache_limit=0x4000,
            with_dcache=True,
            dcache_nways=2,
            dcache_nlines=2,
            dcache_nwords=4,
            dcache_base=0x1000,
            dcache_limit=0x4000)
        m.submodules.unique_spec = unique_spec = UniqueCheck()

        # Wire input ports to Minerva core
        m.d.comb += cpu.external_interrupt.eq(0)
        m.d.comb += cpu.timer_interrupt.eq(0)
        m.d.comb += cpu.software_interrupt.eq(0)
        m.d.comb += cpu.ibus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.ibus.ack.eq(AnySeq(1))
        m.d.comb += cpu.ibus.err.eq(0)
        m.d.comb += cpu.dbus.dat_r.eq(AnySeq(32))
        m.d.comb += cpu.dbus.ack.eq(AnySeq(1))
        m.d.comb += cpu.dbus.err.eq(0)
        
        cycle = Signal(8, reset=0)
        with m.If(cycle != 0xFF):
            m.d.sync += cycle.eq(cycle + 1)
        m.d.comb += unique_spec.reset.eq(cycle < 1)
        m.d.comb += unique_spec.trig.eq(cycle == 11)
        m.d.comb += unique_spec.check.eq(cycle == 24)

        m.d.comb += unique_spec.rvfi_valid.eq(cpu.rvfi.valid)
        m.d.comb += unique_spec.rvfi_order.eq(cpu.rvfi.order)

        return m

class UniqueTestCase(FHDLTestCase):
    def verify(self):
        self.assertFormal(UniqueSpec(), mode="cover", depth=25, spec_name="verify_uniqueness")
        self.assertFormal(UniqueSpec(), mode="bmc", depth=25, spec_name="verify_uniqueness")
        print("verify_uniqueness PASS")

print('*' * 80)
print('*' + ' ' * 78 + '*')
print('*                        Verifying the Minerva core ...                        *')
print('*' + ' ' * 78 + '*')
print('*' * 80)

print("Verifying RV32M instructions ...")
p_insn = Process(target=InsnTestCase().verify)
p_insn.start()

print("Verifying PC forward checks ...")
p_pc_fwd = Process(target=PcFwdTestCase().verify)
p_pc_fwd.start()

print("Verifying PC backward checks ...")
p_pc_bwd = Process(target=PcBwdTestCase().verify)
p_pc_bwd.start()

print("Verifying register checks ...")
p_reg = Process(target=RegTestCase().verify)
p_reg.start()

print("Verifying causal checks ...")
p_causal = Process(target=CausalTestCase().verify)
p_causal.start()

print("Verifying liveness checks ...")
p_liveness = Process(target=LivenessTestCase().verify)
p_liveness.start()

print("Verifying uniqueness checks ...")
p_uniqueness = Process(target=UniqueTestCase().verify)
p_uniqueness.start()

p_insn.join()
p_pc_fwd.join()
p_pc_bwd.join()
p_reg.join()
p_causal.join()
p_liveness.join()
p_uniqueness.join()

print('*' * 80)
print('*' + ' ' * 78 + '*')
print('*                      All verification tasks successful!                      *')
print('*' + ' ' * 78 + '*')
print('*' * 80)