diff --git a/src/merkle/smt/full/leaf.rs b/src/merkle/smt/full/leaf.rs index 095a4fb..585fc40 100644 --- a/src/merkle/smt/full/leaf.rs +++ b/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 { +pub(crate) fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator { 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 /// 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(); diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 5cd510b..69fa9e6 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -263,6 +263,28 @@ impl SparseMerkleTree for Smt { leaf.hash() } + fn construct_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 { let most_significant_felt = key[3]; LeafIndex::new_max_depth(most_significant_felt.as_int()) diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 66cd203..1e605a3 100644 --- a/src/merkle/smt/full/tests.rs +++ b/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,119 @@ fn test_smt_removal() { } } +/// This tests that we can correctly calculate prospective leaves -- that is, we can construct +/// correct [`SmtLeaf`] values for a theoretical insertion on a Merkle tree without mutating or +/// cloning the tree. +#[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)]); + // Sort key_3 before key_1, to test non-append insertion. + let key_3: RpoDigest = + RpoDigest::from([0_u32.into(), 0_u32.into(), 0_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.construct_prospective_leaf(smt.get_leaf(&key_1), &key_1, &value_1).hash(); + smt.insert(key_1, value_1); + + let leaf = smt.get_leaf(&key_1); + assert_eq!( + prospective, + leaf.hash(), + "prospective hash for leaf {leaf:?} did not match actual hash", + ); + } + + // insert key-value 2 + { + let prospective = + smt.construct_prospective_leaf(smt.get_leaf(&key_2), &key_2, &value_2).hash(); + smt.insert(key_2, value_2); + + let leaf = smt.get_leaf(&key_2); + assert_eq!( + prospective, + leaf.hash(), + "prospective hash for leaf {leaf:?} did not match actual hash", + ); + } + + // insert key-value 3 + { + let prospective = + smt.construct_prospective_leaf(smt.get_leaf(&key_3), &key_3, &value_3).hash(); + smt.insert(key_3, value_3); + + let leaf = smt.get_leaf(&key_3); + assert_eq!( + prospective, + leaf.hash(), + "prospective hash for leaf {leaf:?} did not match actual hash", + ); + } + + // remove key 3 + { + let old_leaf = smt.get_leaf(&key_3); + let old_value_3 = smt.insert(key_3, EMPTY_WORD); + assert_eq!(old_value_3, value_3); + let prospective_leaf = + smt.construct_prospective_leaf(smt.get_leaf(&key_3), &key_3, &old_value_3); + + assert_eq!( + old_leaf.hash(), + prospective_leaf.hash(), + "removing and prospectively re-adding a leaf didn't yield the original leaf:\ + \n original leaf: {old_leaf:?}\ + \n prospective leaf: {prospective_leaf:?}", + ); + } + + // remove key 2 + { + let old_leaf = smt.get_leaf(&key_2); + let old_value_2 = smt.insert(key_2, EMPTY_WORD); + assert_eq!(old_value_2, value_2); + let prospective_leaf = + smt.construct_prospective_leaf(smt.get_leaf(&key_2), &key_2, &old_value_2); + + assert_eq!( + old_leaf.hash(), + prospective_leaf.hash(), + "removing and prospectively re-adding a leaf didn't yield the original leaf:\ + \n original leaf: {old_leaf:?}\ + \n prospective leaf: {prospective_leaf:?}", + ); + } + + // remove key 1 + { + let old_leaf = smt.get_leaf(&key_1); + let old_value_1 = smt.insert(key_1, EMPTY_WORD); + assert_eq!(old_value_1, value_1); + let prospective_leaf = + smt.construct_prospective_leaf(smt.get_leaf(&key_1), &key_1, &old_value_1); + assert_eq!( + old_leaf.hash(), + prospective_leaf.hash(), + "removing and prospectively re-adding a leaf didn't yield the original leaf:\ + \n original leaf: {old_leaf:?}\ + \n prospective leaf: {prospective_leaf:?}", + ); + } +} + /// 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() { diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index d7d42da..71e3ca8 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -167,6 +167,20 @@ pub(crate) trait SparseMerkleTree { /// Returns the hash of a leaf fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest; + /// Returns what `existing_leaf` would look like if `key` and `value` WERE inserted into the + /// tree, without mutating the tree itself. + /// + /// `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 construct_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 f1ff0dc..a168589 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -302,6 +302,15 @@ impl SparseMerkleTree for SimpleSmt { leaf.into() } + fn construct_prospective_leaf( + &self, + _existing_leaf: Word, + _key: &LeafIndex, + value: &Word, + ) -> Word { + *value + } + fn key_to_leaf_index(key: &LeafIndex) -> LeafIndex { *key }