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 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 ...riscv_formal_parameters import *

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)
        m.submodules.insn_spec = insn_spec = InsnCheck(
            params=RISCVFormalParameters(32, 32, False, False, False),
            insn_model=self.insn_model,
            rvformal_addr_valid=lambda x:Const(1))

        m.d.comb += insn_spec.reset.eq(0)
        m.d.comb += insn_spec.check.eq(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):
        insns = [
            ('LUI', InsnLui),
            ('AUIPC', InsnAuipc),
            ('JAL', InsnJal),
            ('JALR', InsnJalr),
            ('BEQ', InsnBeq),
            ('BNE', InsnBne),
            ('BLT', InsnBlt),
            ('BGE', InsnBge),
            ('BLTU', InsnBltu),
            ('BGEU', InsnBgeu),
            ('LB', InsnLb),
            ('LH', InsnLh),
            ('LW', InsnLw),
            ('LBU', InsnLbu),
            ('LHU', InsnLhu),
            ('SB', InsnSb),
            ('SH', InsnSh),
            ('SW', InsnSw),
            ('ADDI', InsnAddi),
            ('SLTI', InsnSlti),
            ('SLTIU', InsnSltiu),
            ('XORI', InsnXori),
            ('ORI', InsnOri),
            ('ANDI', InsnAndi),
            ('SLLI', InsnSlli),
            ('SRLI', InsnSrli),
            ('SRAI', InsnSrai),
            ('ADD', InsnAdd),
            ('SUB', InsnSub),
            ('SLL', InsnSll),
            ('SLT', InsnSlt),
            ('SLTU', InsnSltu),
            ('XOR', InsnXor),
            ('SRL', InsnSrl),
            ('SRA', InsnSra),
            ('OR', InsnOr),
            ('AND', InsnAnd)
        ]
        for insn_name, insn_model in insns:
            print("- Verifying instruction %s ..." % insn_name)
            self.assertFormal(InsnSpec(insn_model), mode="bmc", depth=12, engine="smtbmc --nopresat")

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

        m.submodules.cpu = cpu = Minerva(with_rvfi=True)
        m.submodules.pc_fwd_spec = pc_fwd_spec = PcFwdCheck(
            params=RISCVFormalParameters(32, 32, False, False, False),
            rvformal_addr_valid=lambda x:Const(1))

        m.d.comb += pc_fwd_spec.reset.eq(0)
        m.d.comb += pc_fwd_spec.check.eq(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="bmc", depth=12, engine="smtbmc --nopresat")

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

        m.submodules.cpu = cpu = Minerva(with_rvfi=True)
        m.submodules.pc_bwd_spec = pc_bwd_spec = PcBwdCheck(
            params=RISCVFormalParameters(32, 32, False, False, False),
            rvformal_addr_valid=lambda x:Const(1))

        m.d.comb += pc_bwd_spec.reset.eq(0)
        m.d.comb += pc_bwd_spec.check.eq(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="bmc", depth=12, engine="smtbmc --nopresat")

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

        m.submodules.cpu = cpu = Minerva(with_rvfi=True)
        m.submodules.reg_spec = reg_spec = RegCheck(params=RISCVFormalParameters(32, 32, False, False, False))

        m.d.comb += reg_spec.reset.eq(0)
        m.d.comb += reg_spec.check.eq(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="bmc", depth=12, engine="smtbmc --nopresat")

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

        m.submodules.cpu = cpu = Minerva(with_rvfi=True)
        m.submodules.causal_spec = causal_spec = CausalCheck()

        m.d.comb += causal_spec.reset.eq(0)
        m.d.comb += causal_spec.check.eq(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="bmc", depth=12, engine="smtbmc --nopresat")

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

print("Verifying RV32I instructions ...")
InsnTestCase().verify()

print("Verifying PC forward checks ...")
PcFwdTestCase().verify()

print("Verifying PC backward checks ...")
PcBwdTestCase().verify()

print("Verifying register checks ...")
RegTestCase().verify()

print("Verifying causal checks ...")
CausalTestCase().verify()

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