variable change abstraction

src/tea3/cli.py

@@ -1,7 +1,7 @@
 from tea3.pretty_print import pretty_print
 from tea3.cliutils import prompt_int, prompt_choice, prompt_list
 from tea3.tea3model import Tea3Model
-from tea3.variable_search import run_exhaustive
+from tea3.variable_search import run_exhaustive, run_exhaustive_staircase
 from tea3.sbox import run_sbox
 from tea3.variable_xor import run_variable_xor, run_exhaustive_xor
 from tea3.f31f32 import run_f31f32
@@ -37,6 +37,12 @@ def run_classic_cli():
     print("Done.")
 
 def run_exhaustive_cli():
+    print("\nChoose the variable-change family:")
+    print("  1) First-row matrix")
+    print("  2) Staircase matrix")
+
+    family = prompt_choice("Your choice (1 or 2): ", {1, 2})
+
     print("\nR registers are indexed 0–7; bits within each register are 0–7.")
     print("Enter -1 to print all bits in the chosen register.")
 
@@ -45,7 +51,11 @@ def run_exhaustive_cli():
     target_bit = prompt_int("Target bit (-1 or 0–7): ", -1, 7)
 
     print("-" * 50)
+
+    if family == 1:
         run_exhaustive(steps, target_reg, target_bit)
+    else:
+        run_exhaustive_staircase(steps, target_reg, target_bit)
 
     print("\n" + "=" * 50)
     print("Done.")

src/tea3/variable_search.py

@@ -1,74 +1,144 @@
 from itertools import product as iproduct
+
+from sage.all import GF, Matrix, identity_matrix
+
 from tea3.tea3model import Tea3Model
 from tea3.pretty_print import pretty_print
 
+# utils
 
-def count_monomials(poly):
+def count_monomials(poly) -> int:
     return len(poly.monomials())
 
 
-def apply_variable_change(model, coeffs):
+def gf2_matrix_inverse(M_binary):
+    M = Matrix(GF(2), M_binary)
+    if M.is_singular():
+        raise ValueError("Matrix is singular over GF(2).")
+    return M.inverse()
+
+
+def apply_variable_change_matrix(model, M_binary):
+    M_inv = gf2_matrix_inverse(M_binary)
+
     orig_x = [list(model.v[r * 8 : (r + 1) * 8]) for r in range(5)]
     orig_y = model.y_bits
+    ring = orig_y[0][0].parent()
 
     subs = {}
-
     for r in range(5):
-        for j in range(1, 8):
+        for j in range(8):
-            subs[orig_x[r][j]] = orig_y[r][j]
+            terms = [orig_y[r][k] for k in range(8) if M_inv[j, k] == 1]
+            if terms:
+                expr = sum(terms[1:], terms[0])
+            else:
+                expr = ring(0)
+            subs[orig_x[r][j]] = expr
 
-        new_expr = orig_y[r][0]
+    return [[poly.subs(subs) for poly in model.R_bits[i]] for i in range(8)]
-        for j, ai in enumerate(coeffs, start=1):
-            if ai:
-                new_expr += orig_y[r][j]
-        subs[orig_x[r][0]] = 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, target_bit: int = -1):
+# matrices
+
+def make_first_row_matrix(coeffs):
+    """
+    | 1   a1  a2  a3  a4   a5   a6   a7 |
+    | 0   1   0   0   0    0    0    0  |
+    | 0   0   1   0   0    0    0    0  |
+    | 0   0   0   1   0    0    0    0  |
+    | 0   0   0   0   1    0    0    0  |
+    | 0   0   0   0   0    1    0    0  |
+    | 0   0   0   0   0    0    1    0  |
+    | 0   0   0   0   0    0    0    1  |
+    """
+    M = identity_matrix(GF(2), 8)
+    for j, a in enumerate(coeffs, start=1):
+        M[0, j] = GF(2)(a)
+    return M
+
+
+def make_staircase_matrix(coeffs):
+    """
+    | 1   a1  0   0   0   0   0   0  |
+    | 0   1   a2  0   0   0   0   0  |
+    | 0   0   1   a3  0   0   0   0  |
+    | 0   0   0   1   a4  0   0   0  |
+    | 0   0   0   0   1   a5  0   0  |
+    | 0   0   0   0   0   1   a6  0  |
+    | 0   0   0   0   0   0   1   a7 |
+    | a8  0   0   0   0   0   0   1  |
+    """
+    M = identity_matrix(GF(2), 8)
+    for i, a in enumerate(coeffs[:7]):
+        M[i, i + 1] = GF(2)(a)
+    M[7, 0] = GF(2)(coeffs[7])
+    return M
+
+# wrappers
+
+def run_exhaustive_generic(
+    steps: int,
+    make_matrix,
+    n_params: int,
+    *,
+    param_label: str = "a",
+    target_reg: int = 0,
+    target_bit: int = -1,
+    verbose: bool = True,
+):
+    if target_bit not in range(-1, 8):
+        raise ValueError("target_bit must be in [0, 7] or -1.")
+    bits = list(range(8)) if target_bit == -1 else [target_bit]
+
     model = Tea3Model()
     for _ in range(steps):
         model.step()
-    snapshot = model
 
-    if target_bit == -1:
+    best_coeffs = None
-        bits = list(range(8))
+    best_total  = None
-    else:
+    best_R      = None
-        if not (0 <= target_bit < 8):
+    pnames      = ", ".join(f"{param_label}{i}" for i in range(1, n_params + 1))
-            raise ValueError("target_bit must be in [0, 7] or -1")
-        bits = [target_bit]
 
-    best_coeffs: tuple | None = None
+    for idx, coeffs in enumerate(iproduct([0, 1], repeat=n_params)):
-    best_total: int = -1
-    best_R: list | None = None
-
-    for idx, coeffs in enumerate(iproduct([0, 1], repeat=7)):
-        new_R = apply_variable_change(snapshot, coeffs)
         label = "".join(map(str, coeffs))
-        print(f"\n[{idx:03d}] (a1,a2,a3,a4,a5,a6,a7) = {label}")
+        try:
+            new_R = apply_variable_change_matrix(model, make_matrix(coeffs))
+        except ValueError:
+            if verbose:
+                print(f"[{idx:03d}] ({pnames}) = {label}  [singular - skipped]")
+            continue
 
-        total = 0
+        total = sum(count_monomials(new_R[target_reg][j]) for j in bits)
+
+        if verbose:
+            print(f"\n[{idx:03d}] ({pnames}) = {label}")
             for j in bits:
-            poly = new_R[target_reg][j]
+                print(f"  R[{target_reg}][{j}] = "
-            pp = pretty_print(poly)
+                      f"{pretty_print(new_R[target_reg][j])}")
-            print(f"  R[{target_reg}][{j}] = {pp}")
-            total += count_monomials(poly)
 
-        if best_total == -1 or total < best_total:
+        if best_total is None or total < best_total:
             best_total  = total
             best_coeffs = coeffs
             best_R      = new_R
 
     print("\nBest variable change (fewest total monomials)")
-    print(f"  (a1,a2,a3,a4,a5,a6,a7) = ({', '.join(map(str, best_coeffs))})")
+    print(f"  ({pnames}) = ({', '.join(map(str, best_coeffs))})")
     print(f"  Total monomials : {best_total}")
-    print("  Polynomial:\n")
     for j in bits:
-        poly = best_R[target_reg][j]
+        print(f"  R[{target_reg}][{j}] = {pretty_print(best_R[target_reg][j])}")
-        print(f"R[{target_reg}][{j}] = {pretty_print(poly)}")
+
+    return best_coeffs, best_total, best_R
+
+
+def run_exhaustive(steps: int, target_reg: int = 0, target_bit: int = -1):
+    return run_exhaustive_generic(
+        steps, make_first_row_matrix, n_params=7,
+        target_reg=target_reg, target_bit=target_bit,
+    )
+
+
+def run_exhaustive_staircase(steps: int, target_reg: int = 0, target_bit: int = -1):
+    return run_exhaustive_generic(
+        steps, make_staircase_matrix, n_params=8,
+        target_reg=target_reg, target_bit=target_bit,
+    )