#!/usr/bin/env python
# coding: utf-8

# ## Barrel Shifter
# 
# A Barrel shifter performs the operation `v << s` (or `v >> s` with flipped bits) where `s` is not constant.
# An elegant implementation is the approach via cascaded multiplexer elements.

# This `cshift` primitive is later instanced procedurally, iterating through the variable `i` 

# ### Shifter element in MyHDL

# In[1]:


import sys
sys.path.insert(0, "../..")


# In[2]:


from myirl.emulation.myhdl import *

@block
def cshift(q : Signal.Output, a : Signal, b : Signal, sbit : Signal, 
           msb : Signal,
           # These can take a signal or a fixed bool:
           asr, rotate,
           WRAP : bool):
    carry = Signal(bool()) # Carry bit
    u, v = [ Signal(bool()) for i in range(2) ]
        
    @always_comb
    def assign_carry():
        if asr: # arithmetic shift right
            carry.next = sbit & msb
        else:
            carry.next = 0

    @always_comb
    def assign():
        u.next = a & ~sbit
        
        if rotate == False:         
            if WRAP:
                v.next = carry
            else:
                v.next = b & sbit
        else:
            v.next = b & sbit
            
    @always_comb
    def assign_q():    
        q.next = u | v
        
    return instances()


# ### Shifter stage

# In[3]:


@block
def shifter_stage(
    shifter,
    w_in : Signal,
    w_out : Signal.Output,
    msb : Signal,
    nmux : int, sbit : Signal, DATA_WIDTH : int, W_POWER : int, asr : bool, rotate : bool
):
    inst = []

    # Create signal array
    w = [ Signal(bool()) for i in range(DATA_WIDTH) ]
    wo = concat(*reversed(w))
    wi = [ w_in[i] for i in range(DATA_WIDTH) ]

    MUX_W = DATA_WIDTH // nmux

    for imux in range(nmux):
        tmp = imux * MUX_W
        # print(imux)
        for i in range(tmp, tmp + MUX_W):
            j = i + MUX_W//2
            m = j % DATA_WIDTH
            inst.append(shifter(w[m], wi[m], wi[i], sbit, msb, asr, rotate,
                         j >= DATA_WIDTH ))

    @always_comb
    def assign():
        w_out.next = wo

    return instances()


# ### The barrel shifter implementation

# In[4]:


@block
def barrel_shifter(shifter, clk : ClkSignal, ce : Signal, val : Signal, s : Signal, result : Signal.Output, \
                   rotate = False, W_POWER = 5 ):
    
    DATA_WIDTH = 2 ** W_POWER
    print("DATA WIDTH", DATA_WIDTH, "ROTATE", rotate)
        
    worker = [ val ]
    worker = worker + [ Signal(intbv()[DATA_WIDTH:]) for i in range(W_POWER) ]
    msb = val[DATA_WIDTH-1]

    sbit = [ s[i] for i in range(len(s))]
    
    shifter_stages = []
    for stage in range(W_POWER):
        K = W_POWER - stage - 1
        print("Stage %d" % stage)
        shifter_stages.append( \
                shifter_stage(shifter, worker[stage], worker[stage + 1], msb, 2 ** stage, sbit[K], \
                            DATA_WIDTH, W_POWER, False, rotate) \
                             )
        
    @always(clk.posedge)
    def assign():
        if ce == True:
            result.next = worker[W_POWER]
        
    return instances()


# ## Translate and test
# 
# The translation into the target RTL or HDL takes place via a test bench factory function. This allows us to switch between different simulator back end methods. Here, we can choose between `GHDL` and `CXXRTL`.

# In[5]:


from simulation import *

from myirl.simulation import print_
from myirl.emulation.factory_class import factory
from myirl.targets import pyosys

class example_design(factory.Module):
    def __init__(self, name, simclass, *args, **kwargs):
        super().__init__(name, simclass, *args, **kwargs)
        self.W_POWER = 4
        self.debug = False

        # Top level signal set
        self.clk = self.ClkSignal(name="clk")
        self.ce = self.Signal(bool())
        self.val, self.result = [ self.Signal(intbv(0xaa00)[2 ** self.W_POWER:]) for i in range(2) ]
        self.result.rename("result")
        self.s = self.Signal(intbv()[self.W_POWER:])
    
    def build(self):
        return barrel_shifter(cshift,
            clk = self.clk,
            ce = self.ce,
            val = self.val,
            s = self.s,
            result = self.result,
            rotate = False,
            W_POWER = self.W_POWER
        )
    
    def emit_rtlil(self, fileprefix):
        tgt = pyosys.RTLIL("barrel_shifter")
        inst = self.build()
        d = inst.elab(tgt)
        d[0].write_rtlil(fileprefix)

    @factory.testbench('ns')
    def tb_rtl(self):
        
        inst = self.build()

        clk = self.clk
        ce = self.ce
        val = self.val
        result = self.result
        s = self.s
        
        @self.always(delay(2))
        def clkgen():
            clk.next = ~clk

        TEST_VALUES = [
            (0xdead, 8, 0xad00),
            (0x8f01, 15, 0x8000),
        ]

        @self.sequence
        def stim():
            for item in TEST_VALUES:
                ce.next = False
                s.next = item[1]
                val.next = item[0]
                yield(clk.posedge)
                ce.next = True
                yield(clk.posedge)
                yield(clk.posedge)

                print(result, type(result))
                assert result == item[2]

            raise StopSimulation

        return instances()    


# In[6]:


from yosys.simulator import CXXRTL

d = example_design("rtlil_bs", CXXRTL)
d.emit_rtlil("test")


# In[7]:


from myirl.test.ghdl import GHDL

# XXX Note: we must carry out the simulation backend sequence in this order. Otherwise,
# a yet undetermined sticky effect will report a combinatorial loop error.
for sim in CXXRTL, GHDL:
    d = example_design("testbench", sim)
    tb = d.tb_rtl()
    print(76 * '=')
    tb.run(60)


# ## Further Optimization
# 
# A bit more optimization can be done by moving the code into a cythonizeable library. See
# [Cythonized barrel shifter](example_cython_barrelshifter.ipynb)
# 
# This primitive approach however is not very compact and this design should not be mass-instanced using CXXRTL due to issues with combinatorial logic.
# 
# See [Compact barrel shifter implementation](bshift_compact.ipynb) for an improved version.