#!/usr/bin/env python
#
# Implements the tests for the add/store unit.  The add/store unit
# consists of two halves: the function unit (FU) and the computation
# unit (CU).


import inspect

from nmigen.test.tools import FHDLTestCase
from nmigen.back.pysim import Simulator, Delay
from nmigen import (
    Signal, Module, ClockDomain, Elaboratable, ResetSignal, Const, Repl, Cat,
)
from nmigen.hdl.ast import Assume, Past, Assert, Stable, Rose, Fell


SRC2_IMM12 = 0
SRC2_REG2 = 1
SRC2_unused = 2
SRC2_NOT_REG2 = 3


def create_cu_interface(self, xlen=64, fumax=7, xmax=31):
    """Creates the computation unit interface for the add/store unit.

    If specified, xlen specifies the maximum register width.  If not
    specified, xlen defaults to 64.

    If specified, fumax specifies the largest integer that can
    uniquely identify a FU.  If not specified, fumax defaults to 7.

    xmax specifies the highest addressable X-register.  If left
    unspecified, this defaults to 31.  Don't change this unless you
    know exactly what you're doing!
    """

    # The following signals come from the corresponding function unit.

    self.i_go_write = Signal()
    self.i_go_sum = Signal()
    self.i_inst_src2 = Signal(min=0, max=SRC2_NOT_REG2)
    self.i_inst_imm12 = Signal(min=-2048, max=2047)
    self.i_inst_cin = Signal(1)

    # The following come from the "trunk" (really, the register file).

    self.i_trunk_data1 = Signal(xlen)
    self.i_trunk_data2 = Signal(xlen)
    self.o_trunk_data = Signal(xlen)

    # The following are diagnostic outputs only.

    self.o_fv_addend1 = Signal(xlen)
    self.o_fv_addend2 = Signal(xlen)
    self.o_fv_ea = Signal(xlen)
    self.o_fv_src2 = Signal(self.i_inst_src2.nbits)
    self.o_fv_imm12 = Signal(self.i_inst_imm12.nbits)
    self.o_fv_cin = Signal(self.i_inst_cin.nbits)


class AddStoreCU(Elaboratable):
    """Implements the data flow logic for one add/store unit.  See also
    the AddStoreFU class for the function unit that drives instances
    of this class.

    NOTE: All signals in the fv_ bus are reserved for external formal
    verification.  DO NOT use these signals for production purposes.
    """

    def __init__(self):
        """
        Instantiates a new add/store computation unit.
        """

        create_cu_interface(self)

    def elaborate(self, platform):
        m = Module()
        comb = m.d.comb
        sync = m.d.sync

        xlen = self.o_trunk_data.nbits

        # The 2nd addend to the adder can be taken from one of three
        # sources, as set by the instruction's src2 bits.

        addend2 = Signal(xlen)

        comb += [
            self.o_fv_src2.eq(self.i_inst_src2),
            self.o_fv_imm12.eq(self.i_inst_imm12),
            self.o_fv_cin.eq(self.i_inst_cin),
        ]

        with m.If(self.i_inst_src2 == SRC2_IMM12):
            comb += addend2.eq(Cat(
                self.i_inst_imm12,
                Repl(self.i_inst_imm12[11], xlen-12)
            ))

        with m.If(self.i_inst_src2 == SRC2_REG2):
            comb += addend2.eq(self.i_trunk_data2)

        with m.If(self.i_inst_src2 == SRC2_NOT_REG2):
            comb += addend2.eq(~self.i_trunk_data2)

        # The adder passively produces a sum from its two inputs.

        sum = Signal(xlen)
        comb += [
            self.o_fv_addend1.eq(self.i_trunk_data1),
            self.o_fv_addend2.eq(addend2),
            sum.eq(
                (self.i_trunk_data1 + addend2 + self.i_inst_cin)[0:xlen]
            ),
        ]

        # The adder places its results in the effective address
        # register, reg_ea.

        reg_ea = Signal(xlen)

        with m.If(self.i_go_sum):
            sync += reg_ea.eq(sum)

        comb += self.o_fv_ea.eq(reg_ea)

        # If we don't have permission to drive the trunk, then make
        # sure we don't drive it.  In an FPGA, at least, the trunk
        # is implemented using a giant OR network, so that means a
        # quiet trunk is a zero trunk.

        comb += [
            self.o_trunk_data.eq(
                reg_ea & Repl(self.i_go_write, self.o_trunk_data.nbits)
            ),
        ]

        return m


class AddStoreCUFormal(Elaboratable):
    def __init__(self):
        create_cu_interface(self)

    def elaborate(self, platform):
        m = Module()
        sync = m.d.sync
        comb = m.d.comb

        dut = AddStoreCU()
        m.submodules.dut = dut
        comb += [
            dut.i_go_write.eq(self.i_go_write),
            dut.i_go_sum.eq(self.i_go_sum),

            dut.i_inst_src2.eq(self.i_inst_src2),
            dut.i_inst_imm12.eq(self.i_inst_imm12),
            dut.i_inst_cin.eq(self.i_inst_cin),

            self.o_trunk_data.eq(dut.o_trunk_data),
            dut.i_trunk_data1.eq(self.i_trunk_data1),
            dut.i_trunk_data2.eq(self.i_trunk_data2),
            
            self.o_fv_ea.eq(dut.o_fv_ea),
            self.o_fv_addend1.eq(dut.o_fv_addend1),
            self.o_fv_addend2.eq(dut.o_fv_addend2),
            self.o_fv_src2.eq(dut.o_fv_src2),
            self.o_fv_imm12.eq(dut.o_fv_imm12),
            self.o_fv_cin.eq(dut.o_fv_cin),
        ]

        xlen = self.o_trunk_data.nbits

        # Create a flag that helps prevent making assertions and
        # assumptions that spans before simulation time began.

        z_past_valid = Signal(1, reset=0)
        sync += z_past_valid.eq(1)

        # Everything must be reset before we can do anything else.

        rst = ResetSignal()
        with m.If(~z_past_valid):
            sync += Assume(rst)

        # The first addend of the adder should always reflect the
        # first source register on the trunk.

        comb += Assert(self.o_fv_addend1 == self.i_trunk_data1)

        # The second addend depends on the precise instruction
        # being executed.  For ADDI and store instructions,
        # the immediate input should be reflected.  For ADD
        # instructions, it should come from the second source
        # register.  Additionally, if being taken from the 2nd
        # source register, it may optionally be inverted so as
        # to effect a subtraction.

        with m.If(self.o_fv_src2 == SRC2_IMM12):
            comb += [
                Assert(self.o_fv_addend2[0:12] == self.o_fv_imm12),
                Assert(
                    self.o_fv_addend2[12:xlen] == Repl(self.o_fv_imm12[11], xlen-12)
                )
            ]

        with m.If(self.o_fv_src2 == SRC2_REG2):
            comb += Assert(self.o_fv_addend2 == self.i_trunk_data2)

        with m.If(self.o_fv_src2 == SRC2_NOT_REG2):
            comb += Assert(self.o_fv_addend2 == ~self.i_trunk_data2)

        with m.If(self.o_fv_src2 == SRC2_unused):
            comb += Assert(self.o_fv_addend2 == Const(0, xlen))

        # If GO_SUM is asserted, then the sum computed by the
        # adder must be preserved.  We don't really care about
        # 32-bit sign extension here, as that'll be located at
        # the register file data input.

        with m.If(z_past_valid & Stable(rst) & ~rst & Past(self.i_go_sum)):
            sync += [
                Assert(
                    self.o_fv_ea == (
                        Past(self.o_fv_addend1) +
                        Past(self.o_fv_addend2) +
                        Past(self.o_fv_cin)
                    )[0:xlen]
                )
            ]

        with m.If(z_past_valid & Stable(rst) & ~rst & ~Past(self.i_go_sum)):
            sync += Assert(Stable(self.o_fv_ea))

        # If GO_WRITE is negated, then we don't drive the trunk.

        with m.If(z_past_valid & Stable(rst) & ~rst & ~self.i_go_write):
            sync += [
                Assert(self.o_trunk_data == Const(0, xlen)),
            ]

        # If GO_WRITE is asserted, however, we must drive the trunk.

        with m.If(z_past_valid & Stable(rst) & ~rst & self.i_go_write):
            sync += [
                Assert(self.o_trunk_data == self.o_fv_ea),
            ]

        return m

class AddStoreCUTestCase(FHDLTestCase):

    # "P"eriod -- how long the clock is high or low.  You might not
    # think so, but for my wretched joints, I find typing "self.P"
    # easier than "10e-9" all the time.
    P = 10e-9

    def stdfiles(self):
        """Opens correctly named VCD and GTKWave files for writing."""

        caller = inspect.currentframe().f_back.f_code.co_name
        return {
            'vcd_file': open("{}.vcd".format(caller), "w"),
            'gtkw_file': open("{}.gtkw".format(caller), "w"),
        }

    def setup_cu(self, reset_pc=0):
        create_cu_interface(self)
        ports = [
        ]
        ports = [self.__getattribute__(p) for p in ports]

        self.dut = AddStoreCU()
        
        self.m = Module()
        self.sync = ClockDomain()
        self.m.domains += self.sync
        self.m.submodules += self.dut

        self.m.d.comb += [
        ]

        return ports

    def tick(self):
        yield self.sync.clk.eq(1)
        yield Delay(self.P)
        yield self.sync.clk.eq(0)
        yield Delay(self.P)

    def reset(self):
        yield self.sync.rst.eq(1)
        yield from self.tick()
        yield self.sync.rst.eq(0)
        yield from self.tick()

    def test_add_store_cu_formally(self):
        self.assertFormal(AddStoreCUFormal(), mode="bmc", depth=100)
        self.assertFormal(AddStoreCUFormal(), mode="prove", depth=100)

    #def test_idle_state(self):
    #    traces = self.setup_cu()
    #    with Simulator(
    #        self.m.elaborate(platform=None), traces=traces, **(self.stdfiles())
    #    ) as sim:
    #        def process():
    #            yield from self.reset()
    #            self.assertEqual((yield self.o_trunk_rd), 0)
    #            self.assertEqual((yield self.o_trunk_rd_dat), 0)
    #            self.assertEqual((yield self.o_trunk_fuid), 0)
    #            
    #        sim.add_process(process)
    #        sim.run()
