split_bits utility function.

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;

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])]
+    };