recursion program: merge whir Merkle verification and root-of-unity exponanciation

Author: TomWambsgans

crates/rec_aggregation/hashing.py

@@ -259,67 +259,6 @@ def whir_do_1_merkle_level(b, state_in, path_chunk, state_out):
     return
 
 
-@inline
-def hash_and_verify_merkle_hint(leaf_position_nibbles, root, height, num_chunks):
-    # Hint and hash leaf
-    leaf_data = Array(num_chunks * DIGEST_LEN)
-    hint_witness("merkle_leaf", leaf_data)
-    leaf_hash = slice_hash_rtl(leaf_data, num_chunks)
-
-    # Hint and verify merkle path (processing 4 levels per nibble)
-    merkle_path = Array(height * DIGEST_LEN)
-    hint_witness("merkle_path", merkle_path)
-
-    states = Array((div_ceil(height, 4) - 1) * DIGEST_LEN)
-
-    # First full nibble: leaf_hash -> states[0]
-    match_range(leaf_position_nibbles[0], range(0, 16), lambda b: whir_do_4_merkle_levels(b, leaf_hash, merkle_path, states))
-
-    # Middle nibble chunks: states[k-1] -> states[k]
-    for k in unroll(1, div_ceil(height, 4) - 1):
-        match_range(
-            leaf_position_nibbles[k],
-            range(0, 16),
-            lambda b: whir_do_4_merkle_levels(b, states + (k - 1) * DIGEST_LEN, merkle_path + 4 * k * DIGEST_LEN, states + k * DIGEST_LEN),
-        )
-
-    # Last chunk -> root
-    if height % 4 == 0:
-        match_range(
-            leaf_position_nibbles[div_ceil(height, 4) - 1],
-            range(0, 16),
-            lambda b: whir_do_4_merkle_levels(
-                b, states + (div_ceil(height, 4) - 2) * DIGEST_LEN, merkle_path + 4 * (div_ceil(height, 4) - 1) * DIGEST_LEN, root
-            ),
-        )
-    elif height % 4 == 1:
-        match_range(
-            leaf_position_nibbles[(height - height % 4) / 4],
-            range(0, 16),
-            lambda b: whir_do_1_merkle_level(
-                b, states + (div_ceil(height, 4) - 2) * DIGEST_LEN, merkle_path + 4 * ((height - height % 4) / 4) * DIGEST_LEN, root
-            ),
-        )
-    elif height % 4 == 2:
-        match_range(
-            leaf_position_nibbles[(height - height % 4) / 4],
-            range(0, 16),
-            lambda b: whir_do_2_merkle_levels(
-                b, states + (div_ceil(height, 4) - 2) * DIGEST_LEN, merkle_path + 4 * ((height - height % 4) / 4) * DIGEST_LEN, root
-            ),
-        )
-    elif height % 4 == 3:
-        match_range(
-            leaf_position_nibbles[(height - height % 4) / 4],
-            range(0, 16),
-            lambda b: whir_do_3_merkle_levels(
-                b, states + (div_ceil(height, 4) - 2) * DIGEST_LEN, merkle_path + 4 * ((height - height % 4) / 4) * DIGEST_LEN, root
-            ),
-        )
-
-    return leaf_data
-
-
 def merkle_verif_batch(merkle_paths, leaves_digests, leave_positions, root, height, num_queries):
     match_range(
         height,

crates/rec_aggregation/utils.py

@@ -473,41 +473,116 @@ def checked_decompose_bits(a):
 
 
 @inline
-def checked_decompose_bits_and_compute_root_pow_const(a, domain_size):
-    # Hint 6 nibbles (4 bits each) + 1 top-7-bit value = 7 hints
+def whir_4_merkle_step_and_pow(v, state_in, path_chunk, state_out, power_shift):
+    whir_do_4_merkle_levels(v, state_in, path_chunk, state_out)
+    return ROOT ** (power_shift * v)
+
+
+@inline
+def whir_3_merkle_step_and_pow(v, state_in, path_chunk, state_out, power_shift):
+    whir_do_3_merkle_levels(v, state_in, path_chunk, state_out)
+    return ROOT ** (power_shift * (v % 8))
+
+
+@inline
+def whir_2_merkle_step_and_pow(v, state_in, path_chunk, state_out, power_shift):
+    whir_do_2_merkle_levels(v, state_in, path_chunk, state_out)
+    return ROOT ** (power_shift * (v % 4))
+
+
+@inline
+def whir_1_merkle_step_and_pow(v, state_in, path_chunk, state_out, power_shift):
+    whir_do_1_merkle_level(v, state_in, path_chunk, state_out)
+    return ROOT ** (power_shift * (v % 2))
+
+
+@inline
+def decompose_and_verify_merkle_query(a, domain_size, prev_root, num_chunks):
     nibbles = Array(6)
     top7: Imu
     hint_decompose_bits_merkle_whir(nibbles, top7, a, 4)
 
     for i in unroll(0, 6):
         assert nibbles[i] < 16
-
     assert top7 < 2**7
 
     partial_sum: Mut = nibbles[0]
     for i in unroll(1, 6):
         partial_sum += nibbles[i] * 16**i
-
     if top7 == 2**7 - 1:
         assert partial_sum == 0
-
     assert partial_sum + top7 * 2**24 == a
 
-    # Compute domain_generator^index
+    leaf_data = Array(num_chunks * DIGEST_LEN)
+    hint_witness("merkle_leaf", leaf_data)
+    leaf_hash = slice_hash_rtl(leaf_data, num_chunks)
+
+    merkle_path = Array(domain_size * DIGEST_LEN)
+    hint_witness("merkle_path", merkle_path)
+
+    n_nibbles = div_ceil(domain_size, 4)
+    states = Array((n_nibbles - 1) * DIGEST_LEN)
+
     prod: Mut = 1
-    for k in unroll(0, (domain_size - domain_size % 4) / 4):
-        nib_pow = match_range(nibbles[k], range(0, 16), lambda v: ROOT ** (2 ** (TWO_ADICITY - domain_size + 4 * k) * v))
+
+    # First nibble: leaf_hash -> states[0]
+    nib_pow = match_range(
+        nibbles[0],
+        range(0, 16),
+        lambda v: whir_4_merkle_step_and_pow(v, leaf_hash, merkle_path, states, 2 ** (TWO_ADICITY - domain_size)),
+    )
+    prod *= nib_pow
+
+    # Middle nibbles: states[k-1] -> states[k]
+    for k in unroll(1, n_nibbles - 1):
+        nib_pow = match_range(
+            nibbles[k],
+            range(0, 16),
+            lambda v: whir_4_merkle_step_and_pow(
+                v,
+                states + (k - 1) * DIGEST_LEN,
+                merkle_path + 4 * k * DIGEST_LEN,
+                states + k * DIGEST_LEN,
+                2 ** (TWO_ADICITY - domain_size + 4 * k),
+            ),
+        )
         prod *= nib_pow
 
-    if domain_size % 4 != 0:
-        edge_pow = match_range(
-            nibbles[(domain_size - domain_size % 4) / 4],
+    # Last nibble: states[-1] -> prev_root
+    last_k = n_nibbles - 1
+    last_state_in = states + (last_k - 1) * DIGEST_LEN
+    last_path = merkle_path + 4 * last_k * DIGEST_LEN
+    last_power_shift = 2 ** (TWO_ADICITY - domain_size + 4 * last_k)
+    if domain_size % 4 == 0:
+        nib_pow = match_range(
+            nibbles[last_k],
+            range(0, 16),
+            lambda v: whir_4_merkle_step_and_pow(v, last_state_in, last_path, prev_root, last_power_shift),
+        )
+        prod *= nib_pow
+    elif domain_size % 4 == 1:
+        nib_pow = match_range(
+            nibbles[last_k],
+            range(0, 16),
+            lambda v: whir_1_merkle_step_and_pow(v, last_state_in, last_path, prev_root, last_power_shift),
+        )
+        prod *= nib_pow
+    elif domain_size % 4 == 2:
+        nib_pow = match_range(
+            nibbles[last_k],
             range(0, 16),
-            lambda v: ROOT ** (2 ** (TWO_ADICITY - domain_size + 4 * ((domain_size - domain_size % 4) / 4)) * (v % 2 ** (domain_size % 4))),
+            lambda v: whir_2_merkle_step_and_pow(v, last_state_in, last_path, prev_root, last_power_shift),
         )
-        prod *= edge_pow
+        prod *= nib_pow
+    elif domain_size % 4 == 3:
+        nib_pow = match_range(
+            nibbles[last_k],
+            range(0, 16),
+            lambda v: whir_3_merkle_step_and_pow(v, last_state_in, last_path, prev_root, last_power_shift),
+        )
+        prod *= nib_pow
 
-    return nibbles, prod
+    return leaf_data, prod
 
 
 def checked_decompose_bits_small_value_const(to_decompose, n_bits: Const):

crates/rec_aggregation/whir.py

@@ -235,8 +235,7 @@ def decompose_and_verify_merkle_batch_with_height(num_queries, sampled, root, he
 
 def decompose_and_verify_merkle_batch_const(num_queries, sampled, root, height: Const, num_chunks: Const, circle_values, merkle_leaves):
     for i in range(0, num_queries):
-        nibbles, circle_values[i] = checked_decompose_bits_and_compute_root_pow_const(sampled[i], height)
-        merkle_leaves[i] = hash_and_verify_merkle_hint(nibbles, root, height, num_chunks)
+        merkle_leaves[i], circle_values[i] = decompose_and_verify_merkle_query(sampled[i], height, root, num_chunks)
     return