variable change exhaustive search

src/tea3/cli.py

@@ -1,5 +1,6 @@
 from tea3.pretty_print import pretty_print
 from tea3.tea3model import Tea3Model
+from tea3.variable_search import run_exhaustive
 
 
 def prompt_int(message: str, lo: int, hi: int) -> int:
@@ -15,28 +16,28 @@ def prompt_int(message: str, lo: int, hi: int) -> int:
         print(f"Value must be between {lo} and {hi} (inclusive).")
 
 
-def main() -> None:
+def prompt_choice(message: str, choices: set[int]) -> int:
-    print("=" * 50)
+    choices_str = ", ".join(map(str, sorted(choices)))
-    print(" Tea3 Model ")
+    while True:
-    print("=" * 50)
+        raw = input(message).strip()
+        try:
+            value = int(raw)
+        except ValueError:
+            print("Please enter a number.")
+            continue
+        if value in choices:
+            return value
+        print(f"Value must be one of: {choices_str}")
 
-    steps = prompt_int("\nHow many steps do you want to run? (1–100): ", 1, 100)
 
+def run_classic_cli():
     print("\nR registers are indexed 0–7; bits within each register are 0–7.")
     print("Enter -1 to print all registers, or all bits.")
 
+    steps = prompt_int("How many steps do you want to run? (1–100): ", 1, 100)
     reg = prompt_int("Which R register do you want to inspect? (-1 or 0–7): ", -1, 7)
     bit = prompt_int("Which bit of that register?              (-1 or 0–7): ", -1, 7)
 
-    if reg == -1 and bit == -1:
-        print(f"\nRunning {steps} step(s), watching all R registers and all bits")
-    elif reg == -1:
-        print(f"\nRunning {steps} step(s), watching all R registers, bit [{bit}]")
-    elif bit == -1:
-        print(f"\nRunning {steps} step(s), watching R[{reg}] and all bits")
-    else:
-        print(f"\nRunning {steps} step(s), watching R[{reg}][{bit}]")
-
     print("-" * 50)
 
     model = Tea3Model()
@@ -59,4 +60,28 @@ def main() -> None:
     print("Done.")
 
 
+def main():
+    print("=" * 50)
+    print(" Tea3 Model ")
+    print("=" * 50)
+
+    print("\nChoose a mode:")
+    print("  1) Classic inspection")
+    print("  2) Exhaustive variable-change search")
+
+    mode = prompt_choice("Your choice (1 or 2): ", {1, 2})
+
+    if mode == 1:
+        run_classic_cli()
+    else:
+        steps = prompt_int("How many steps? (1–100): ", 1, 100)
+        target_reg = prompt_int("Target register (0–7): ", 0, 7)
+
+        print("-" * 50)
+        run_exhaustive(steps, target_reg)
+
+        print("\n" + "=" * 50)
+        print("Done.")
+
+
 main()

src/tea3/tea3model.py

@@ -12,10 +12,11 @@ class Tea3Model:
         self.step_count = 0
 
         names = (
-            [f"x{i}{j}" for i in range(5) for j in range(8)] +
+            [f"x{i}{j}" for i in range(5) for j in range(8)] +  # 0–39
-            [f"r{i}{j}" for i in range(5) for j in range(8)] +
+            [f"y{i}{j}" for i in range(5) for j in range(8)] +  # 40–79
-            [f"R{i}{j}" for i in range(8) for j in range(8)] +
+            [f"r{i}{j}" for i in range(5) for j in range(8)] +  # 80–119
-            ["g"]
+            [f"R{i}{j}" for i in range(8) for j in range(8)] +  # 120–183
+            ["g"]                                               # 184
         )
 
         name_string = ",".join(names)
@@ -23,13 +24,12 @@ class Tea3Model:
         self.v = self.S.gens()
 
         self.x_bits = [list(self.v[i*8       : (i+1)*8      ]) for i in range(5)]
-        self.r_bits = [list(self.v[40 + i*8 : 40 + (i+1)*8]) for i in range(5)]
+        self.y_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)]
+        self.r_bits = [list(self.v[80 + i*8  : 80 + (i+1)*8]) for i in range(5)]
-
+        self.R_bits = [list(self.v[120 + i*8 : 120 + (i+1)*8]) for i in range(8)]
         self.g = self.v[-1]
 
     def _abstract_R(self):
-        s    = self.step_count
         one  = self.S.one()
         zero = self.S.zero()
 
@@ -38,7 +38,7 @@ class Tea3Model:
                 poly = self.R_bits[i][j]
 
                 groups  = {}
-                pure_xr = zero
+                pure_xyr = zero
                 const    = one if bool(poly.constant_coefficient()) else zero
 
                 for monom in poly:
@@ -46,20 +46,20 @@ class Tea3Model:
                     if not term_vars:
                         continue
 
-                    xr_vars = [v for v in term_vars if str(v)[0] in ('x', 'r')]
+                    xyr_vars = [v for v in term_vars if str(v)[0] in ('x', 'y', 'r')]
-                    Rf_vars = [v for v in term_vars if str(v)[0] in ('R', 'f', 'g')]
+                    Rg_vars  = [v for v in term_vars if str(v)[0] in ('R', 'g')]
 
-                    xr_mono = reduce(mul, (self.S(v) for v in xr_vars), one)
+                    xyr_mono = reduce(mul, (self.S(v) for v in xyr_vars), one)
-                    xr_key  = frozenset(str(v) for v in xr_vars)
+                    xyr_key  = frozenset(str(v) for v in xyr_vars)
 
-                    if not Rf_vars:
+                    if not Rg_vars:
-                        pure_xr += xr_mono
+                        pure_xyr += xyr_mono
                     else:
-                        groups[xr_key] = xr_mono
+                        groups[xyr_key] = xyr_mono
 
-                result = pure_xr + const
+                result = pure_xyr + const
-                for xr_key, xr_mono in groups.items():
+                for xyr_key, xyr_mono in groups.items():
-                    result += xr_mono * self.g
+                    result += xyr_mono * self.g
 
                 self.R_bits[i][j] = result
 
@@ -101,6 +101,7 @@ class Tea3Model:
 
         return R7
 
+
 def S(r):
     # placeholder
     return r

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)}")