add SmtForest::insert_path()

CHANGELOG.md

@@ -8,6 +8,7 @@
 - Fixed undefined `BaseElement` in rescue arch optimizations ([#644](https://github.com/0xMiden/crypto/pull/644)).
 - Fixed bugs in Merkle tree capacity checks for `SimpleSmt` and `PartialMerkleTree` ([#648](https://github.com/0xMiden/crypto/pull/648)).
 - Added `MerkleStore::has_path()` ([#649](https://github.com/0xMiden/crypto/pull/649)).
+- Added `SmtForest::insert_proof()` to support partial SMTs in `SmtForest` ([#650](https://github.com/0xMiden/crypto/pull/650)).
 - Refactored `StorageUpdates` to use explicit `SubtreeUpdate` enum for storage operations ([#654](https://github.com/0xMiden/crypto/issues/654)).
 - Refactored `LargeSmt` into smaller focused modules ([#658](https://github.com/0xMiden/crypto/pull/658)).
 - [BREAKING] Organized `merkle` module into public submodules (`mmr`, `smt`, `store`) ([#660](https://github.com/0xMiden/crypto/pull/660)).

miden-crypto/src/merkle/smt/forest/mod.rs

@@ -118,6 +118,26 @@ impl SmtForest {
     // STATE MUTATORS
     // --------------------------------------------------------------------------------------------
 
+    /// Inserts all nodes present in the provided [`SmtProof`] into the forest and returns
+    /// the root computed from the proof.
+    ///
+    /// If the computed root already exists, returns without modifying the forest.
+    pub fn insert_proof(&mut self, proof: SmtProof) -> Word {
+        let root = proof.compute_root();
+        let path_nodes: Vec<_> = proof.authenticated_nodes().collect();
+        let (_path, leaf) = proof.into_parts();
+
+        if !self.roots.insert(root) {
+            return root;
+        }
+
+        let leaf_hash = leaf.hash();
+        self.leaves.insert(leaf_hash, leaf);
+        self.store.insert_nodes_from_path(root, path_nodes);
+
+        root
+    }
+
     /// Inserts the specified key-value pair into an SMT with the specified root. This will also
     /// add a new root to the forest. Returns the new root.
     ///

miden-crypto/src/merkle/smt/forest/store.rs

@@ -3,7 +3,10 @@ use alloc::vec::Vec;
 use crate::{
     Map, Word,
     hash::rpo::Rpo256,
-    merkle::{EmptySubtreeRoots, MerkleError, MerklePath, MerkleProof, NodeIndex, smt::SMT_DEPTH},
+    merkle::{
+        EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerkleProof, NodeIndex,
+        smt::SMT_DEPTH,
+    },
 };
 
 // SMT FOREST STORE
@@ -46,6 +49,11 @@ impl SmtStore {
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
+    #[cfg(test)]
+    pub(super) fn num_nodes(&self) -> usize {
+        self.nodes.len()
+    }
+
     /// Returns the node at `index` rooted on the tree `root`.
     ///
     /// # Errors
@@ -242,31 +250,26 @@ impl SmtStore {
             .ok_or(MerkleError::NodeIndexNotFoundInStore(root, NodeIndex::root()))?;
 
         // The update was computed successfully, update ref counts and insert into the store
-        fn dfs(
-            node: Word,
-            store: &mut Map<Word, ForestInnerNode>,
-            new_nodes: &mut Map<Word, ForestInnerNode>,
-        ) {
-            if node == Word::empty() {
-                return;
-            }
-            if let Some(node) = store.get_mut(&node) {
-                // This node already exists in the store, increase its reference count.
-                // Stops the dfs descent here to leave children ref counts unchanged.
-                node.rc += 1;
-            } else if let Some(mut smt_node) = new_nodes.remove(&node) {
-                // This is a non-leaf node, insert it into the store and process its children.
-                smt_node.rc = 1;
-                store.insert(node, smt_node);
-                dfs(smt_node.left, store, new_nodes);
-                dfs(smt_node.right, store, new_nodes);
-            }
-        }
-        dfs(new_root, &mut self.nodes, &mut new_nodes);
+        self.insert_node_recursive(new_root, &mut new_nodes);
 
         Ok(new_root)
     }
 
+    /// Inserts the nodes described by `path_nodes` into the store, updating reference counts and
+    /// ensuring that shared subtrees remain tracked only once.
+    pub(super) fn insert_nodes_from_path(
+        &mut self,
+        root: Word,
+        path_nodes: impl IntoIterator<Item = InnerNodeInfo>,
+    ) {
+        let mut new_nodes: Map<Word, ForestInnerNode> = Map::new();
+        for InnerNodeInfo { value, left, right } in path_nodes {
+            new_nodes.insert(value, ForestInnerNode { left, right, rc: 0 });
+        }
+
+        self.insert_node_recursive(root, &mut new_nodes);
+    }
+
     /// Decreases the reference count of the specified node and releases memory if the count
     /// reached zero.
     ///
@@ -286,6 +289,8 @@ impl SmtStore {
         let left = smt_node.left;
         let right = smt_node.right;
 
+        self.nodes.remove(&node);
+
         let mut result = Vec::new();
         result.extend(self.remove_node(left));
         result.extend(self.remove_node(right));
@@ -303,6 +308,22 @@ impl SmtStore {
         }
         removed_leaves
     }
+
+    /// Inserts nodes from `new_nodes` into the store, increasing reference counts for existing
+    /// nodes and recursively adding unseen nodes together with their subtrees.
+    fn insert_node_recursive(&mut self, node: Word, new_nodes: &mut Map<Word, ForestInnerNode>) {
+        if node == Word::empty() {
+            return;
+        }
+        if let Some(existing) = self.nodes.get_mut(&node) {
+            existing.rc += 1;
+        } else if let Some(mut smt_node) = new_nodes.remove(&node) {
+            smt_node.rc = 1;
+            self.nodes.insert(node, smt_node);
+            self.insert_node_recursive(smt_node.left, new_nodes);
+            self.insert_node_recursive(smt_node.right, new_nodes);
+        }
+    }
 }
 
 // HELPER FUNCTIONS

miden-crypto/src/merkle/smt/forest/tests.rs

@@ -4,12 +4,18 @@ use itertools::Itertools;
 use super::{EmptySubtreeRoots, MerkleError, SmtForest, Word};
 use crate::{
     Felt, ONE, WORD_SIZE, ZERO,
-    merkle::{int_to_node, smt::SMT_DEPTH},
+    merkle::{
+        int_to_node,
+        smt::{SMT_DEPTH, Smt},
+    },
 };
 
 // TESTS
 // ================================================================================================
 
+// Number of nodes in an empty forest.
+const EMPTY_NODE_COUNT: usize = SMT_DEPTH as usize;
+
 #[test]
 fn test_insert_root_not_in_store() -> Result<(), MerkleError> {
     let mut forest = SmtForest::new();
@@ -89,6 +95,43 @@ fn test_insert_multiple_values() -> Result<(), MerkleError> {
     Ok(())
 }
 
+#[test]
+fn test_insert_proof() -> Result<(), MerkleError> {
+    // Create an SMT with multiple entries to test partial forest view
+    let key1 = Word::new([ZERO, ZERO, ZERO, ONE]);
+    let value1 = Word::new([ONE; WORD_SIZE]);
+
+    let key2 = Word::new([ZERO, ZERO, ONE, ZERO]);
+    let value2 = Word::new([ONE; WORD_SIZE]);
+
+    let key3 = Word::new([ZERO, ONE, ZERO, Felt::new(2)]);
+    let value3 = Word::new([ONE; WORD_SIZE]);
+
+    let smt = Smt::with_entries(vec![(key1, value1), (key2, value2), (key3, value3)])?;
+    let proof = smt.open(&key1);
+
+    let mut forest = SmtForest::new();
+    assert_eq!(forest.store.num_nodes(), EMPTY_NODE_COUNT);
+    let root = forest.insert_proof(proof);
+    assert_eq!(root, smt.root());
+
+    // key1 should be accessible as we inserted its proof
+    let stored_proof = forest.open(root, key1)?;
+    assert!(stored_proof.verify_membership(&key1, &value1, &root));
+
+    // key2 path is available, but the key is not tracked in the forest.
+    assert_matches!(forest.open(root, key2), Err(MerkleError::UntrackedKey(_)));
+    // key3 path is not available in the forest.
+    assert_matches!(forest.open(root, key3), Err(MerkleError::NodeIndexNotFoundInStore(_, _)));
+
+    forest.pop_smts(vec![root]);
+    assert_eq!(forest.store.num_nodes(), EMPTY_NODE_COUNT);
+    assert!(forest.roots.is_empty());
+    assert!(forest.leaves.is_empty());
+
+    Ok(())
+}
+
 #[test]
 fn test_batch_insert() -> Result<(), MerkleError> {
     let forest = SmtForest::new();