From 2fd1136154abcb9882f494296a05bf3d89b38acc Mon Sep 17 00:00:00 2001
From: Byron Lathi <bslathi19@gmail.com>
Date: Mon, 27 Oct 2025 20:13:21 -0700
Subject: [PATCH] Actually randomize testing

---
 .../sim/poly1305_friendly_modulo.py           | 53 ++++++++++++++-----
 1 file changed, 41 insertions(+), 12 deletions(-)

diff --git a/ChaCha20_Poly1305_64/sim/poly1305_friendly_modulo.py b/ChaCha20_Poly1305_64/sim/poly1305_friendly_modulo.py
index a1a2a31..0294d5f 100644
--- a/ChaCha20_Poly1305_64/sim/poly1305_friendly_modulo.py
+++ b/ChaCha20_Poly1305_64/sim/poly1305_friendly_modulo.py
@@ -22,8 +22,16 @@ class TB:
         self.log = logging.getLogger("cocotb.tb")
         self.log.setLevel(logging.INFO)
 
+        self.input_queue = Queue()
+
+        self.expected_queue = Queue()
+        self.output_queue = Queue()
+
         cocotb.start_soon(Clock(self.dut.i_clk, CLK_PERIOD, units="ns").start())
 
+        cocotb.start_soon(self.run_input())
+        cocotb.start_soon(self.run_output())
+
     async def cycle_reset(self):
         await self._cycle_reset(self.dut.i_rst, self.dut.i_clk)
 
@@ -38,25 +46,46 @@ class TB:
         await RisingEdge(clk)
         await RisingEdge(clk)
 
+    async def write_input(self, value: int, shift_amount: int):
+        await self.input_queue.put((value, shift_amount))
+        await self.expected_queue.put((value << (shift_amount*26)) % PRIME)
+
+    async def run_input(self):
+        while True:
+            value, shift_amount = await self.input_queue.get()
+            self.dut.i_valid.value = 1
+            self.dut.i_val.value = value
+            self.dut.i_shift_amount.value = shift_amount
+            await RisingEdge(self.dut.i_clk)
+            self.dut.i_valid.value = 0
+            self.dut.i_shift_amount.value = 0
+            self.dut.i_val.value = 0
+
+    async def run_output(self):
+        while True:
+            await RisingEdge(self.dut.i_clk)
+            if self.dut.o_valid.value:
+                await self.output_queue.put(self.dut.o_result.value.integer)
+
 @cocotb.test
 async def test_sanity(dut):
     tb = TB(dut)
 
     await tb.cycle_reset()
 
-    value_a = random.randint(1,2**(130+16))
+    count = 1024
 
-    # value_a = PRIME + 1000000
-    tb.dut.i_valid.value = 1
-    tb.dut.i_val.value = value_a
-    await RisingEdge(tb.dut.i_clk)
-    tb.dut.i_valid.value = 0
-    tb.dut.i_val.value = 0
+    for _ in range(count):
+        await tb.write_input(random.randint(1,2**(130+16)), random.randint(0, 4))
 
-    await RisingEdge(tb.dut.o_valid)
-    value = tb.dut.o_result.value.integer
+    fail = False
 
-    print(value_a % PRIME)
-    print(value)
+    for _ in range(count):
+        sim_val = await tb.expected_queue.get()
+        dut_val = await tb.output_queue.get()
 
-    await Timer(1, "us")
\ No newline at end of file
+        if sim_val != dut_val:
+            tb.log.info(f"{sim_val:x} -> {dut_val:x}")
+            fail = True
+
+    assert not fail
\ No newline at end of file