abstract pool

src/tea3/cli.py

@@ -29,7 +29,7 @@ def main() -> None:
     print(f"\nRunning {steps} step(s), watching R[{reg}][{bit}]")
     print("-" * 50)
 
-    model = Tea3Model(max_steps=steps)
+    model = Tea3Model()
 
     for i in range(steps):
         model.step()

src/tea3/tea3model.py

@@ -1,11 +1,13 @@
 from sage.all import GF, BooleanPolynomialRing
+from functools import reduce
+from operator import mul
 
 from tea3.constants import TEA3_SBOX, T_F1, T_F2
 from tea3.pretty_print import pretty_print, pretty_print_vec
 
 
 class Tea3Model:
-    def __init__(self, max_steps=20):
+    def __init__(self, max_abstract=40960):
         self.F = GF(2)
         self.step_count = 0
 
@@ -13,7 +15,7 @@ class Tea3Model:
             [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"f{s}_{i}{j}" for s in range(max_steps) for i in range(8) for j in range(8)]
+            [f"g{n}"    for n in range(max_abstract)]
         )
 
         name_string = ",".join(names)
@@ -24,60 +26,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
+        self._g_pool  = list(self.v[80 + 64:])
-        base = 80 + 64
+        self._g_index = 0
-        self.fR_bits = [
+        self.abstractions = {}
-            [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):
+    def _alloc(self, step, i, j, xr_key):
-        """
+        if self._g_index >= len(self._g_pool):
-        Split a polynomial into:
+            raise RuntimeError("Abstract variable pool exhausted — increase max_abstract.")
-          - R_f_part: monomials involving only 'R' or 'f' (abstract) variables
+        var = self._g_pool[self._g_index]
-          - xr_part:  monomials involving 'x' or 'r' variables
+        self.abstractions[(step, i, j, xr_key)] = (var, self._g_index)
-
+        self._g_index += 1
-        constant term is grouped with R_f_part when R_f_part is non-zero
+        return var
-        """
-        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
+        one  = self.S.one()
+        zero = self.S.zero()
+
         for i in range(8):
             for j in range(8):
-                R_f_part, xr_part = self._split_poly(self.R_bits[i][j])
+                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
+                groups  = {}
+                pure_xr = zero
+                const   = one if bool(poly.constant_coefficient()) else zero
+
+                for monom in poly:
+                    term_vars = monom.variables()
+                    if not term_vars:
+                        continue
+
+                    xr_vars = [v for v in term_vars if str(v)[0] in ('x', 'r')]
+                    Rf_vars = [v for v in term_vars if str(v)[0] in ('R', 'f', 'g')]  # 'g' added
+
+                    xr_mono = reduce(mul, (self.S(v) for v in xr_vars), one)
+                    Rf_mono = reduce(mul, (self.S(v) for v in Rf_vars), one)
+                    xr_key  = frozenset(str(v) for v in xr_vars)
+
+                    if not Rf_vars:
+                        pure_xr += xr_mono
                     else:
-                    self.R_bits[i][j] = xr_part
+                        if xr_key not in groups:
+                            groups[xr_key] = {'xr_mono': xr_mono, 'Rf_sum': zero}
+                        groups[xr_key]['Rf_sum'] += Rf_mono
+
+                result = pure_xr + const
+
+                for xr_key, bucket in groups.items():
+                    Rf_sum  = bucket['Rf_sum']
+                    xr_mono = bucket['xr_mono']
+                    if Rf_sum == zero:
+                        continue
+                    g = self._alloc(s, i, j, xr_key)
+                    result += xr_mono * g
+
+                self.R_bits[i][j] = result
 
     def step(self):
         R = self.R_bits.copy()