From 0dd85cf79dc41cc9aaa95a6336fcac71f48a0fc7 Mon Sep 17 00:00:00 2001
From: chriseth <chris@ethereum.org>
Date: Wed, 7 Feb 2024 16:56:27 +0100
Subject: [PATCH] split_bits utility function.

---
 std/binary.asm | 22 ++++++++++++++--------
 std/utils.asm  | 19 +++++++++++++++++++
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/std/binary.asm b/std/binary.asm
index 8dba7677cc..7fd7c6b59f 100644
--- a/std/binary.asm
+++ b/std/binary.asm
@@ -1,4 +1,5 @@
 use std::convert::int;
+use std::utils::split_bits;
 
 machine Binary(latch, operation_id) {
 
@@ -15,15 +16,20 @@ machine Binary(latch, operation_id) {
     col fixed latch(i) { if (i % 4) == 3 { 1 } else { 0 } };
     col fixed FACTOR(i) { 1 << (((i + 1) % 4) * 8) };
 
-    col fixed P_A(i) { i % 256 };
-    col fixed P_B(i) { (i >> 8) % 256 };
-    col fixed P_operation(i) { (i / (256 * 256)) % 3 };
+    // TOOD would be nice with destructuring assignment for arrays.
+    let inputs: (int -> int)[] = split_bits([2, 8, 8]);
+    let a = inputs[2];
+    let b = inputs[1];
+    let op = inputs[0];
+    col fixed P_A(i) { a(i) };
+    col fixed P_B(i) { b(i) };
+    col fixed P_operation(i) { op(i)};
     col fixed P_C(i) {
-        match P_operation(i) {
-            0 => int(P_A(i)) & int(P_B(i)),
-            1 => int(P_A(i)) | int(P_B(i)),
-            2 => int(P_A(i)) ^ int(P_B(i)),
-        } & 0xff
+        match op(i) {
+            0 => a(i) & b(i),
+            1 => a(i) | b(i),
+            2 => a(i) ^ b(i),
+        }
     };
 
     col witness A_byte;
diff --git a/std/utils.asm b/std/utils.asm
index 0a54aa2f97..63e301a84e 100644
--- a/std/utils.asm
+++ b/std/utils.asm
@@ -23,3 +23,22 @@ let unchanged_until = |c, latch| (c' - c) * (1 - latch) = 0;
 
 /// Evaluates to a constraint that forces `c` to be either 0 or 1.
 let force_bool: expr -> constr = |c| c * (1 - c) = 0;
+
+/// Returns an array of functions such that the range of the `i`th function is exactly the
+/// `bits[i]`-bit numbers (i.e. 0 until 2**bits[i] - 1, inclusive), such that all combinations
+/// of values of these functions appear.
+/// In other words, each of the functions returns a different `bits[i]`-bit slice of the input,
+/// shifted to the right so that the bit slice starts at the first bit.
+/// This function is useful for combined range checks or building the inputs for function
+/// that is implemented in a lookup.
+/// See binary.asm for an example.
+let split_bits: int[] -> (int -> int)[] = |bits| split_bits_internal(0, std::array::len(bits), bits);
+
+let split_bits_internal: int, int[] -> (int -> int)[] = |used_bits, len, bits|
+    if len == 0 {
+        // We could assert here that the degree is at least 2**used_bits
+        []
+    } else {
+        split_bits_internal(used_bits + bits[len - 1], len - 1, bits)
+            + [|i| (i >> used_bits) % (1 << bits[len - 1])]
+    };