sygus_am.py

"""
Amaranth implementation of the bit operation synthesizer for FPGA.

This module implements the same bit operations (identity, not, shift_left, shift_right)
as the Python version, but in synthesizable Amaranth HDL for FPGA deployment.
"""

from amaranth import *
from amaranth.sim import Simulator, Tick, Delay


class BitOperationUnit(Elaboratable):
    """
    A single bit operation unit that can perform one of four operations:
    - identity: pass through
    - not: bitwise NOT
    - shift_left: shift left by 1
    - shift_right: shift right by 1
    """
    
    def __init__(self, width=5):
        self.width = width
        # Input signal
        self.data_in = Signal(width)
        # Operation select: 0=identity, 1=not, 2=shift_left, 3=shift_right
        self.op_sel = Signal(2)
        # Output signal
        self.data_out = Signal(width)
    
    def elaborate(self, platform):
        m = Module()
        
        with m.Switch(self.op_sel):
            with m.Case(0):  # identity
                m.d.comb += self.data_out.eq(self.data_in)
            with m.Case(1):  # not
                m.d.comb += self.data_out.eq(~self.data_in)
            with m.Case(2):  # shift_left 1
                # Shift left: bits move to higher positions, 0 at LSB
                # Cat(0, bits[:-1]) -> 0 goes to bit 0, old bits 0-3 go to bits 1-4
                shifted = Cat(0, self.data_in[:-1])
                m.d.comb += self.data_out.eq(shifted)
            with m.Case(3):  # shift_right 1
                # Shift right: bits move to lower positions, 0 at MSB
                # Cat(bits[1:], 0) -> old bits 1-4 go to bits 0-3, 0 goes to bit 4
                shifted = Cat(self.data_in[1:], 0)
                m.d.comb += self.data_out.eq(shifted)
        
        return m


class BitOperationPipeline(Elaboratable):
    """
    A pipeline of bit operation units that can apply a sequence of operations.
    
    This module applies a sequence of operations to transform input to output.
    The operations are controlled by op_sequence input.
    """
    
    def __init__(self, width=5, depth=4):
        self.width = width
        self.depth = depth
        
        # Input/output
        self.data_in = Signal(width)
        self.data_out = Signal(width)
        
        # Operation sequence: depth * 2 bits (2 bits per operation)
        self.op_sequence = Signal(depth * 2)
        
        # Enable signal
        self.enable = Signal()
        
        # Internal pipeline registers
        self.pipeline = [Signal(width) for _ in range(depth + 1)]
        self.op_units = [BitOperationUnit(width) for _ in range(depth)]
    
    def elaborate(self, platform):
        m = Module()
        
        # Instantiate operation units
        for i, op_unit in enumerate(self.op_units):
            m.submodules[f"op_unit_{i}"] = op_unit
        
        # Connect pipeline
        m.d.comb += self.pipeline[0].eq(self.data_in)
        
        for i in range(self.depth):
            # Extract 2-bit operation code for this stage
            op_code = self.op_sequence[i*2:(i+1)*2]
            
            # Connect operation unit
            m.d.comb += self.op_units[i].data_in.eq(self.pipeline[i])
            m.d.comb += self.op_units[i].op_sel.eq(op_code)
            m.d.comb += self.pipeline[i+1].eq(self.op_units[i].data_out)
        
        # Output
        m.d.comb += self.data_out.eq(self.pipeline[self.depth])
        
        return m


class SynthesizerCore(Elaboratable):
    """
    Main synthesizer core that can search for operation sequences.
    
    This is a simplified version that applies a given operation sequence.
    For full brute-force search, you would need to add a state machine
    that iterates through all possible operation sequences.
    """
    
    def __init__(self, width=5, depth=4):
        self.width = width
        self.depth = depth
        
        # Input/output
        self.input_val = Signal(width)
        self.target_output = Signal(width)
        self.output_val = Signal(width)
        
        # Operation sequence control
        self.op_sequence = Signal(depth * 2)
        self.sequence_valid = Signal()
        
        # Status
        self.match = Signal()  # 1 if output matches target
        
        # Pipeline
        self.pipeline = BitOperationPipeline(width, depth)
    
    def elaborate(self, platform):
        m = Module()
        
        m.submodules.pipeline = self.pipeline
        
        # Connect pipeline
        m.d.comb += self.pipeline.data_in.eq(self.input_val)
        m.d.comb += self.pipeline.op_sequence.eq(self.op_sequence)
        m.d.comb += self.pipeline.enable.eq(1)
        m.d.comb += self.output_val.eq(self.pipeline.data_out)
        
        # Check if output matches target
        m.d.comb += self.match.eq(self.output_val == self.target_output)
        
        return m


def operation_name_to_code(name):
    """Convert operation name to 2-bit code."""
    op_map = {
        "identity": 0,
        "not": 1,
        "shift_left 1": 2,
        "shift_right 1": 3,
    }
    return op_map.get(name, 0)


def operation_sequence_to_bits(sequence, depth):
    """Convert a list of operation names to a bit sequence."""
    bits = 0
    for i, op_name in enumerate(sequence):
        if i >= depth:
            break
        op_code = operation_name_to_code(op_name)
        bits |= op_code << (i * 2)
    return bits


def test_bit_operations():
    """Test the bit operation unit."""
    print("=" * 50)
    print("Testing BitOperationUnit")
    print("=" * 50)

    width = 5
    dut = BitOperationUnit(width)

    def test_case(name, op_sel, input_val, expected):
        async def testbench(ctx):
            ctx.set(dut.data_in, input_val)
            ctx.set(dut.op_sel, op_sel)
            result = ctx.get(dut.data_out)
            status = "✓" if result == expected else "✗"
            print(f"{status} {name}: input={input_val:05b}, op={op_sel}, "
                  f"output={result:05b}, expected={expected:05b}")

        return testbench

    sim = Simulator(dut)
    sim.add_testbench(test_case("Identity", 0, 0b10100, 0b10100))
    sim.run()

    sim = Simulator(dut)
    sim.add_testbench(test_case("NOT", 1, 0b11111, 0b00000))
    sim.run()

    sim = Simulator(dut)
    sim.add_testbench(test_case("Shift Left", 2, 0b10000, 0b00000))
    sim.run()

    sim = Simulator(dut)
    sim.add_testbench(test_case("Shift Right", 3, 0b10000, 0b01000))
    sim.run()


def test_pipeline():
    """Test the operation pipeline with a sequence."""
    print("\n" + "=" * 50)
    print("Testing BitOperationPipeline")
    print("=" * 50)

    width = 5
    depth = 4

    # Test case: identity -> identity -> shift_right -> not
    # Input: 11111, Expected: 10000
    sequence = ["identity", "identity", "shift_right 1", "not"]
    op_bits = operation_sequence_to_bits(sequence, depth)

    dut = BitOperationPipeline(width, depth)

    async def testbench(ctx):
        ctx.set(dut.data_in, 0b11111)
        ctx.set(dut.op_sequence, op_bits)
        ctx.set(dut.enable, 1)
        result = ctx.get(dut.data_out)
        expected = 0b10000
        status = "✓" if result == expected else "✗"
        print(f"{status} Pipeline test: input=11111, sequence={' -> '.join(sequence)}")
        print(f"   Output: {result:05b}, Expected: {expected:05b}")

    sim = Simulator(dut)
    sim.add_testbench(testbench)
    sim.run()


def test_synthesizer_core():
    """Test the main synthesizer core."""
    print("\n" + "=" * 50)
    print("Testing SynthesizerCore")
    print("=" * 50)

    width = 5
    depth = 4

    # Test case: verify the operation sequence produces expected output
    sequence = ["identity", "identity", "shift_right 1", "not"]
    op_bits = operation_sequence_to_bits(sequence, depth)

    dut = SynthesizerCore(width, depth)

    async def testbench(ctx):
        ctx.set(dut.input_val, 0b11111)
        ctx.set(dut.target_output, 0b10000)
        ctx.set(dut.op_sequence, op_bits)
        output = ctx.get(dut.output_val)
        match = ctx.get(dut.match)
        status = "✓" if match else "✗"
        print(f"{status} Synthesizer test: input=11111, target=10000")
        print(f"   Output: {output:05b}, Match: {match}")

    sim = Simulator(dut)
    sim.add_testbench(testbench)
    sim.run()


def generate_verilog_example():
    """Example of how to generate Verilog for FPGA synthesis."""
    from amaranth.back import verilog
    
    # Create a top-level module
    class TopLevel(Elaboratable):
        def __init__(self):
            self.width = 5
            self.depth = 4
            self.core = SynthesizerCore(self.width, self.depth)
            
            # Expose signals
            self.input_val = self.core.input_val
            self.target_output = self.core.target_output
            self.op_sequence = self.core.op_sequence
            self.output_val = self.core.output_val
            self.match = self.core.match
        
        def elaborate(self, platform):
            m = Module()
            m.submodules.core = self.core
            return m
    
    # Generate Verilog
    top = TopLevel()
    v_output = verilog.convert(top, ports=[
        top.input_val,
        top.target_output,
        top.op_sequence,
        top.output_val,
        top.match,
    ])
    
    print("Generated Verilog:")
    print("=" * 50)
    print(v_output)
    print("=" * 50)
    print("\nSave this to a .v file and use it in your FPGA toolchain.")


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1 and sys.argv[1] == "--test":
        test_bit_operations()
        test_pipeline()
        test_synthesizer_core()
    elif len(sys.argv) > 1 and sys.argv[1] == "--verilog":
        generate_verilog_example()
    else:
        print("Amaranth SYGUS Bit Operation Synthesizer")
        print("=" * 50)
        print("\nThis module implements bit operations for FPGA synthesis.")
        print("\nUsage:")
        print("  python amaranth_sygus.py --test     # Run test cases")
        print("  python amaranth_sygus.py --verilog  # Generate Verilog")
        print("\nTo use this in your FPGA project:")
        print("1. Install: pip install amaranth")
        print("2. Import SynthesizerCore or BitOperationPipeline")
        print("3. Add it to your top-level Amaranth module")
        print("4. Connect input/output signals")
        print("5. Generate Verilog with --verilog or use amaranth.build")
        print("6. Synthesize with your FPGA toolchain (Vivado, Quartus, etc.)")