From 13e80115caaf5d8d6803363eaea8c284073ef439 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Fri, 9 Aug 2024 13:05:43 -0600 Subject: [PATCH 1/3] chore: refresh Cargo.lock --- Cargo.lock | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e10bdb2..85a2fcc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -147,9 +147,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.1.7" +version = "1.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26a5c3fd7bfa1ce3897a3a3501d362b2d87b7f2583ebcb4a949ec25911025cbc" +checksum = "504bdec147f2cc13c8b57ed9401fd8a147cc66b67ad5cb241394244f2c947549" dependencies = [ "jobserver", "libc", @@ -190,9 +190,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.13" +version = "4.5.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fbb260a053428790f3de475e304ff84cdbc4face759ea7a3e64c1edd938a7fc" +checksum = "c937d4061031a6d0c8da4b9a4f98a172fc2976dfb1c19213a9cf7d0d3c837e36" dependencies = [ "clap_builder", "clap_derive", @@ -200,9 +200,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.13" +version = "4.5.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64b17d7ea74e9f833c7dbf2cbe4fb12ff26783eda4782a8975b72f895c9b4d99" +checksum = "85379ba512b21a328adf887e85f7742d12e96eb31f3ef077df4ffc26b506ffed" dependencies = [ "anstream", "anstyle", @@ -845,18 +845,18 @@ checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" [[package]] name = "serde" -version = "1.0.204" +version = "1.0.205" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bc76f558e0cbb2a839d37354c575f1dc3fdc6546b5be373ba43d95f231bf7c12" +checksum = "e33aedb1a7135da52b7c21791455563facbbcc43d0f0f66165b42c21b3dfb150" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.204" +version = "1.0.205" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e0cd7e117be63d3c3678776753929474f3b04a43a080c744d6b0ae2a8c28e222" +checksum = "692d6f5ac90220161d6774db30c662202721e64aed9058d2c394f451261420c1" dependencies = [ "proc-macro2", "quote", From f8d6af5effb9c2a783150d1954148af029aa4c8e Mon Sep 17 00:00:00 2001 From: Qyriad Date: Fri, 9 Aug 2024 17:26:29 -0600 Subject: [PATCH 2/3] refactor: make Smt's node recomputation pure And do mutations in its callers instead. --- src/merkle/smt/mod.rs | 44 +++++++++++++++++++++++++++++------- src/merkle/smt/simple/mod.rs | 12 +++++++++- 2 files changed, 47 insertions(+), 9 deletions(-) diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 52ed1d2..bd6a6e9 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -104,23 +104,41 @@ pub(crate) trait SparseMerkleTree { leaf_index.into() }; - self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf)); + let mut mutations = + self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf)); + for index in mutations.removals.drain(..) { + self.remove_inner_node(index); + } + + for (index, new_node) in mutations.additions.drain(..) { + self.insert_inner_node(index, new_node); + } + + self.set_root(mutations.new_root); old_value } /// Recomputes the branch nodes (including the root) from `index` all the way to the root. /// `node_hash_at_index` is the hash of the node stored at index. + /// + /// This method is pure, and only computes the mutations to apply. fn recompute_nodes_from_index_to_root( - &mut self, + &self, mut index: NodeIndex, node_hash_at_index: RpoDigest, - ) { + ) -> Mutations { let mut node_hash = node_hash_at_index; + + let mut removals: Vec = Vec::new(); + let mut additions: Vec<(NodeIndex, InnerNode)> = Vec::new(); + for node_depth in (0..index.depth()).rev() { let is_right = index.is_value_odd(); index.move_up(); + let InnerNode { left, right } = self.get_inner_node(index); + let (left, right) = if is_right { (left, node_hash) } else { @@ -129,14 +147,15 @@ pub(crate) trait SparseMerkleTree { node_hash = Rpo256::merge(&[left, right]); if node_hash == *EmptySubtreeRoots::entry(DEPTH, node_depth) { - // If a subtree is empty, when can remove the inner node, since it's equal to the - // default value - self.remove_inner_node(index) + // If a subtree is empty, we can remove the inner node, since it's equal to the + // default value. + removals.push(index); } else { - self.insert_inner_node(index, InnerNode { left, right }); + additions.push((index, InnerNode { left, right })); } } - self.set_root(node_hash); + + Mutations { removals, additions, new_root: node_hash } } // REQUIRED METHODS @@ -243,3 +262,12 @@ impl TryFrom for LeafIndex { Self::new(node_index.value()) } } + +// MUTATIONS +// ================================================================================================ + +pub(crate) struct Mutations { + removals: Vec, + additions: Vec<(NodeIndex, InnerNode)>, + new_root: RpoDigest, +} diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 2fa5ae4..ef26322 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -242,7 +242,17 @@ impl SimpleSmt { // recompute nodes starting from subtree root // -------------- - self.recompute_nodes_from_index_to_root(subtree_root_index, subtree.root); + let mut mutations = + self.recompute_nodes_from_index_to_root(subtree_root_index, subtree.root); + for index in mutations.removals.drain(..) { + self.remove_inner_node(index); + } + + for (index, new_node) in mutations.additions.drain(..) { + self.insert_inner_node(index, new_node); + } + + self.set_root(mutations.new_root); Ok(self.root) } From 3bc4f68e639679ed1f5179ea45988ba16d8467d8 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Mon, 12 Aug 2024 15:56:07 -0600 Subject: [PATCH 3/3] WIP: implement hash_prospective_leaf() --- src/merkle/smt/full/leaf.rs | 4 +- src/merkle/smt/full/mod.rs | 43 +++++++++++++++++ src/merkle/smt/full/tests.rs | 92 +++++++++++++++++++++++++++++++++++- src/merkle/smt/mod.rs | 8 ++++ src/merkle/smt/simple/mod.rs | 4 ++ 5 files changed, 148 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..a844898 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -262,6 +262,27 @@ 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 of that leaf, + // so this hash doesn't depend on any existing leaf with this key + // at all. + return SmtLeaf::new_empty(leaf_index).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 +377,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 }