Files
tea3-py/src/tea3/variable_search.py
T

84 lines
2.4 KiB
Python

from itertools import product as iproduct
from tea3.tea3model import Tea3Model
from tea3.pretty_print import pretty_print
def count_monomials(poly):
s = pretty_print(poly).strip()
if not s or s == "0":
return 0
return sum(1 for term in s.split("+") if term.strip())
def apply_variable_change(model, coeffs):
"""
Variable change inside each 8-bit register:
x_{r,0} = y_{r,0} + a1*y_{r,1} + ... + a7*y_{r,7}
x_{r,j} = y_{r,j} for j = 1..7
for every register r in {0,...,4}.
"""
orig_x = [list(model.v[r * 8 : (r + 1) * 8]) for r in range(5)]
orig_y = model.y_bits
subs = {}
for r in range(5):
for j in range(1, 8):
subs[orig_x[r][j]] = orig_y[r][j]
new_expr = orig_y[r][0]
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):
model = Tea3Model()
for _ in range(steps):
model.step()
snapshot = model
if target_bit == -1:
bits = list(range(8))
else:
if not (0 <= target_bit < 8):
raise ValueError("target_bit must be in [0, 7] or -1")
bits = [target_bit]
best_coeffs: tuple | None = None
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}")
total = 0
for j in bits:
poly = new_R[target_reg][j]
pp = pretty_print(poly)
print(f" R[{target_reg}][{j}] = {pp}")
total += count_monomials(poly)
if best_total == -1 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" 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(poly)}")