diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 52b416f..2de5786 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -262,6 +262,32 @@ impl SparseMerkleTree 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 { 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 e852811..644584c 100644 --- a/src/merkle/smt/full/tests.rs +++ b/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 // -------------------------------------------------------------------------------------------- diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 52ed1d2..c6699ef 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -166,6 +166,15 @@ pub(crate) trait SparseMerkleTree { /// 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; diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 2fa5ae4..9f0fed1 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -301,6 +301,10 @@ impl SparseMerkleTree for SimpleSmt { leaf.into() } + fn hash_prospective_leaf(&self, _: &LeafIndex, value: &Word) -> RpoDigest { + Self::hash_leaf(value) + } + fn key_to_leaf_index(key: &LeafIndex) -> LeafIndex { *key } diff --git a/src/merkle/smt/simple/tests.rs b/src/merkle/smt/simple/tests.rs index 96ec7df..8825713 100644 --- a/src/merkle/smt/simple/tests.rs +++ b/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::; + const EMPTY: Word = SimpleSmt::EMPTY_VALUE; + let mut smt = SimpleSmt::new().unwrap(); + + let key = |num: u64| LeafIndex::::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 // --------------------------------------------------------------------------------------------