simplify hint_decompose_bits_xmss

Author: TomWambsgans

crates/lean_vm/src/isa/hint.rs

@@ -129,7 +129,7 @@ impl CustomHint {
 
     pub fn n_args(&self) -> usize {
         match self {
-            Self::DecomposeBitsXMSS => 5,
+            Self::DecomposeBitsXMSS => 4,
             Self::DecomposeBitsMerkleWhir => 3,
             Self::DecomposeBits => 4,
             Self::LessThan => 3,
@@ -145,13 +145,11 @@ impl CustomHint {
         match self {
             Self::DecomposeBitsXMSS => {
                 let decomposed_ptr = args[0].read_value(ctx.memory, ctx.fp)?.to_usize();
-                let remaining_ptr = args[1].read_value(ctx.memory, ctx.fp)?.to_usize();
-                let to_decompose_ptr = args[2].read_value(ctx.memory, ctx.fp)?.to_usize();
-                let num_to_decompose = args[3].read_value(ctx.memory, ctx.fp)?.to_usize();
-                let chunk_size = args[4].read_value(ctx.memory, ctx.fp)?.to_usize();
+                let to_decompose_ptr = args[1].read_value(ctx.memory, ctx.fp)?.to_usize();
+                let num_to_decompose = args[2].read_value(ctx.memory, ctx.fp)?.to_usize();
+                let chunk_size = args[3].read_value(ctx.memory, ctx.fp)?.to_usize();
                 assert!(24_usize.is_multiple_of(chunk_size));
                 let mut memory_index_decomposed = decomposed_ptr;
-                let mut memory_index_remaining = remaining_ptr;
                 #[allow(clippy::explicit_counter_loop)]
                 for i in 0..num_to_decompose {
                     let value = ctx.memory.get(to_decompose_ptr + i)?.to_usize();
@@ -160,8 +158,6 @@ impl CustomHint {
                         ctx.memory.set(memory_index_decomposed, value)?;
                         memory_index_decomposed += 1;
                     }
-                    ctx.memory.set(memory_index_remaining, F::from_usize(value >> 24))?;
-                    memory_index_remaining += 1;
                 }
             }
             Self::DecomposeBitsMerkleWhir => {

crates/rec_aggregation/xmss_aggregate.py

@@ -39,22 +39,23 @@ def xmss_verify(merkle_root, message, slot_lo, slot_hi, merkle_chunks):
     poseidon16_compress(b_input, b_input + DIGEST_LEN, encoding_fe)
 
     encoding = Array(NUM_ENCODING_FE * 24 / (2 * W))
-    remaining = Array(NUM_ENCODING_FE)
 
-    hint_decompose_bits_xmss(encoding, remaining, encoding_fe, NUM_ENCODING_FE, 2 * W)
+    hint_decompose_bits_xmss(encoding, encoding_fe, NUM_ENCODING_FE, 2 * W)
 
     # check that the decomposition is correct
     for i in unroll(0, NUM_ENCODING_FE):
         for j in unroll(0, 24 / (2 * W)):
             assert encoding[i * (24 / (2 * W)) + j] < CHAIN_LENGTH**2
 
-        assert remaining[i] < 2**7 - 1  # ensures uniformity + prevent overflow
-
-        partial_sum: Mut = remaining[i] * 2**24
-        partial_sum += encoding[i * (24 / (2 * W))]
+        partial_sum: Mut = encoding[i * (24 / (2 * W))]
         for j in unroll(1, 24 / (2 * W)):
             partial_sum += encoding[i * (24 / (2 * W)) + j] * (CHAIN_LENGTH**2) ** j
-        assert partial_sum == encoding_fe[i]
+
+        # p = 2^31 - 2^24 + 1, so inv(2^24) = -127 (mod p).
+        # Deduce remaining_i from partial_sum + remaining_i * 2^24 == encoding_fe[i]:
+        # remaining_i = (encoding_fe[i] - partial_sum) * inv(2^24) = (partial_sum - encoding_fe[i]) * 127
+        remaining_i = (partial_sum - encoding_fe[i]) * 127
+        assert remaining_i < 2**7 - 1  # ensures uniformity + prevent overflow
 
     # grinding
     debug_assert(V_GRINDING % 2 == 0)