feat(smt): impl constructing leaves that don't yet exist (WIP)

src/merkle/smt/full/leaf.rs

@@ -350,7 +350,7 @@ impl Deserializable for SmtLeaf {
 // ================================================================================================
 
 /// Converts a key-value tuple to an iterator of `Felt`s
-fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator<Item = Felt> {
+pub(crate) fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator<Item = Felt> {
     let key_elements = key.into_iter();
     let value_elements = value.into_iter();
 
@@ -359,7 +359,7 @@ fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator<Item = Felt>
 
 /// Compares two keys, compared element-by-element using their integer representations starting with
 /// the most significant element.
-fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering {
+pub(crate) fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering {
     for (v1, v2) in key_1.iter().zip(key_2.iter()).rev() {
         let v1 = v1.as_int();
         let v2 = v2.as_int();

src/merkle/smt/full/mod.rs

@@ -263,6 +263,28 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
         leaf.hash()
     }
 
+    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 {
+                    existing_leaf.insert(*key, *value);
+                } else {
+                    existing_leaf.remove(*key);
+                }
+
+                existing_leaf
+            },
+        }
+    }
+
     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

@@ -2,7 +2,7 @@ use alloc::vec::Vec;
 
 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,
 };
@@ -258,6 +258,104 @@ fn test_smt_removal() {
     }
 }
 
+#[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 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
+    let key_2: RpoDigest =
+        RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
+    let key_3: RpoDigest =
+        RpoDigest::from([3_u32.into(), 3_u32.into(), 3_u32.into(), Felt::new(raw)]);
+
+    let value_1 = [ONE; WORD_SIZE];
+    let value_2 = [2_u32.into(); WORD_SIZE];
+    let value_3: [Felt; 4] = [3_u32.into(); WORD_SIZE];
+
+    // insert key-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);
+
+        assert_eq!(prospective, smt.get_leaf(&key_1).hash());
+
+        assert_eq!(smt.get_leaf(&key_1), SmtLeaf::Single((key_1, value_1)));
+    }
+
+    // insert key-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);
+
+        assert_eq!(prospective, smt.get_leaf(&key_2).hash());
+
+        assert_eq!(
+            smt.get_leaf(&key_2),
+            SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)]),
+        );
+    }
+
+    // insert key-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);
+
+        assert_eq!(prospective, smt.get_leaf(&key_3).hash());
+
+        assert_eq!(
+            smt.get_leaf(&key_3),
+            SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2), (key_3, value_3)]),
+        );
+    }
+
+    // remove key 3
+    {
+        let old_hash = smt.get_leaf(&key_3).hash();
+        let old_value_3 = smt.insert(key_3, EMPTY_WORD);
+        assert_eq!(old_value_3, 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),
+            SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)]),
+        );
+    }
+
+    // remove key 2
+    {
+        let old_hash = smt.get_leaf(&key_2).hash();
+        let old_value_2 = smt.insert(key_2, EMPTY_WORD);
+        assert_eq!(old_value_2, 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)));
+    }
+
+    // remove key 1
+    {
+        let old_hash = smt.get_leaf(&key_1).hash();
+        let old_value_1 = smt.insert(key_1, EMPTY_WORD);
+        assert_eq!(old_value_1, 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()));
+    }
+}
+
 /// Tests that 2 key-value pairs stored in the same leaf have the same path
 #[test]
 fn test_smt_path_to_keys_in_same_leaf_are_equal() {

src/merkle/smt/mod.rs

@@ -167,6 +167,18 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     /// Returns the hash of a leaf
     fn hash_leaf(leaf: &Self::Leaf) -> 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.
+    #[cfg_attr(not(test), allow(dead_code))]
+    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

@@ -302,6 +302,15 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
         leaf.into()
     }
 
+    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> {
         *key
     }