From 5f54b8acd8978fd8281b3dc8359041049b94c747 Mon Sep 17 00:00:00 2001
From: saji <saji@saji.dev>
Date: Wed, 23 Oct 2024 15:14:53 -0500
Subject: [PATCH] fetcher basic doesn't support variable functions

---
 src/groovylight/fetcher.py            | 44 +++++++++++++++++++++------
 src/groovylight/tests/test_fetcher.py | 43 +++++++++++++++++++++-----
 src/groovylight/tests/test_hub75.py   |  1 -
 3 files changed, 70 insertions(+), 18 deletions(-)

diff --git a/src/groovylight/fetcher.py b/src/groovylight/fetcher.py
index bbd9413..56ba523 100644
--- a/src/groovylight/fetcher.py
+++ b/src/groovylight/fetcher.py
@@ -3,17 +3,15 @@
 # during operation, it is given a row index, and responds with the data.
 
 
-from amaranth import Module, Cat, Mux, ShapeLike, Signal, Assert, Array, unsigned
+from amaranth import Module, Signal, unsigned, Cat
 from amaranth.build import Platform
 from amaranth.lib import wiring, data
 from amaranth.lib.wiring import In, Out
-from amaranth.lib.memory import Memory, ReadPort, WritePort
 from amaranth.lib import stream
-from amaranth.utils import ceil_log2
 import logging
 
-from .common import Rgb666Layout, Hub75Stream, Hub75Ctrl, Hub75Data, Rgb888Layout
-from .geom import DisplayRotation, DisplayString
+from .common import Rgb888Layout
+from .geom import DisplayString
 
 
 logger = logging.getLogger(__name__)
@@ -91,7 +89,8 @@ class AddressGenerator(wiring.Component):
                 # stream data out as long as it's valid.
 
                 with m.If(
-                    self.coordstream.ready & (counter == self.geom.dimensions.length - 1)
+                    self.coordstream.ready
+                    & (counter == self.geom.dimensions.length - 1)
                 ):
                     m.next = "done"
                 with m.Elif(self.coordstream.ready):
@@ -107,17 +106,42 @@ class AddressGenerator(wiring.Component):
 
 
 class BasicFetcher(wiring.Component):
-    """A generic fetcher. Takes a function of the form f(x,y: int) -> RGB."""
+    """A generic function-based fetcher. Takes a function of the form f(x,y: int) -> RGB."""
 
     def __init__(
         self, geom: DisplayString, dfunc, data_shape=Rgb888Layout, *, src_loc_at=0
     ):
+        self.geom = geom
         self.dfunc = dfunc
         super().__init__(
             {
-                "pixstream": Out(stream.Signature(data.ArrayLayout(Rgb888Layout, 2))),
-                "start": In(1),
-                "addr": In(geom.dimensions.addr_bits),
+                "input": In(
+                    stream.Signature(data.ArrayLayout(CoordLayout, geom.dimensions.mux))
+                ),
+                "pixstream": Out(
+                    stream.Signature(data.ArrayLayout(data_shape, geom.dimensions.mux))
+                ),
             },
             src_loc_at=src_loc_at,
         )
+
+    def elaborate(self, platform: Platform) -> Module:
+        m = Module()
+
+        # test mode - pass through, r = x + y, g = x - y, b = {y,x}
+
+        colors = self.pixstream.payload
+        m.d.comb += [
+            self.input.valid.eq(self.pixstream.valid),
+            self.input.ready.eq(self.pixstream.ready),
+        ]
+
+        for i in range(self.geom.dimensions.mux):
+            inp = self.input.payload[i]
+            m.d.comb += [
+                colors[i].red.eq(inp.x + inp.y),
+                colors[i].green.eq(inp.x - inp.y),
+                colors[i].blue.eq(inp.x ^ inp.y),
+            ]
+
+        return m
diff --git a/src/groovylight/tests/test_fetcher.py b/src/groovylight/tests/test_fetcher.py
index e105fff..728d2f3 100644
--- a/src/groovylight/tests/test_fetcher.py
+++ b/src/groovylight/tests/test_fetcher.py
@@ -4,7 +4,7 @@ import random
 from random import randrange
 import pytest
 
-from groovylight.fetcher import AddressConverter, AddressGenerator
+from groovylight.fetcher import AddressConverter, AddressGenerator, BasicFetcher
 from groovylight.geom import DisplayString, Coord, DisplayDimensions, DisplayRotation
 
 ds_testdata = [
@@ -51,11 +51,11 @@ async def stream_put(ctx, stream, payload):
 
 
 generator_tests = [
-        (0, DisplayRotation.R0),
-        (0, DisplayRotation.R90),
-        (4, DisplayRotation.R90),
-        
-        ]
+    (0, DisplayRotation.R0),
+    (0, DisplayRotation.R90),
+    (4, DisplayRotation.R90),
+]
+
 
 @pytest.mark.parametrize("addr, rot", generator_tests)
 def test_generator(addr, rot):
@@ -74,7 +74,8 @@ def test_generator(addr, rot):
         await ctx.tick().until(dut.done == 1)
 
     expected = [
-        [ds.translate_coord(x, addr, 0), ds.translate_coord(x, addr, 1)] for x in range(128)
+        [ds.translate_coord(x, addr, 0), ds.translate_coord(x, addr, 1)]
+        for x in range(128)
     ]
 
     expected.reverse()
@@ -88,3 +89,31 @@ def test_generator(addr, rot):
     sim.add_testbench(stream_checker)
     with sim.write_vcd("generator.vcd"):
         sim.run()
+
+
+basicFetcher_tests = [
+    ({"x": 0, "y": 0}, {"red": 0, "green": 0, "blue": 0}),
+    ({"x": 1, "y": 1}, {"red": 2, "green": 0, "blue": 0}),
+    ({"x": 9, "y": 1}, {"red": 10, "green": 8, "blue": 8}),
+]
+
+
+@pytest.mark.parametrize("inp, expected", basicFetcher_tests)
+def test_basic_fetcher(inp, expected):
+    ds = DisplayString(
+        Coord(3, 0), DisplayDimensions(128, 64, mux=1), DisplayRotation.R0
+    )
+    dut = BasicFetcher(ds, None)
+    sim = Simulator(dut)
+
+    async def test(ctx):
+        ctx.set(dut.input.payload[0], inp)
+        res = ctx.get(dut.pixstream.payload)[0]
+        assert res["red"] == expected["red"]
+        assert res["green"] == expected["green"]
+        assert res["blue"] == expected["blue"]
+
+    sim.add_testbench(test)
+
+    with sim.write_vcd("fetcher.vcd"):
+        sim.run()
diff --git a/src/groovylight/tests/test_hub75.py b/src/groovylight/tests/test_hub75.py
index b31951e..80664ef 100644
--- a/src/groovylight/tests/test_hub75.py
+++ b/src/groovylight/tests/test_hub75.py
@@ -146,7 +146,6 @@ def test_datadriver_single(bcm):
     with sim.write_vcd("output.vcd"):
         sim.run()
 
-
 @pytest.mark.skip()
 def test_hub75_coordinator():
     m = Module()