From d6a062baa06aa3c656c3bc28978bea55ed24e674 Mon Sep 17 00:00:00 2001
From: Byron Lathi <bslathi19@gmail.com>
Date: Tue, 28 Oct 2025 21:59:28 -0700
Subject: [PATCH] Make modular mult work

---
 ChaCha20_Poly1305_64/sim/modulo_theory.py     |  4 ----
 .../sim/poly1305_friendly_modular_mult.py     |  6 ++++--
 .../src/poly1305_friendly_modular_mult.sv     | 20 +++++++++++--------
 3 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/ChaCha20_Poly1305_64/sim/modulo_theory.py b/ChaCha20_Poly1305_64/sim/modulo_theory.py
index dd8929c..f311095 100644
--- a/ChaCha20_Poly1305_64/sim/modulo_theory.py
+++ b/ChaCha20_Poly1305_64/sim/modulo_theory.py
@@ -57,10 +57,6 @@ def friendly_modular_mult(value_a: int, value_b: int) -> int:
 
     mods = [friendly_modulo(prod, 26*i) for i, prod in enumerate(prods)]
 
-    if sum(mods) >= 2*PRIME:
-        print("Saw greater than 2x prime!!!")
-
-
     mod_sum = friendly_modulo(sum(mods), 0)
 
     return mod_sum
diff --git a/ChaCha20_Poly1305_64/sim/poly1305_friendly_modular_mult.py b/ChaCha20_Poly1305_64/sim/poly1305_friendly_modular_mult.py
index 180f00b..f521981 100644
--- a/ChaCha20_Poly1305_64/sim/poly1305_friendly_modular_mult.py
+++ b/ChaCha20_Poly1305_64/sim/poly1305_friendly_modular_mult.py
@@ -8,6 +8,8 @@ from cocotb.queue import Queue
 
 from cocotbext.axi import AxiStreamBus, AxiStreamSource
 
+from modulo_theory import friendly_modular_mult
+
 import random
 
 PRIME = 2**130-5
@@ -48,7 +50,7 @@ class TB:
 
     async def write_input(self, data: int, h: int):
         await self.input_queue.put((data, h))
-        await self.expected_queue.put((data * h) % PRIME)
+        await self.expected_queue.put(friendly_modular_mult(h, data))
 
     async def run_input(self):
         while True:
@@ -76,7 +78,7 @@ async def test_sanity(dut):
 
     await tb.cycle_reset()
 
-    count = 1
+    count = 16
 
     for _ in range(count):
         await tb.write_input(random.randint(1,2**128-1), random.randint(0, 2**130-6))
diff --git a/ChaCha20_Poly1305_64/src/poly1305_friendly_modular_mult.sv b/ChaCha20_Poly1305_64/src/poly1305_friendly_modular_mult.sv
index 7a9dc1e..7b50b51 100644
--- a/ChaCha20_Poly1305_64/src/poly1305_friendly_modular_mult.sv
+++ b/ChaCha20_Poly1305_64/src/poly1305_friendly_modular_mult.sv
@@ -21,6 +21,7 @@ logic [2:0] state_counter, state_counter_next;
 logic [2:0] state_counter_p [4];
 
 logic [ACC_WIDTH-1:0] accumulator, accumulator_next;    // accumulator is outgoing
+logic [ACC_WIDTH:0] accumulator_intermediate;
 
 logic [DATA_WIDTH-1:0] data, data_next;
 logic [ACC_WIDTH-1:0] h, h_next;    // h is incoming
@@ -61,26 +62,29 @@ always_comb begin
 
     accumulator_next = '0;
     mult_product_next = '0;
+    accumulator_intermediate = '0;
 
 
+    if (state_counter < 3'h5) begin
+        mult_product_next = h[state_counter*26 +: 26] * data;
+        state_counter_next = state_counter + 1;
+    end
+
     if (state_counter >= 3'h4 && i_valid) begin
         data_next = i_data;
         h_next = i_accumulator;
         state_counter_next = '0;
     end
 
-    if (state_counter < 3'h5) begin
-        mult_product_next = h[state_counter*26 +: 26] * data;
-        state_counter_next = state_counter + 1;
-    end
-
     if (state_counter_p[3] == '0) begin
         accumulator_next = modulo_result;
     end else begin
-        if (accumulator + modulo_result > PRIME) begin
-            accumulator_next = accumulator + modulo_result - PRIME;
+        accumulator_intermediate = accumulator + modulo_result;
+        if (accumulator_intermediate[130]) begin
+            // if we wrapped around, we need to add 5
+            accumulator_next = accumulator_intermediate[129:0] + 5;
         end else begin
-            accumulator_next = accumulator + modulo_result;
+            accumulator_next = accumulator_intermediate[129:0];
         end
     end