diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 9dd5d62..60b7bd2 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -292,21 +292,23 @@ impl SparseMerkleTree for Smt { } } - fn hash_prospective_leaf(&self, key: &RpoDigest, value: &Word) -> RpoDigest { - // Future work could avoid cloning the leaf by mirroring some of the insertion logic and - // hashing without an intermediate leaf, but cloning is only expensive for multi-leaves, - // which should be really rare. - let leaf_index: LeafIndex = Self::key_to_leaf_index(key); - match self.leaves.get(&leaf_index.value()) { - None => SmtLeaf::new_single(*key, *value).hash(), - Some(existing_leaf) => { - let mut new_leaf = existing_leaf.clone(); + fn get_prospective_leaf( + &self, + mut existing_leaf: SmtLeaf, + key: &RpoDigest, + value: &Word, + ) -> SmtLeaf { + debug_assert_eq!(existing_leaf.index(), Self::key_to_leaf_index(key)); + + match existing_leaf { + SmtLeaf::Empty(_) => SmtLeaf::new_single(*key, *value), + _ => { if *value != EMPTY_WORD { - new_leaf.insert(*key, *value); + existing_leaf.insert(*key, *value); } else { - new_leaf.remove(*key); + existing_leaf.remove(*key); } - new_leaf.hash() + existing_leaf }, } } diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index dc3faf3..2c4e8c8 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -282,7 +282,7 @@ fn test_checked_insertion() { // insert key-value 1 let root_1 = { - let prospective = smt.hash_prospective_leaf(&key_1, &value_1); + let prospective = smt.get_prospective_leaf(smt.get_leaf(&key_1), &key_1, &value_1).hash(); let old_value_1 = smt.insert(key_1, value_1); assert_eq!(old_value_1, EMPTY_WORD); @@ -313,7 +313,7 @@ fn test_checked_insertion() { // insert key-value 2 let root_2 = { - let prospective = smt.hash_prospective_leaf(&key_2, &value_2); + let prospective = smt.get_prospective_leaf(smt.get_leaf(&key_2), &key_2, &value_2).hash(); let old_value_2 = smt.insert(key_2, value_2); assert_eq!(old_value_2, EMPTY_WORD); @@ -342,7 +342,7 @@ fn test_checked_insertion() { // insert key-value 3 let root_3 = { - let prospective = smt.hash_prospective_leaf(&key_3, &value_3); + let prospective = smt.get_prospective_leaf(smt.get_leaf(&key_3), &key_3, &value_3).hash(); let old_value_3 = smt.insert(key_3, value_3); assert_eq!(old_value_3, EMPTY_WORD); @@ -374,7 +374,10 @@ fn test_checked_insertion() { let old_hash = smt.get_leaf(&key_3).hash(); let old_value_3 = smt.insert_ensure_root(key_3, EMPTY_WORD, root_2).unwrap(); assert_eq!(old_value_3, value_3); - assert_eq!(old_hash, smt.hash_prospective_leaf(&key_3, &old_value_3)); + assert_eq!( + old_hash, + smt.get_prospective_leaf(smt.get_leaf(&key_3), &key_3, &old_value_3).hash() + ); assert_eq!( smt.get_leaf(&key_3), @@ -389,7 +392,10 @@ fn test_checked_insertion() { let old_hash = smt.get_leaf(&key_2).hash(); let old_value_2 = smt.insert_ensure_root(key_2, EMPTY_WORD, root_1).unwrap(); assert_eq!(old_value_2, value_2); - assert_eq!(old_hash, smt.hash_prospective_leaf(&key_2, &old_value_2)); + assert_eq!( + old_hash, + smt.get_prospective_leaf(smt.get_leaf(&key_2), &key_2, &old_value_2).hash() + ); assert_eq!(smt.get_leaf(&key_2), SmtLeaf::Single((key_1, value_1))); @@ -401,7 +407,10 @@ fn test_checked_insertion() { let old_hash = smt.get_leaf(&key_1).hash(); let old_value_1 = smt.insert_ensure_root(key_1, EMPTY_WORD, root_empty).unwrap(); assert_eq!(old_value_1, value_1); - assert_eq!(old_hash, smt.hash_prospective_leaf(&key_1, &old_value_1)); + assert_eq!( + old_hash, + smt.get_prospective_leaf(smt.get_leaf(&key_1), &key_1, &old_value_1).hash() + ); assert_eq!(smt.get_leaf(&key_1), SmtLeaf::new_empty(key_1.into())); diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 6ebd895..a36c314 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -140,7 +140,8 @@ pub(crate) trait SparseMerkleTree { NodeIndex::from(leaf_index) }; - let mut new_child_hash = self.hash_prospective_leaf(&key, &value); + let mut new_child_hash = + Self::hash_leaf(&self.get_prospective_leaf(self.get_leaf(&key), &key, &value)); for node_depth in (0..node_index.depth()).rev() { let is_right = node_index.is_value_odd(); node_index.move_up(); @@ -260,13 +261,19 @@ pub(crate) trait SparseMerkleTree { #[allow(dead_code)] fn is_leaf_empty(leaf: &Self::Leaf) -> bool; - /// Returns the hash of a leaf if the leaf WERE inserted into the tree, - /// without performing any insertion or other mutation. + /// Returns what `existing_leaf` would look like if `key` and `value` WERE inserted into the + /// tree, without mutating the tree itself. /// - /// Note: calling this function after actually performing an insert with the same arguments - /// will *not* return the same result, as inserting multiple times with the same key mutates - /// the leaf each time. - fn hash_prospective_leaf(&self, key: &Self::Key, value: &Self::Value) -> RpoDigest; + /// `existing_leaf` must have the same index as the key, or the result will be meaningless. To + /// get a prospective leaf based on the current state of the tree, use `self.get_leaf(key)` as + /// the argument for `existing_leaf`. The return value from this function can be chained back + /// into this function as the first argument to continue making prospective changes. + fn get_prospective_leaf( + &self, + existing_leaf: Self::Leaf, + key: &Self::Key, + value: &Self::Value, + ) -> Self::Leaf; /// Maps a key to a leaf index fn key_to_leaf_index(key: &Self::Key) -> LeafIndex; diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 5463fe5..80bc2f1 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -328,8 +328,13 @@ impl SparseMerkleTree for SimpleSmt { *leaf == Self::EMPTY_VALUE } - fn hash_prospective_leaf(&self, _key: &LeafIndex, value: &Word) -> RpoDigest { - Self::hash_leaf(value) + fn get_prospective_leaf( + &self, + _existing_leaf: Word, + _key: &LeafIndex, + value: &Word, + ) -> Word { + *value } fn key_to_leaf_index(key: &LeafIndex) -> LeafIndex {