add SmtForest::insert_path()

CHANGELOG.md

@@ -7,6 +7,7 @@
 - Add serialization of sealing and unsealing keys in IES ([#637](https://github.com/0xMiden/crypto/pull/637)).
 - Fixed undefined `BaseElement` in rescue arch optimizations ([#644](https://github.com/0xMiden/crypto/pull/644)).
 - Added `MerkleStore::has_path()` ([#649](https://github.com/0xMiden/crypto/pull/649)).
+- Added `SmtForest::insert_path()` to support partial SMTs in `SmtForest` ([#650](https://github.com/0xMiden/crypto/pull/650)).
 
 ## 0.18.2 (2025-11-08)
 

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_path(&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_DEPTH},
+    merkle::{
+        EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerkleProof, NodeIndex,
+        SMT_DEPTH,
+    },
 };
 
 // SMT FOREST STORE
@@ -242,31 +245,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.
     ///
@@ -303,6 +301,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,7 +4,10 @@ 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
@@ -89,6 +92,28 @@ fn test_insert_multiple_values() -> Result<(), MerkleError> {
     Ok(())
 }
 
+#[test]
+fn test_insert_path() -> Result<(), MerkleError> {
+    let mut smt = Smt::new();
+    let key = Word::new([ZERO, ZERO, ZERO, ONE]);
+    let value = Word::new([ONE; WORD_SIZE]);
+    smt.insert(key, value)?;
+    let proof = smt.open(&key);
+
+    let mut forest = SmtForest::new();
+    let root = forest.insert_path(proof);
+    assert_eq!(root, smt.root());
+
+    let stored_proof = forest.open(root, key)?;
+    assert!(stored_proof.verify_membership(&key, &value, &root));
+
+    forest.pop_smts(vec![root]);
+    assert!(forest.roots.is_empty());
+    assert!(forest.leaves.is_empty());
+
+    Ok(())
+}
+
 #[test]
 fn test_batch_insert() -> Result<(), MerkleError> {
     let forest = SmtForest::new();