variable change exhaustive search

src/tea3/variable_search.py

@@ -0,0 +1,40 @@
+from itertools import product as iproduct
+from tea3.tea3model import Tea3Model
+from tea3.pretty_print import pretty_print
+
+def apply_variable_change(model, coeffs):
+    orig_x = [list(model.v[i*8 : (i+1)*8]) for i in range(5)]
+    orig_y = model.y_bits
+
+    subs = {}
+
+    for i in range(1, 5):
+        for j in range(8):
+            subs[orig_x[i][j]] = orig_y[i][j]
+
+    for j in range(8):
+        new_expr = orig_y[0][j]
+        for i, ai in enumerate(coeffs, start=1):
+            if ai:
+                new_expr = new_expr + orig_y[i][j]
+        subs[orig_x[0][j]] = new_expr
+
+    new_R = []
+    for i in range(8):
+        row = [poly.subs(subs) for poly in model.R_bits[i]]
+        new_R.append(row)
+    return new_R
+
+def run_exhaustive(steps: int, target_reg: int = 0):
+    for idx, coeffs in enumerate(iproduct([0, 1], repeat=4)):
+        model = Tea3Model()
+        for _ in range(steps):
+            model.step()
+
+        new_R = apply_variable_change(model, coeffs)
+
+        label = "".join(map(str, coeffs))
+        print(f"\n[{idx:02d}] (a1,a2,a3,a4) = {label}")
+        for j in range(8):
+            poly = new_R[target_reg][j]
+            print(f"  R[{target_reg}][{j}] = {pretty_print(poly)}")