From 781ec5960b43ed62cd2eba85ca48907a4cbdfbfe Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@proton.me>
Date: Tue, 7 Jul 2026 12:06:41 +0200
Subject: [PATCH] refactor, move set_bits to new file

---
 src/tea3/cli.py   | 46 +++------------------------------------------
 src/tea3/utils.py | 48 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 43 deletions(-)
 create mode 100644 src/tea3/utils.py

diff --git a/src/tea3/cli.py b/src/tea3/cli.py
index f3ff561..96673d8 100644
--- a/src/tea3/cli.py
+++ b/src/tea3/cli.py
@@ -1,6 +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.utils import set_bits
 from tea3.variable_search import run_exhaustive, run_exhaustive_staircase, run_exhaustive_staircase2, run_exhaustive_bp
 from tea3.sbox import run_sbox
 from tea3.variable_xor import run_variable_xor, run_exhaustive_xor
@@ -38,47 +39,6 @@ def run_classic_cli():
     print("\n" + "=" * 50)
     print("Done.")
 
-def _set_bits(model, bit_names, value):
-    if value not in (0, 1):
-        raise ValueError("value must be 0 or 1")
-
-    const = model.S.one() if value else model.S.zero()
-
-    for raw in bit_names:
-        name = raw.strip()
-        if not name:
-            continue
-
-        if len(name) != 3 or not name[1:].isdigit():
-            raise ValueError(
-                f"Invalid bit name '{name}'. Expected formats like R00, x00, r10."
-            )
-
-        prefix = name[0]
-        reg = int(name[1])
-        bit = int(name[2])
-
-        if prefix == "R":
-            if not (0 <= reg <= 7 and 0 <= bit <= 7):
-                raise ValueError(f"Invalid R bit '{name}'. Use R00..R77.")
-            model.R_bits[reg][bit] = const
-
-        elif prefix == "x":
-            if not (0 <= reg <= 4 and 0 <= bit <= 7):
-                raise ValueError(f"Invalid x bit '{name}'. Use x00..x47.")
-            model.x_bits[reg][bit] = const
-
-        elif prefix == "r":
-            if not (0 <= reg <= 4 and 0 <= bit <= 7):
-                raise ValueError(f"Invalid r bit '{name}'. Use r00..r47.")
-            model.r_bits[reg][bit] = const
-
-        else:
-            raise ValueError(
-                f"Invalid bit name '{name}'. Expected R, x, or r."
-            )
-
-
 def run_advanced_cli():
     print("\nR registers are indexed 0–7; bits within each register are 0–7.")
     print("x and r registers are indexed 0–4; bits within each register are 0–7.")
@@ -99,8 +59,8 @@ def run_advanced_cli():
 
     try:
         # 1 overrides 0
-        _set_bits(model, zero_bits, 0)
-        _set_bits(model, one_bits, 1)
+        set_bits(model, zero_bits, 0)
+        set_bits(model, one_bits, 1)
     except ValueError as e:
         print(f"Error: {e}")
         return
diff --git a/src/tea3/utils.py b/src/tea3/utils.py
new file mode 100644
index 0000000..c5ffe9c
--- /dev/null
+++ b/src/tea3/utils.py
@@ -0,0 +1,48 @@
+from sage.all import GF, BooleanPolynomialRing
+from functools import reduce
+from operator import mul
+
+def set_bits(model, bit_names, value):
+    """
+    Set the given bits to a constant value (0 or 1) in the model.
+    bit_names: list of strings like 'R00', 'x12', 'r47'.
+    value: 0 or 1.
+    """
+    if value not in (0, 1):
+        raise ValueError("value must be 0 or 1")
+
+    const = model.S.one() if value else model.S.zero()
+
+    for raw in bit_names:
+        name = raw.strip()
+        if not name:
+            continue
+
+        if len(name) != 3 or not name[1:].isdigit():
+            raise ValueError(
+                f"Invalid bit name '{name}'. Expected formats like R00, x00, r10."
+            )
+
+        prefix = name[0]
+        reg = int(name[1])
+        bit = int(name[2])
+
+        if prefix == "R":
+            if not (0 <= reg <= 7 and 0 <= bit <= 7):
+                raise ValueError(f"Invalid R bit '{name}'. Use R00..R77.")
+            model.R_bits[reg][bit] = const
+
+        elif prefix == "x":
+            if not (0 <= reg <= 4 and 0 <= bit <= 7):
+                raise ValueError(f"Invalid x bit '{name}'. Use x00..x47.")
+            model.x_bits[reg][bit] = const
+
+        elif prefix == "r":
+            if not (0 <= reg <= 4 and 0 <= bit <= 7):
+                raise ValueError(f"Invalid r bit '{name}'. Use r00..r47.")
+            model.r_bits[reg][bit] = const
+
+        else:
+            raise ValueError(
+                f"Invalid bit name '{name}'. Expected R, x, or r."
+            )