diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e89d61b1..32d6bdec4 100644 --- a/CHANGELOG.md +++ b/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)). diff --git a/miden-crypto/src/merkle/smt/forest/mod.rs b/miden-crypto/src/merkle/smt/forest/mod.rs index 597cdabcb..81fbfe4fe 100644 --- a/miden-crypto/src/merkle/smt/forest/mod.rs +++ b/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. /// diff --git a/miden-crypto/src/merkle/smt/forest/store.rs b/miden-crypto/src/merkle/smt/forest/store.rs index aedb0513b..b80737e3e 100644 --- a/miden-crypto/src/merkle/smt/forest/store.rs +++ b/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, - new_nodes: &mut Map, - ) { - 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, + ) { + let mut new_nodes: Map = 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) { + 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 diff --git a/miden-crypto/src/merkle/smt/forest/tests.rs b/miden-crypto/src/merkle/smt/forest/tests.rs index 93e011eb3..a1662f8d5 100644 --- a/miden-crypto/src/merkle/smt/forest/tests.rs +++ b/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();