from amaranth import *

class InOutBuff(Elaboratable):
    '''
    This module wraps another modules input and output with a buffer
    This is usefull for doing timeing analysis on combinational logic 

    An instance of a module should be passed, not the module itself
    '''
    def __init__(self, sub_module: Elaboratable):
        assert sub_module.ports is not None, 'sub_module must have ports'

        self.sub_module = sub_module
        ports_in = [Signal(port.width, name=port.name + '_inbuf') for port in sub_module.ports['in']]
        ports_out = [Signal(port.width, name=port.name + '_outbuf') for port in sub_module.ports['out']]
        self.ports = {'in': ports_in, 'out': ports_out}

    def elaborate(self, platform):
        m = Module()
        m.submodules.sub = self.sub_module

        for i in range(len(self.ports['in'])):
            m.d.sync += self.sub_module.ports['in'][i].eq(self.ports['in'][i])

        for i in range(len(self.ports['out'])):
            m.d.sync += self.ports['out'][i].eq(self.sub_module.ports['out'][i])

        return m