test: Add example program for merkle_step_mem

triton-vm/src/example_programs.rs

@@ -13,6 +13,7 @@ lazy_static! {
         calculate_new_mmr_peaks_from_append_with_safe_lists();
     pub static ref MERKLE_TREE_AUTHENTICATION_PATH_VERIFY: Program =
         merkle_tree_authentication_path_verify();
+    pub static ref MERKLE_TREE_UPDATE: Program = merkle_tree_update();
 }
 
 fn fibonacci_sequence() -> Program {
@@ -174,6 +175,47 @@ fn merkle_tree_authentication_path_verify() -> Program {
     )
 }
 
+/// Triton program to verifiably change a Merkle tree's leaf. That is:
+/// 1. verify that the supplied `old_leaf` is indeed a leaf in the Merkle tree
+///     defined by the `merkle_root` and the `tree_height`,
+/// 2. update the leaf at the specified `leaf_index` with the `new_leaf`, and
+/// 3. return the new Merkle root.
+///
+/// The authentication path for the leaf to update has to be supplied via RAM.
+///
+/// - input:
+///     - RAM address of leaf's authentication path
+///     - leaf index to update
+///     - Merkle tree's height
+///     - old leaf
+///     - (current) merkle root
+///     - new leaf
+/// - output:
+///     - new root
+fn merkle_tree_update() -> Program {
+    triton_program! {
+        read_io 3           // _ *ap leaf_index tree_height
+        push 2 pow add      // _ *ap node_index
+        dup 1 push 1 dup 2  // _ *ap node_index *ap 1 node_index
+        read_io 5           // _ *ap node_index *ap 1 node_index [old_leaf; 5]
+        call compute_root   // _ *ap node_index *ap' 1 1 [root; 5]
+        read_io 5           // _ *ap node_index *ap' 1 1 [root; 5] [presumed_root; 5]
+        assert_vector       // _ *ap node_index *ap' 1 1 [root; 5]
+        pop 5 pop 3         // _ *ap node_index
+        push 1 swap 1       // _ *ap 1 node_index
+        read_io 5           // _ *ap 1 node_index [new_leaf; 5]
+        call compute_root   // _ *ap' 1 1 [new_root; 5]
+        write_io 5          // _ *ap' 1 1
+        pop 3 halt          // _
+
+        // BEFORE: _ *ap                     1 node_index [leaf; 5]
+        // AFTER:  _ (*ap + 5 * tree_height) 1          1 [root; 5]
+        compute_root:
+            merkle_step_mem
+            recurse_or_return
+    }
+}
+
 fn verify_sudoku() -> Program {
     // RAM layout:
     // 0..=8: primes for mapping digits 1..=9

triton-vm/src/vm.rs

@@ -2365,6 +2365,53 @@ pub(crate) mod tests {
         assert!(let Ok(_) = program.run(public_input.into(), non_determinism));
     }
 
+    #[proptest]
+    fn triton_assembly_merkle_tree_update(
+        leaved_merkle_tree: LeavedMerkleTreeTestData,
+        #[strategy(arb())] new_leaf: Digest,
+        #[strategy(arb())] auth_path_address: BFieldElement,
+    ) {
+        let merkle_tree = leaved_merkle_tree.merkle_tree;
+        let leaf_index = leaved_merkle_tree.revealed_indices[0];
+        let auth_path = merkle_tree.authentication_structure(&[leaf_index]).unwrap();
+
+        let mut ram = HashMap::new();
+        {
+            let mut auth_path_address = auth_path_address;
+            for &Digest(digest) in &auth_path {
+                for element in digest {
+                    ram.insert(auth_path_address, element);
+                    auth_path_address.increment();
+                }
+            }
+        }
+        let non_determinism = NonDeterminism::default().with_ram(ram);
+
+        let old_leaf = Digest::from(leaved_merkle_tree.leaves[leaf_index]);
+        let old_root = merkle_tree.root();
+        let mut public_input = vec![
+            auth_path_address,
+            u64::try_from(leaf_index).unwrap().into(),
+            u64::try_from(merkle_tree.height()).unwrap().into(),
+        ];
+        public_input.extend(old_leaf.reversed().values());
+        public_input.extend(old_root.reversed().values());
+        public_input.extend(new_leaf.reversed().values());
+        let public_input = PublicInput::new(public_input);
+
+        let program = MERKLE_TREE_UPDATE.clone();
+        let_assert!(Ok(output) = program.run(public_input, non_determinism));
+        let new_root = Digest(output.try_into().unwrap());
+
+        let inclusion_proof = MerkleTreeInclusionProof::<Tip5> {
+            tree_height: merkle_tree.height(),
+            indexed_leaves: vec![(leaf_index, new_leaf)],
+            authentication_structure: auth_path,
+            _hasher: std::marker::PhantomData,
+        };
+        assert!(inclusion_proof.verify(new_root));
+    }
+
     #[proptest]
     fn run_tvm_get_collinear_y(
         #[strategy(arb())] p0: (BFieldElement, BFieldElement),