WIP: change hash pros to get pros

src/merkle/smt/full/mod.rs

@@ -292,21 +292,23 @@ impl SparseMerkleTree<SMT_DEPTH> 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<SMT_DEPTH> = 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
             },
         }
     }

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()));
 

src/merkle/smt/mod.rs

@@ -140,7 +140,8 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
             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<const DEPTH: u8> {
     #[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<DEPTH>;

src/merkle/smt/simple/mod.rs

@@ -328,8 +328,13 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
         *leaf == Self::EMPTY_VALUE
     }
 
-    fn hash_prospective_leaf(&self, _key: &LeafIndex<DEPTH>, value: &Word) -> RpoDigest {
-        Self::hash_leaf(value)
+    fn get_prospective_leaf(
+        &self,
+        _existing_leaf: Word,
+        _key: &LeafIndex<DEPTH>,
+        value: &Word,
+    ) -> Word {
+        *value
     }
 
     fn key_to_leaf_index(key: &LeafIndex<DEPTH>) -> LeafIndex<DEPTH> {