improve `hint_decompose_bits_merkle_whir

Author: TomWambsgans

crates/lean_vm/src/isa/hint.rs

@@ -130,7 +130,7 @@ impl CustomHint {
     pub fn n_args(&self) -> usize {
         match self {
             Self::DecomposeBitsXMSS => 5,
-            Self::DecomposeBitsMerkleWhir => 4,
+            Self::DecomposeBitsMerkleWhir => 3,
             Self::DecomposeBits => 4,
             Self::LessThan => 3,
             Self::Log2Ceil => 2,
@@ -166,8 +166,8 @@ impl CustomHint {
             }
             Self::DecomposeBitsMerkleWhir => {
                 let decomposed_ptr = args[0].read_value(ctx.memory, ctx.fp)?.to_usize();
-                let value = args[2].read_value(ctx.memory, ctx.fp)?.to_usize();
-                let chunk_size = args[3].read_value(ctx.memory, ctx.fp)?.to_usize();
+                let value = args[1].read_value(ctx.memory, ctx.fp)?.to_usize();
+                let chunk_size = args[2].read_value(ctx.memory, ctx.fp)?.to_usize();
                 assert!(24_usize.is_multiple_of(chunk_size));
                 let mut memory_index_decomposed = decomposed_ptr;
                 #[allow(clippy::explicit_counter_loop)]
@@ -176,8 +176,6 @@ impl CustomHint {
                     ctx.memory.set(memory_index_decomposed, value)?;
                     memory_index_decomposed += 1;
                 }
-                ctx.memory
-                    .set(args[1].memory_address(ctx.fp)?, F::from_usize(value >> 24))?;
             }
             Self::DecomposeBits => {
                 let to_decompose = args[0].read_value(ctx.memory, ctx.fp)?.to_usize();

crates/rec_aggregation/utils.py

@@ -516,19 +516,22 @@ def whir_1_merkle_step_and_pow(v, state_in, path_chunk, state_out, power_shift):
 @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)
+    hint_decompose_bits_merkle_whir(nibbles, 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
+
+    # p = 2^31 - 2^24 + 1, so 2^24 * 127 = p - 1 ≡ -1 (mod p), hence inv(2^24) = -127.
+    # Deduce top7 from the identity partial_sum + top7 * 2^24 == a:
+    # top7 = (a - partial_sum) * inv(2^24) = (partial_sum - a) * 127
+    top7 = (partial_sum - a) * 127
+    assert top7 < 2**7
     if top7 == 2**7 - 1:
         assert partial_sum == 0
-    assert partial_sum + top7 * 2**24 == a
 
     leaf_data = Array(num_chunks * DIGEST_LEN)
     hint_witness("merkle_leaf", leaf_data)