From a8fe97afe97d1e613f7e568eb177d6abcfe9ed18 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Mon, 12 Aug 2024 15:56:07 -0600 Subject: [PATCH] WIP: implement hash_prospective_leaf() --- src/merkle/smt/full/leaf.rs | 4 +- src/merkle/smt/full/mod.rs | 45 ++++++++++++++++++ src/merkle/smt/full/tests.rs | 92 +++++++++++++++++++++++++++++++++++- src/merkle/smt/mod.rs | 8 ++++ src/merkle/smt/simple/mod.rs | 4 ++ 5 files changed, 150 insertions(+), 3 deletions(-) diff --git a/src/merkle/smt/full/leaf.rs b/src/merkle/smt/full/leaf.rs index 23e1ee4..e82a3f3 100644 --- a/src/merkle/smt/full/leaf.rs +++ b/src/merkle/smt/full/leaf.rs @@ -349,7 +349,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(); @@ -358,7 +358,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 52b416f..1f43634 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -262,6 +262,29 @@ impl SparseMerkleTree for Smt { leaf.hash() } + fn hash_prospective_leaf(&self, key: &RpoDigest, value: &Word) -> RpoDigest { + // If this key already has a value, then the hash will be based off a + // prospective mutation on the leaf. + let leaf_index: LeafIndex = Self::key_to_leaf_index(&key); + match self.leaves.get(&leaf_index.value()) { + Some(existing_leaf) => { + if value == &Self::EMPTY_VALUE { + // A leaf with an empty value is conceptually a removal the + // value in that leaf with this key. + // TODO: avoid cloning the leaf. + let mut cloned = existing_leaf.clone(); + cloned.remove(*key); + return cloned.hash(); + } + // TODO: avoid cloning the leaf. + let mut cloned = existing_leaf.clone(); + cloned.insert(*key, *value); + cloned.hash() + }, + None => 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()) @@ -356,3 +379,25 @@ fn test_smt_serialization_deserialization() { let bytes = smt.to_bytes(); assert_eq!(smt, Smt::read_from_bytes(&bytes).unwrap()); } + +#[test] +fn test_prospective_hash() { + // Smt with values + let smt_leaves_2: [(RpoDigest, Word); 2] = [ + ( + RpoDigest::new([Felt::new(101), Felt::new(102), Felt::new(103), Felt::new(104)]), + [Felt::new(1_u64), Felt::new(2_u64), Felt::new(3_u64), Felt::new(4_u64)], + ), + ( + RpoDigest::new([Felt::new(105), Felt::new(106), Felt::new(107), Felt::new(108)]), + [Felt::new(5_u64), Felt::new(6_u64), Felt::new(7_u64), Felt::new(8_u64)], + ), + ]; + let smt = Smt::with_entries(smt_leaves_2).unwrap(); + + for (key, value) in &smt_leaves_2 { + let expected = smt.get_leaf(key).hash(); + let actual = smt.hash_prospective_leaf(key, value); + assert_eq!(expected, actual); + } +} diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index e852811..27d24bd 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, }; @@ -257,6 +257,96 @@ 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.hash_prospective_leaf(&key_1, &value_1); + let old_value_1 = smt.insert(key_1, value_1); + assert_eq!(old_value_1, EMPTY_WORD); + + assert_eq!(smt.get_leaf(&key_1).hash(), prospective); + + assert_eq!(smt.get_leaf(&key_1), SmtLeaf::Single((key_1, value_1))); + + } + + // insert key-value 2 + { + let prospective = smt.hash_prospective_leaf(&key_2, &value_2); + let old_value_2 = smt.insert(key_2, value_2); + assert_eq!(old_value_2, EMPTY_WORD); + + assert_eq!(smt.get_leaf(&key_2).hash(), prospective); + + assert_eq!( + smt.get_leaf(&key_2), + SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)]) + ); + } + + // insert key-value 3 + { + let prospective_hash = smt.hash_prospective_leaf(&key_3, &value_3); + let old_value_3 = smt.insert(key_3, value_3); + assert_eq!(old_value_3, EMPTY_WORD); + + assert_eq!(smt.get_leaf(&key_3).hash(), prospective_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.hash_prospective_leaf(&key_3, &old_value_3)); + + 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.hash_prospective_leaf(&key_2, &old_value_2)); + + 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.hash_prospective_leaf(&key_1, &old_value_1)); + + 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() { diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index bd6a6e9..bcfdb7b 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -185,6 +185,14 @@ 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. + 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 ef26322..a6fe794 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -311,6 +311,10 @@ impl SparseMerkleTree for SimpleSmt { leaf.into() } + fn hash_prospective_leaf(&self, _key: &LeafIndex, value: &Word) -> RpoDigest { + Self::hash_leaf(value) + } + fn key_to_leaf_index(key: &LeafIndex) -> LeafIndex { *key }