refactor: use functions and constraint comprehension

air-script/tests/bitwise/bitwise.air

@@ -4,7 +4,7 @@ public_inputs:
     stack_inputs: [16]
 
 trace_columns:
-    main: [s, a, b, a_limb[4], b_limb[4], zp, z, dummy]
+    main: [s, a, b, a0, a1, a2, a3, b0, b1, b2, b3, zp, z, dummy]
 
 periodic_columns:
     k0: [1, 0, 0, 0, 0, 0, 0, 0]
@@ -21,30 +21,27 @@ integrity_constraints:
     # Enforce that selector should stay the same throughout the cycle.
     enf k1 * (s' - s) = 0
 
-    # Enforce that the input is decomposed into valid bits
-    enf a_limb[0]^2 - a_limb[0] = 0
-    enf a_limb[1]^2 - a_limb[1] = 0
-    enf a_limb[2]^2 - a_limb[2] = 0
-    enf a_limb[3]^2 - a_limb[3] = 0
-    enf b_limb[0]^2 - b_limb[0] = 0
-    enf b_limb[1]^2 - b_limb[1] = 0
-    enf b_limb[2]^2 - b_limb[2] = 0
-    enf b_limb[3]^2 - b_limb[3] = 0
-
-    # Enforce that the value in the first row of column `a` of the current 8-row cycle should be 
-    # the aggregation of the decomposed bit columns `a_limb`.
-    let a_aggr = sum([2^i * a for (i, a) in (0..4, a_limb)])
-    enf k0 * (a - a_aggr) = 0
-    # Enforce that the value in the first row of column `b` of the current 8-row cycle should be 
-    # the aggregation of the decomposed bit columns `b_limb`.
-    let b_aggr = sum([2^i * b for (i, b) in (0..4, b_limb)])
-    enf k0 * (b - b_aggr) = 0
-
-    # Enforce that for all rows in an 8-row cycle, except for the last one, the values in a and b
-    # columns are increased by the values contained in the individual bit columns a_limb and 
-    # b_limb.
-    enf k1 * (a' - (a * 16 + a_aggr)) = 0
-    enf k1 * (b' - (b * 16 + b_aggr)) = 0
+    # Enforce that input is decomposed into valid bits
+    enf a0^2 - a0 = 0
+    enf a1^2 - a1 = 0
+    enf a2^2 - a2 = 0
+    enf a3^2 - a3 = 0
+    enf b0^2 - b0 = 0
+    enf b1^2 - b1 = 0
+    enf b2^2 - b2 = 0
+    enf b3^2 - b3 = 0
+
+    # Enforce that the values in the column a in the first row should be the aggregation of the
+    # decomposed bit columns a0..a3.
+    enf k0 * (a - (2^0 * a0 + 2^1 * a1 + 2^2 * a2 + 2^3 * a3)) = 0
+    # Enforce that the values in the column b in the first row should be the aggregation of the
+    # decomposed bit columns b0..b3.
+    enf k0 * (b - (2^0 * b0 + 2^1 * b1 + 2^2 * b2 + 2^3 * b3)) = 0
+
+    # Enforce that for all rows in an 8-row cycle except for the last one, the values in a and b
+    # columns are increased by the values contained in the individual bit columns a and b.
+    enf k1 * (a' - (a * 16 + 2^0 * a0 + 2^1 * a1 + 2^2 * a2 + 2^3 * a3)) = 0
+    enf k1 * (b' - (b * 16 + 2^0 * b0 + 2^1 * b1 + 2^2 * b2 + 2^3 * b3)) = 0
 
     # Enforce that in the first row, the aggregated output value of the previous row should be 0.
     enf k0 * zp = 0
@@ -55,12 +52,10 @@ integrity_constraints:
 
     # Enforce that for all rows the value in the z column is computed by multiplying the previous
     # output value (from the zp column in the current row) by 16 and then adding it to the bitwise
-    # operation applied to the row's set of bits of a_limb and b_limb. The entire constraint must 
-    # also be multiplied by the operation selector flag to ensure it is only applied for the 
-    # appropriate operation. The constraint for AND is enforced when s = 0 and the constraint for 
-    # XOR is enforced when s = 1. Because the selectors for the AND and XOR operations are mutually
+    # operation applied to the row's set of bits of a and b. The entire constraint must also be
+    # multiplied by the operation selector flag to ensure it is only applied for the appropriate
+    # operation. The constraint for AND is enforced when s = 0 and the constraint for XOR is
+    # enforced when s = 1. Because the selectors for the AND and XOR operations are mutually
     # exclusive, the constraints for different operations can be aggregated into the same result
     # indices.
-    let a_and_b = sum([2^i * a * b for (i, a, b) in (0..4, a_limb, b_limb)])
-    let a_xor_b = sum([2^i * (a + b - 2 * a * b) for (i, a, b) in (0..4, a_limb, b_limb)])
-    enf (1 - s) * (z - (zp * 16 + a_and_b)) + s * (z - (zp * 16 + a_xor_b)) = 0
+    enf (1 - s) * (z - (zp * 16 + 2^0 * a0 * b0 + 2^1 * a1 * b1 + 2^2 * a2 * b2 + 2^3 * a3 * b3)) + s * (z - (zp * 16 + 2^0 * (a0 + b0 - 2 * a0 * b0) + 2^1 * (a1 + b1 - 2 * a1 * b1) + 2^2 * (a2 + b2 - 2 * a2 * b2) + 2^3 * (a3 + b3 - 2 * a3 * b3))) = 0

constraints/miden-vm/bitwise.air

@@ -0,0 +1,70 @@
+def BitwiseAir
+
+public_inputs:
+    stack_inputs: [16]
+
+trace_columns:
+    main: [s, a, b, a_limb[4], b_limb[4], zp, z, dummy]
+
+periodic_columns:
+    k0: [1, 0, 0, 0, 0, 0, 0, 0]
+    k1: [1, 1, 1, 1, 1, 1, 1, 0]
+
+boundary_constraints:
+    # This is a dummy trace column to satisfy requirement of at least one boundary constraint.
+    enf dummy.first = 0
+
+integrity_constraints:
+    # Enforces that column must be binary
+    ev check_binary(main: [a]):
+        enf a^2 - a = 0
+
+    # Returns value aggregated from limbs
+    fn aggregate(limb: vector[4]) -> scalar:
+        let result = sum([2^i * a for (i, a) in (0..4, limb)])
+        result
+
+    # Enforce that selector must be binary
+    enf check_binary(s)
+
+    # Enforce that selector should stay the same throughout the cycle.
+    enf k1 * (s' - s) = 0
+
+    # Enforce that the input is decomposed into valid bits
+    enf check_binary(a_limb[i]) for i in 0..4
+    enf check_binary(b_limb[i]) for i in 0..4  
+
+    # Enforce that the value in the first row of column `a` of the current 8-row cycle should be 
+    # the aggregation of the decomposed bit columns `a_limb`.
+    let a_aggr = aggregate(a_limb)
+    enf k0 * (a - a_aggr) = 0
+
+    # Enforce that the value in the first row of column `b` of the current 8-row cycle should be 
+    # the aggregation of the decomposed bit columns `b_limb`.
+    let b_aggr = aggregate(b_limb)
+    enf k0 * (b - b_aggr) = 0
+
+    # Enforce that for all rows in an 8-row cycle, except for the last one, the values in a and b
+    # columns are increased by the values contained in the individual bit columns a_limb and 
+    # b_limb.
+    enf k1 * (a' - (a * 16 + a_aggr)) = 0
+    enf k1 * (b' - (b * 16 + b_aggr)) = 0
+
+    # Enforce that in the first row, the aggregated output value of the previous row should be 0.
+    enf k0 * zp = 0
+
+    # Enforce that for each row except the last, the aggregated output value must equal the
+    # previous aggregated output value in the next row.
+    enf k1 * (z - zp') = 0
+
+    # Enforce that for all rows the value in the z column is computed by multiplying the previous
+    # output value (from the zp column in the current row) by 16 and then adding it to the bitwise
+    # operation applied to the row's set of bits of a_limb and b_limb. The entire constraint must 
+    # also be multiplied by the operation selector flag to ensure it is only applied for the 
+    # appropriate operation. The constraint for AND is enforced when s = 0 and the constraint for 
+    # XOR is enforced when s = 1. Because the selectors for the AND and XOR operations are mutually
+    # exclusive, the constraints for different operations can be aggregated into the same result
+    # indices.
+    let a_and_b = sum([2^i * a * b for (i, a, b) in (0..4, a_limb, b_limb)])
+    let a_xor_b = sum([2^i * (a + b - 2 * a * b) for (i, a, b) in (0..4, a_limb, b_limb)])
+    enf (1 - s) * (z - (zp * 16 + a_and_b)) + s * (z - (zp * 16 + a_xor_b)) = 0