From 4d5dfbfef4baae9151e596366aa63f28ec3127f3 Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@proton.me>
Date: Tue, 21 Apr 2026 12:18:55 +0200
Subject: [PATCH] symbolic abstraction

---
 src/tea3/tea3model.py | 70 +++++++++++++++++++++++++++++++++++++++----
 1 file changed, 65 insertions(+), 5 deletions(-)

diff --git a/src/tea3/tea3model.py b/src/tea3/tea3model.py
index 3e5d4fa..eb8d4a5 100644
--- a/src/tea3/tea3model.py
+++ b/src/tea3/tea3model.py
@@ -5,13 +5,15 @@ from tea3.pretty_print import pretty_print, pretty_print_vec
 
 
 class Tea3Model:
-    def __init__(self):
+    def __init__(self, max_steps=20):
         self.F = GF(2)
+        self.step_count = 0
 
         names = (
             [f"x{i}{j}" for i in range(5) for j in range(8)] +
             [f"r{i}{j}" for i in range(5) for j in range(8)] +
-            [f"R{i}{j}" for i in range(8) for j in range(8)]
+            [f"R{i}{j}" for i in range(8) for j in range(8)] +
+            [f"f{s}_{i}{j}" for s in range(max_steps) for i in range(8) for j in range(8)]
         )
 
         name_string = ",".join(names)
@@ -22,6 +24,61 @@ class Tea3Model:
         self.r_bits = [list(self.v[40 + i*8 : 40 + (i+1)*8]) for i in range(5)]
         self.R_bits = [list(self.v[80 + i*8 : 80 + (i+1)*8]) for i in range(8)]
 
+        # Abstract variables
+        base = 80 + 64
+        self.fR_bits = [
+            [list(self.v[base + s*64 + i*8 : base + s*64 + i*8 + 8])
+             for i in range(8)]
+            for s in range(max_steps)
+        ]
+
+    def _split_poly(self, poly):
+        """
+        Split a polynomial into:
+          - R_f_part: monomials involving only 'R' or 'f' (abstract) variables
+          - xr_part:  monomials involving 'x' or 'r' variables
+
+        constant term is grouped with R_f_part when R_f_part is non-zero
+        """
+        zero = self.S.zero()
+        R_f_part = zero
+        xr_part = zero
+
+        has_const = bool(poly.constant_coefficient())
+
+        for monom in poly:
+            vars_in_term = monom.variables()
+            if not vars_in_term:
+                continue
+            families = {str(v)[0] for v in vars_in_term}
+            monom_poly = self.S(monom)
+            if families <= {'R', 'f'}:
+                R_f_part += monom_poly
+            else:
+                xr_part += monom_poly
+
+        if has_const:
+            if R_f_part != zero:
+                R_f_part += self.S.one()
+            else:
+                xr_part += self.S.one()
+
+        return R_f_part, xr_part
+
+    def _abstract_R(self):
+        """
+        Replace the R/f-dependent part of every R_bits[i][j] with an
+        abstract variable f{step}_{i}{j}, leaving only x and r terms explicit.
+        """
+        s = self.step_count
+        for i in range(8):
+            for j in range(8):
+                R_f_part, xr_part = self._split_poly(self.R_bits[i][j])
+                if R_f_part != self.S.zero():
+                    self.R_bits[i][j] = self.fR_bits[s][i][j] + xr_part
+                else:
+                    self.R_bits[i][j] = xr_part
+
     def step(self):
         R = self.R_bits.copy()
         x = self.x_bits.copy()
@@ -55,6 +112,9 @@ class Tea3Model:
         self.R_bits[1] = R0
         self.R_bits[0] = xor_vec(x0, xor_vec(R7, xor_vec(BP(R4), F32(R2, R1))))
 
+        self._abstract_R()
+        self.step_count += 1
+
         return R7
 
 def S(r):
@@ -99,9 +159,9 @@ def F32(X, Y):
     ]
 
 
-
-t = Tea3Model()
-for i in range(6):
+step = 5
+t = Tea3Model(step)
+for i in range(step):
     print("step "+str(i))
     t.step()
     print(pretty_print(t.R_bits[0][0]))