Add MMR unpacking

stdlib/asm/collections/mmr.masm

@@ -1,3 +1,5 @@
+use.std::mem
+
 #! Computes the `ilog2(number)` and `2^(ilog2(number))`.
 #!
 #! number must be non-zero, otherwise this will error
@@ -104,3 +106,72 @@ export.get
      # stack: [leaf, ...]
   end
 end
+
+#! Load the MMR peak data based on its hash.
+#!
+#! Input: [HASH, mmr_ptr, ...]
+#! Output: [...]
+#!
+#! Where:
+#!  - HASH: is the MMR peak hash, the hash is expected to be padded to an even
+#!    length and to have a minimum size of 16 elements
+#!  - The advice map must contain a key with HASH, and its value is
+#!    `number_of_leaves || hash_data`, and hash_data is the data used to computed `HASH`
+#!  - mmt_ptr: the memory location where the MMR data will be written to,
+#!    starting with the MMR forest (its total leaves count) followed by its peaks
+#!
+#! Cycles: 156 + 9 * extra_peak_pair cycles
+#!    where `extra_peak` is the number of peak pairs in addition to the first
+#!    16, i.e. `round_up((num_of_peaks - 16) / 2)`
+export.unpack
+  # load the number_of_leaves and peaks to the advice_stack (0 cycles)
+  adv.keyval
+  # operand_stack => [HASH, mmr_ptr, ...]
+  # advice_stack => [number_of_leaves, peaks*, ...]
+
+  # load the size from the advice stack (1 cycles)
+  adv_push.1
+  # operand_stack => [forest, HASH, mmr_ptr, ...]
+  # advice_stack => [peaks*, ...]
+
+  # save the forest to memory (4 cycles)
+  dup dup.6 mem_store
+  # => [forest, HASH, mmr_ptr, ...]
+
+  # count number of peaks (69 cycles)
+  u32split u32unchecked_popcnt swap u32unchecked_popcnt add
+  # => [count, HASH, mmr_ptr, ...]
+
+  # the peaks are padded to a minimum length of 16 (10 cycles)
+  push.16 u32unchecked_max
+  # => [count_min, HASH, mmr_ptr, ...]
+
+  # when the number of peaks is greater than 16, then they are padded to an even number (7 cycles)
+  dup is_odd add
+  # => [even_count_min, HASH, mmr_ptr, ...]
+
+  # compute the end address including the padding data and forest (3 cycles)
+  dup.5 add add.1
+  # => [mmt_ptr_end, HASH, mmr_ptr, ...]
+
+  # update the mmr_ptr to account for the size (2 cycles)
+  movup.5 add.1
+  # => [mmr_ptr+1, mmt_ptr_end, HASH, ...]
+
+  # hash the first 16 words (28 cycles)
+  padw padw padw
+  adv_pipe adv_pipe adv_pipe adv_pipe
+  adv_pipe adv_pipe adv_pipe adv_pipe
+  # => [C, B, A, mmr_ptr+17, mmt_ptr_end, HASH, ...]
+
+  # handle MMR with more than 16 elements (10 + 9 * words cycles)
+  exec.mem::pipe_double_words_to_memory
+  # => [C, B, A, mmr_ptr+17, HASH, ...]
+
+  # drop anything but the hash result, word B (11 cycles)
+  dropw swapw dropw movup.4 drop
+  # => [B, HASH, ...]
+
+  # assert on the resulting hash (11 cycles)
+  assert_eqw
+end

stdlib/docs/mmr_collections.md

@@ -4,3 +4,4 @@
 | ----------- | ------------- |
 | ilog2_checked | Computes the `ilog2(number)` and `2^(ilog2(number))`.<br /><br />number must be non-zero, otherwise this will error<br /><br />Stack transition:<br /><br />Input: [number, ...]<br /><br />Output: [ilog2, power_of_two, ...]<br /><br />Cycles:  12 + 9 * leading_zeros |
 | get | Loads the leaf at the absolute `pos` in the MMR.<br /><br />This MMR implementation supports only u32 positions.<br /><br />Stack transition:<br /><br />Input: [pos, mmr_ptr, ...]<br /><br />Output: [N, ...] where `N` is the leaf and `R` is the MMR peak that owns the leaf.<br /><br />Cycles: 65 + 9 * tree_position (where `tree_position` is 0-indexed bit position from most to least significant) |
+| unpack | Load the MMR peak data based on its hash.<br /><br />Input: [HASH, mmr_ptr, ...]<br /><br />Output: [...]<br /><br />Where:<br /><br />- HASH: is the MMR peak hash, the hash is expected to be padded to an even<br /><br />length and to have a minimum size of 16 elements<br /><br />- mmt_ptr: the memory location where the MMR data will be written to,<br /><br />starting with the MMR forest (its total leaves count) followed by its peaks<br /><br />Cycles: 156 + 9 * extra_peak_pair cycles<br /><br />where `extra_peak` is the number of peak pairs in addition to the first<br /><br />16, i.e. `round_up((num_of_peaks - 16) / 2)` |

stdlib/tests/collections/mmr.rs

@@ -1,6 +1,6 @@
 use test_utils::{
     crypto::{init_merkle_leaves, MerkleError, MerkleStore, NodeIndex},
-    StarkField,
+    hash_elements, stack_to_ints, Felt, StarkField, ZERO,
 };
 
 #[test]
@@ -213,3 +213,205 @@ fn test_mmr_tree_with_one_element() -> Result<(), MerkleError> {
 
     Ok(())
 }
+
+#[test]
+fn test_mmr_unpack() {
+    let number_of_leaves: u64 = 0b10101; // 3 peaks, 21 leaves
+
+    // The hash data is not the same as the peaks, it is padded to 16 elements
+    let hash_data: [[Felt; 4]; 16] = [
+        // 3 peaks. These hashes are invalid, we can't produce data for any of these peaks (only
+        // for testing)
+        [ZERO, ZERO, ZERO, Felt::new(1)],
+        [ZERO, ZERO, ZERO, Felt::new(2)],
+        [ZERO, ZERO, ZERO, Felt::new(3)],
+        // Padding, the MMR is padded to a minimum length o 16
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+    ];
+    let hash = hash_elements(&hash_data.concat());
+
+    // Set up the VM stack with the MMR hash, and its target address
+    let mut stack = stack_to_ints(&*hash);
+    let mmr_ptr = 1000;
+    stack.insert(0, mmr_ptr);
+
+    // both the advice stack and merkle store start empty (data is available in
+    // the map and pushed to the advice stack by the MASM code)
+    let advice_stack = &[];
+    let store = MerkleStore::new();
+
+    let mut map_data: Vec<Felt> = Vec::with_capacity(hash_data.len() + 1);
+    map_data.push(number_of_leaves.into());
+    map_data.extend_from_slice(&hash_data.as_slice().concat());
+
+    let advice_map: &[([u8; 32], Vec<Felt>)] = &[
+        // Under the MMR key is the number_of_leaves, followed by the MMR peaks, and any padding
+        (hash.as_bytes(), map_data),
+    ];
+
+    let source = "
+        use.std::collections::mmr
+        begin exec.mmr::unpack end
+    ";
+    let test = build_test!(source, &stack, advice_stack, store, advice_map.iter().cloned());
+
+    #[rustfmt::skip]
+    let expect_memory = [
+        number_of_leaves, 0, 0, 0, // MMR leaves (only one Felt is used)
+        0, 0, 0, 1,                // first peak
+        0, 0, 0, 2,                // second peak
+        0, 0, 0, 3,                // third peak
+    ];
+    test.expect_stack(&[]);
+    test.expect_stack_and_memory(&[], mmr_ptr, &expect_memory);
+}
+
+#[test]
+fn test_mmr_unpack_invalid_hash() {
+    // The hash data is not the same as the peaks, it is padded to 16 elements
+    let mut hash_data: [[Felt; 4]; 16] = [
+        // 3 peaks. These hashes are invalid, we can't produce data for any of these peaks (only
+        // for testing)
+        [ZERO, ZERO, ZERO, Felt::new(1)],
+        [ZERO, ZERO, ZERO, Felt::new(2)],
+        [ZERO, ZERO, ZERO, Felt::new(3)],
+        // Padding, the MMR is padded to a minimum length o 16
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+        [ZERO, ZERO, ZERO, ZERO],
+    ];
+    let hash = hash_elements(&hash_data.concat());
+
+    // Set up the VM stack with the MMR hash, and its target address
+    let mut stack = stack_to_ints(&*hash);
+    let mmr_ptr = 1000;
+    stack.insert(0, mmr_ptr);
+
+    // both the advice stack and merkle store start empty (data is available in
+    // the map and pushed to the advice stack by the MASM code)
+    let advice_stack = &[];
+    let store = MerkleStore::new();
+
+    // corrupt the data, this changes the hash and the commitment check must fail
+    hash_data[0][0] = hash_data[0][0] + Felt::new(1);
+
+    let mut map_data: Vec<Felt> = Vec::with_capacity(hash_data.len() + 1);
+    map_data.push(Felt::new(0b10101)); // 3 peaks, 21 leaves
+    map_data.extend_from_slice(&hash_data.as_slice().concat());
+
+    let advice_map: &[([u8; 32], Vec<Felt>)] = &[
+        // Under the MMR key is the number_of_leaves, followed by the MMR peaks, and any padding
+        (hash.as_bytes(), map_data),
+    ];
+
+    let source = "
+        use.std::collections::mmr
+        begin exec.mmr::unpack end
+    ";
+    let test = build_test!(source, &stack, advice_stack, store, advice_map.iter().cloned());
+
+    assert!(test.execute().is_err());
+}
+
+/// Tests the case of an MMR with more than 16 peaks
+#[test]
+fn test_mmr_unpack_large_mmr() {
+    let number_of_leaves: u64 = 0b11111111111111111; // 17 peaks
+
+    // The hash data is not the same as the peaks, it is padded to 16 elements
+    let hash_data: [[Felt; 4]; 18] = [
+        // These hashes are invalid, we can't produce data for any of these peaks (only for
+        // testing)
+        [ZERO, ZERO, ZERO, Felt::new(1)],
+        [ZERO, ZERO, ZERO, Felt::new(2)],
+        [ZERO, ZERO, ZERO, Felt::new(3)],
+        [ZERO, ZERO, ZERO, Felt::new(4)],
+        [ZERO, ZERO, ZERO, Felt::new(5)],
+        [ZERO, ZERO, ZERO, Felt::new(6)],
+        [ZERO, ZERO, ZERO, Felt::new(7)],
+        [ZERO, ZERO, ZERO, Felt::new(8)],
+        [ZERO, ZERO, ZERO, Felt::new(9)],
+        [ZERO, ZERO, ZERO, Felt::new(10)],
+        [ZERO, ZERO, ZERO, Felt::new(11)],
+        [ZERO, ZERO, ZERO, Felt::new(12)],
+        [ZERO, ZERO, ZERO, Felt::new(13)],
+        [ZERO, ZERO, ZERO, Felt::new(14)],
+        [ZERO, ZERO, ZERO, Felt::new(15)],
+        [ZERO, ZERO, ZERO, Felt::new(16)],
+        // Padding, peaks greater than 16 are padded to an even number
+        [ZERO, ZERO, ZERO, Felt::new(17)],
+        [ZERO, ZERO, ZERO, ZERO],
+    ];
+    let hash = hash_elements(&hash_data.concat());
+
+    // Set up the VM stack with the MMR hash, and its target address
+    let mut stack = stack_to_ints(&*hash);
+    let mmr_ptr = 1000;
+    stack.insert(0, mmr_ptr);
+
+    // both the advice stack and merkle store start empty (data is available in
+    // the map and pushed to the advice stack by the MASM code)
+    let advice_stack = &[];
+    let store = MerkleStore::new();
+
+    let mut map_data: Vec<Felt> = Vec::with_capacity(hash_data.len() + 1);
+    map_data.push(number_of_leaves.into());
+    map_data.extend_from_slice(&hash_data.as_slice().concat());
+
+    let advice_map: &[([u8; 32], Vec<Felt>)] = &[
+        // Under the MMR key is the number_of_leaves, followed by the MMR peaks, and any padding
+        (hash.as_bytes(), map_data),
+    ];
+
+    let source = "
+        use.std::collections::mmr
+        begin exec.mmr::unpack end
+    ";
+    let test = build_test!(source, &stack, advice_stack, store, advice_map.iter().cloned());
+
+    #[rustfmt::skip]
+    let expect_memory = [
+        number_of_leaves, 0, 0, 0, // MMR leaves (only one Felt is used)
+        0, 0, 0, 1,                // peaks
+        0, 0, 0, 2,
+        0, 0, 0, 3,
+        0, 0, 0, 4,
+        0, 0, 0, 5,
+        0, 0, 0, 6,
+        0, 0, 0, 7,
+        0, 0, 0, 8,
+        0, 0, 0, 9,
+        0, 0, 0, 10,
+        0, 0, 0, 11,
+        0, 0, 0, 12,
+        0, 0, 0, 13,
+        0, 0, 0, 14,
+        0, 0, 0, 15,
+        0, 0, 0, 16,
+        0, 0, 0, 17,
+    ];
+    test.expect_stack(&[]);
+    test.expect_stack_and_memory(&[], mmr_ptr, &expect_memory);
+}

test-utils/src/lib.rs

@@ -8,12 +8,14 @@ extern crate alloc;
 #[cfg(not(target_family = "wasm"))]
 use proptest::prelude::{Arbitrary, Strategy};
 
-use vm_core::chiplets::hasher::{apply_permutation, hash_elements, STATE_WIDTH};
+use vm_core::chiplets::hasher::{apply_permutation, STATE_WIDTH};
 use vm_core::utils::{collections::Vec, string::String};
 
 // EXPORTS
 // ================================================================================================
 
+pub use vm_core::chiplets::hasher::hash_elements;
+
 pub use assembly::{Library, MaslLibrary};
 pub use processor::{
     AdviceInputs, AdviceProvider, ExecutionError, ExecutionTrace, Process, StackInputs,
@@ -23,6 +25,7 @@ pub use prover::{MemAdviceProvider, ProofOptions};
 pub use test_case::test_case;
 pub use verifier::ProgramInfo;
 pub use vm_core::{
+    crypto::hash::RpoDigest,
     stack::STACK_TOP_SIZE,
     utils::{collections, group_slice_elements, group_vector_elements, IntoBytes, ToElements},
     Felt, FieldElement, Program, StarkField, Word, ONE, WORD_SIZE, ZERO,
@@ -138,8 +141,9 @@ impl Test {
     /// Builds a final stack from the provided stack-ordered array and asserts that executing the
     /// test will result in the expected final stack state.
     pub fn expect_stack(&self, final_stack: &[u64]) {
-        let result = self.get_last_stack_state();
-        assert_eq!(stack_to_top_ints(final_stack), stack_to_ints(&result));
+        let result = stack_to_ints(&self.get_last_stack_state());
+        let expected = stack_top_to_ints(final_stack);
+        assert_eq!(expected, result, "Expected stack to be {:?}, found {:?}", expected, result);
     }
 
     /// Executes the test and validates that the process memory has the elements of `expected_mem`
@@ -162,9 +166,12 @@ impl Test {
 
         // validate the memory state
         for data in expected_mem.chunks(WORD_SIZE) {
-            let mem_state = process.get_memory_value(0, mem_start_addr).unwrap();
-            let expected_mem: Vec<Felt> = data.iter().map(|&v| Felt::new(v)).collect();
-            assert_eq!(expected_mem, mem_state);
+            let mem_state = stack_to_ints(&process.get_memory_value(0, mem_start_addr).unwrap());
+            assert_eq!(
+                data, mem_state,
+                "Expected memory [{}] => {:?}, found {:?}",
+                mem_start_addr, data, mem_state
+            );
             mem_start_addr += 1;
         }
 
@@ -181,7 +188,7 @@ impl Test {
         final_stack: &[u64],
     ) -> Result<(), proptest::prelude::TestCaseError> {
         let result = self.get_last_stack_state();
-        proptest::prop_assert_eq!(stack_to_top_ints(final_stack), stack_to_ints(&result));
+        proptest::prop_assert_eq!(stack_top_to_ints(final_stack), stack_to_ints(&result));
 
         Ok(())
     }
@@ -258,7 +265,7 @@ pub fn stack_to_ints(values: &[Felt]) -> Vec<u64> {
     values.iter().map(|e| (*e).as_int()).collect()
 }
 
-pub fn stack_to_top_ints(values: &[u64]) -> Vec<u64> {
+pub fn stack_top_to_ints(values: &[u64]) -> Vec<u64> {
     let mut result: Vec<u64> = values.to_vec();
     result.resize(STACK_TOP_SIZE, 0);
     result