implement hashing for leaves that don't yet exist

This commit allows users of the SparseMerkleTree to determine the hash
of a leaf-to-be before it's actually inserted. This functionality is
currently unused, but in further commits will be used to implement
pre-validated Merkle tree insertion.

src/merkle/smt/full/mod.rs

@@ -262,6 +262,32 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
         leaf.hash()
     }
 
+    fn hash_prospective_leaf(&self, key: &RpoDigest, value: &Word) -> RpoDigest {
+        let leaf_index = Self::key_to_leaf_index(key);
+
+        // The he hash always depends on if we already have a leaf at this
+        // key or not.
+        match self.leaves.get(&leaf_index.value()) {
+            Some(existing_leaf) => {
+                // TODO: avoid cloning the leaf, which is potentially expensive
+                // for leaves with multiple values.
+                let mut cloned = existing_leaf.clone();
+                if value == &Self::EMPTY_VALUE {
+                    cloned.remove(*key);
+                } else {
+                    cloned.insert(*key, *value);
+                }
+                cloned.hash()
+            },
+            None => {
+                if value == &Self::EMPTY_VALUE {
+                    return SmtLeaf::new_empty(leaf_index).hash();
+                }
+                SmtLeaf::new_single(*key, *value).hash()
+            },
+        }
+    }
+
     fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex<SMT_DEPTH> {
         let most_significant_felt = key[3];
         LeafIndex::new_max_depth(most_significant_felt.as_int())

src/merkle/smt/full/tests.rs

@@ -1,6 +1,6 @@
 use super::{Felt, LeafIndex, NodeIndex, Rpo256, RpoDigest, Smt, SmtLeaf, EMPTY_WORD, SMT_DEPTH};
 use crate::{
-    merkle::{EmptySubtreeRoots, MerkleStore},
+    merkle::{smt::SparseMerkleTree, EmptySubtreeRoots, MerkleStore},
     utils::{Deserializable, Serializable},
     Word, ONE, WORD_SIZE,
 };
@@ -326,6 +326,87 @@ fn test_smt_entries() {
     assert!(entries.next().is_none());
 }
 
+/// Tests that we can correctly calculate leaf hashes before actually inserting them.
+#[test]
+fn test_prospective_hash() {
+    let mut smt = Smt::default();
+
+    let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
+
+    let key_1 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
+    let key_2 = RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
+    let key_3 = RpoDigest::from([3_u32.into(), 3_u32.into(), 3_u32.into(), Felt::new(raw)]);
+
+    let value_1 = [ONE; WORD_SIZE];
+    let value_2: [Felt; 4] = [2_u32.into(); WORD_SIZE];
+    let value_3: [Felt; 4] = [3_u32.into(); WORD_SIZE];
+    let value_3_2: [Felt; 4] = [3_u32.into(), 2_u32.into(), 2_u32.into(), 2_u32.into()];
+
+    // Prospective hashes before insertion should be equal to actual hashes
+    // after insertion.
+    {
+        // That includes inserting empty values on an empty tree.
+        let prospective_empty = smt.hash_prospective_leaf(&key_1, &Smt::EMPTY_VALUE);
+        smt.insert(key_1, Smt::EMPTY_VALUE);
+        assert_eq!(smt.get_leaf(&key_1).hash(), prospective_empty);
+
+        let prospective = smt.hash_prospective_leaf(&key_1, &value_1);
+        smt.insert(key_1, value_1);
+        let actual = smt.get_leaf(&key_1).hash();
+        assert_eq!(actual, prospective);
+    }
+
+    {
+        let prospective = smt.hash_prospective_leaf(&key_2, &value_2);
+        smt.insert(key_2, value_2);
+        assert_eq!(smt.get_leaf(&key_2).hash(), prospective);
+    }
+
+    {
+        let prospective = smt.hash_prospective_leaf(&key_3, &value_3);
+        smt.insert(key_3, value_3);
+        assert_eq!(smt.get_leaf(&key_3).hash(), prospective);
+
+        let prospective_multi = smt.hash_prospective_leaf(&key_3, &value_3_2);
+        smt.insert(key_3, value_3_2);
+        assert_eq!(smt.get_leaf(&key_3).hash(), prospective_multi);
+    }
+
+    // Prospective hashes after removal should be equal to actual hashes
+    // before removal, and prospective hashes *for* the removal should be
+    // equal to actual hashes after the removal.
+    {
+        let prospective_empty = smt.hash_prospective_leaf(&key_3, &Smt::EMPTY_VALUE);
+        let old_hash = smt.get_leaf(&key_3).hash();
+        smt.insert(key_3, Smt::EMPTY_VALUE);
+        assert_eq!(old_hash, smt.hash_prospective_leaf(&key_3, &value_3_2));
+        assert_eq!(smt.get_leaf(&key_3).hash(), prospective_empty);
+
+        let prospective_empty_2 = smt.hash_prospective_leaf(&key_3, &Smt::EMPTY_VALUE);
+        let old_hash_2 = smt.get_leaf(&key_3).hash();
+        assert_ne!(old_hash, old_hash_2);
+        assert_eq!(smt.get_leaf(&key_3).hash(), prospective_empty_2);
+        smt.insert(key_3, Smt::EMPTY_VALUE);
+        assert_eq!(old_hash_2, smt.hash_prospective_leaf(&key_3, &Smt::EMPTY_VALUE));
+    }
+
+    {
+        let prospective_empty = smt.hash_prospective_leaf(&key_2, &Smt::EMPTY_VALUE);
+        let old_hash = smt.get_leaf(&key_2).hash();
+        smt.insert(key_2, Smt::EMPTY_VALUE);
+        assert_eq!(old_hash, smt.hash_prospective_leaf(&key_2, &value_2));
+        assert_eq!(smt.get_leaf(&key_2).hash(), prospective_empty);
+    }
+
+    {
+        let prospective_empty = smt.hash_prospective_leaf(&key_1, &Smt::EMPTY_VALUE);
+        let old_hash = smt.get_leaf(&key_1).hash();
+        smt.insert(key_1, Smt::EMPTY_VALUE);
+        assert_eq!(old_hash, smt.hash_prospective_leaf(&key_1, &value_1));
+        assert_eq!(smt.get_leaf(&key_1).hash(), prospective_empty);
+    }
+}
+
 // SMT LEAF
 // --------------------------------------------------------------------------------------------
 

src/merkle/smt/mod.rs

@@ -166,6 +166,15 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     /// Returns the hash of a leaf
     fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest;
 
+    /// Returns the hash of a leaf if the leaf WERE inserted into the tree,
+    /// without performing any insertion or other mutation.
+    ///
+    /// 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.
+    #[allow(dead_code)] // Coming soon :)
+    fn hash_prospective_leaf(&self, key: &Self::Key, value: &Self::Value) -> RpoDigest;
+
     /// Maps a key to a leaf index
     fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
 

src/merkle/smt/simple/mod.rs

@@ -301,6 +301,10 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
         leaf.into()
     }
 
+    fn hash_prospective_leaf(&self, _: &LeafIndex<DEPTH>, value: &Word) -> RpoDigest {
+        Self::hash_leaf(value)
+    }
+
     fn key_to_leaf_index(key: &LeafIndex<DEPTH>) -> LeafIndex<DEPTH> {
         *key
     }

src/merkle/smt/simple/tests.rs

@@ -443,6 +443,90 @@ fn test_simplesmt_set_subtree_entire_tree() {
     assert_eq!(tree.root(), *EmptySubtreeRoots::entry(DEPTH, 0));
 }
 
+/// Tests that we can correctly calculate leaf hashes before actually inserting them.
+#[test]
+fn test_prospective_hash() {
+    const DEPTH: u8 = 3;
+    type SimpleSmt = super::SimpleSmt::<DEPTH>;
+    const EMPTY: Word = SimpleSmt::EMPTY_VALUE;
+    let mut smt = SimpleSmt::new().unwrap();
+
+    let key = |num: u64| LeafIndex::<DEPTH>::new(num).unwrap();
+
+    let key_1 = key(3);
+    let value_1 = int_to_leaf(9);
+
+    let key_2 = key(4);
+    let value_2 = int_to_leaf(12);
+
+    let key_3 = key(5);
+    let value_3 = int_to_leaf(15);
+    let value_3_2 = int_to_leaf(20);
+
+    // Prospective hashes before insertion should be equal to actual hashes
+    // after insertion.
+    {
+        // That includes inserting empty values on an empty tree.
+        let prospective_empty = smt.hash_prospective_leaf(&key_1, &EMPTY);
+        smt.insert(key_1, EMPTY);
+        assert_eq!(SimpleSmt::hash_leaf(&smt.get_leaf(&key_1)), prospective_empty);
+
+        let prospective = smt.hash_prospective_leaf(&key_1, &value_1);
+        smt.insert(key_1, value_1);
+        assert_eq!(SimpleSmt::hash_leaf(&smt.get_leaf(&key_1)), prospective);
+    }
+
+    {
+        let prospective = smt.hash_prospective_leaf(&key_2, &value_2);
+        smt.insert(key_2, value_2);
+        assert_eq!(SimpleSmt::hash_leaf(&smt.get_leaf(&key_2)), prospective);
+    }
+
+    {
+        let prospective = smt.hash_prospective_leaf(&key_3, &value_3);
+        smt.insert(key_3, value_3);
+        assert_eq!(SimpleSmt::hash_leaf(&smt.get_leaf(&key_3)), prospective);
+
+        let prospective_multi = smt.hash_prospective_leaf(&key_3, &value_3_2);
+        smt.insert(key_3, value_3_2);
+        assert_eq!(SimpleSmt::hash_leaf(&smt.get_leaf(&key_3)), prospective_multi);
+    }
+
+    // Prospective hashes after removal should be equal to actual hashes
+    // before removal, and prospective hashes *for* the removal should be
+    // equal to actual hashes after the removal.
+    {
+        let prospective_empty = smt.hash_prospective_leaf(&key_3, &EMPTY);
+        let old_hash = SimpleSmt::hash_leaf(&smt.get_leaf(&key_3));
+        smt.insert(key_3, EMPTY);
+        assert_eq!(old_hash, smt.hash_prospective_leaf(&key_3, &value_3_2));
+        assert_eq!(SimpleSmt::hash_leaf(&smt.get_leaf(&key_3)), prospective_empty);
+
+        let prospective_empty_2 = smt.hash_prospective_leaf(&key_3, &EMPTY);
+        let old_hash_2 = SimpleSmt::hash_leaf(&smt.get_leaf(&key_3));
+        assert_ne!(old_hash, old_hash_2);
+        assert_eq!(SimpleSmt::hash_leaf(&smt.get_leaf(&key_3)), prospective_empty_2);
+        smt.insert(key_3, EMPTY);
+        assert_eq!(old_hash_2, smt.hash_prospective_leaf(&key_3, &EMPTY));
+    }
+
+    {
+        let prospective_empty = smt.hash_prospective_leaf(&key_2, &EMPTY);
+        let old_hash = SimpleSmt::hash_leaf(&smt.get_leaf(&key_2));
+        smt.insert(key_2, EMPTY);
+        assert_eq!(old_hash, smt.hash_prospective_leaf(&key_2, &value_2));
+        assert_eq!(SimpleSmt::hash_leaf(&smt.get_leaf(&key_2)), prospective_empty);
+    }
+
+    {
+        let prospective_empty = smt.hash_prospective_leaf(&key_1, &EMPTY);
+        let old_hash = SimpleSmt::hash_leaf(&smt.get_leaf(&key_1));
+        smt.insert(key_1, EMPTY);
+        assert_eq!(old_hash, smt.hash_prospective_leaf(&key_1, &value_1));
+        assert_eq!(SimpleSmt::hash_leaf(&smt.get_leaf(&key_1)), prospective_empty);
+    }
+}
+
 // HELPER FUNCTIONS
 // --------------------------------------------------------------------------------------------