From eb407cd8fcf923aca5809178c37d2a52333a8681 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Mon, 12 Aug 2024 15:56:07 -0600 Subject: [PATCH 1/4] feat(smt): impl constructing leaves that don't yet exist (WIP) --- src/merkle/smt/full/leaf.rs | 4 +- src/merkle/smt/full/mod.rs | 22 ++++++++ src/merkle/smt/full/tests.rs | 101 ++++++++++++++++++++++++++++++++++- src/merkle/smt/mod.rs | 12 +++++ src/merkle/smt/simple/mod.rs | 9 ++++ 5 files changed, 145 insertions(+), 3 deletions(-) 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..2c38256 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 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 { 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..4a95a9e 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,105 @@ 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)]); + // 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.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_3, value_3), (key_1, value_1), (key_2, value_2)]), + ); + } + + // 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() { diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index d7d42da..a2f3735 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -167,6 +167,18 @@ pub(crate) trait SparseMerkleTree { /// 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; diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index f1ff0dc..66c2a62 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 get_prospective_leaf( + &self, + _existing_leaf: Word, + _key: &LeafIndex, + value: &Word, + ) -> Word { + *value + } + fn key_to_leaf_index(key: &LeafIndex) -> LeafIndex { *key } From 04ef7205d539508ed539f8bab89848ee1ab047e0 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Thu, 29 Aug 2024 16:30:37 -0600 Subject: [PATCH 2/4] feat(smt): make leaf-emptiness knowable at the trait level --- src/merkle/smt/full/mod.rs | 7 +++++++ src/merkle/smt/mod.rs | 7 +++++++ src/merkle/smt/simple/mod.rs | 4 ++++ 3 files changed, 18 insertions(+) diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 2c38256..adf0a0a 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -285,6 +285,13 @@ impl SparseMerkleTree for Smt { } } + fn is_leaf_empty(leaf: &SmtLeaf) -> bool { + match leaf { + SmtLeaf::Empty(_) => true, + _ => false, + } + } + 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/mod.rs b/src/merkle/smt/mod.rs index a2f3735..c23c855 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -167,6 +167,13 @@ pub(crate) trait SparseMerkleTree { /// Returns the hash of a leaf fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest; + /// Returns true if a leaf is completely sparse and does not hold any key-value pairs. + #[allow(dead_code)] + fn is_leaf_empty(leaf: &Self::Leaf) -> bool; + + /// 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 diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 66c2a62..6ab2f00 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -311,6 +311,10 @@ impl SparseMerkleTree for SimpleSmt { *value } + fn is_leaf_empty(leaf: &Word) -> bool { + *leaf == Self::EMPTY_VALUE + } + fn key_to_leaf_index(key: &LeafIndex) -> LeafIndex { *key } From 32d96f0dcec3061a4d9e942d408aafcf9bc8f3b0 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Wed, 21 Aug 2024 13:22:11 -0600 Subject: [PATCH 3/4] feat(smt): export get_value() in the trait --- src/merkle/smt/full/mod.rs | 16 ++++++++++------ src/merkle/smt/mod.rs | 4 ++++ src/merkle/smt/simple/mod.rs | 4 ++++ 3 files changed, 18 insertions(+), 6 deletions(-) diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index adf0a0a..85051bd 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -121,12 +121,7 @@ impl Smt { /// Returns the value associated with `key` pub fn get_value(&self, key: &RpoDigest) -> Word { - let leaf_pos = LeafIndex::::from(*key).value(); - - match self.leaves.get(&leaf_pos) { - Some(leaf) => leaf.get_value(key).unwrap_or_default(), - None => EMPTY_WORD, - } + >::get_value(self, key) } /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle @@ -250,6 +245,15 @@ impl SparseMerkleTree for Smt { } } + fn get_value(&self, key: &Self::Key) -> Self::Value { + let leaf_pos = LeafIndex::::from(*key).value(); + + match self.leaves.get(&leaf_pos) { + Some(leaf) => leaf.get_value(key).unwrap_or_default(), + None => EMPTY_WORD, + } + } + fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf { let leaf_pos = LeafIndex::::from(*key).value(); diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index c23c855..4c9a175 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -161,6 +161,10 @@ pub(crate) trait SparseMerkleTree { /// Inserts a leaf node, and returns the value at the key if already exists fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option; + /// Returns the value at the specified key. Recall that by definition, any key that hasn't been + /// updated is associated with [`Self::EMPTY_VALUE`]. + fn get_value(&self, key: &Self::Key) -> Self::Value; + /// Returns the leaf at the specified index. fn get_leaf(&self, key: &Self::Key) -> Self::Leaf; diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 6ab2f00..4d27b5c 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -289,6 +289,10 @@ impl SparseMerkleTree for SimpleSmt { } } + fn get_value(&self, key: &LeafIndex) -> Word { + self.get_leaf(key) + } + fn get_leaf(&self, key: &LeafIndex) -> Word { let leaf_pos = key.value(); match self.leaves.get(&leaf_pos) { From 2d9317ead1ca5b469fcfbb3479919c5476c69e00 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Wed, 4 Sep 2024 15:03:24 -0600 Subject: [PATCH 4/4] WIP: feat(smt): implement generic prospective insertions This commit adds a type, MutationSet, which represents a set of changes to a SparseMerkleTree that haven't happened yet, and can be queried on to ensure a set of insertions result in the correct tree root before finalizing and committing the mutation. This is a direct step towards issue 222, and will directly enable removing Merkle tree clones in miden-node InnerState::apply_block(). As part of this change, SparseMerkleTree now requires its Key to be Ord and its Leaf to be Clone (both bounds which were already met by existing implementations). The Ord bound could instead be changed to Eq + Hash, if MutationSet were changed to use a HashMap instead of a BTreeMap. Additionally, as MutationSet is a generic type which works on any type that implements SparseMerkleTree, but is intended for public API use, the SparseMerkleTree trait and InnerNode type have been made public so MutationSet can be used outside of this crate. --- src/merkle/mod.rs | 4 +- src/merkle/smt/full/mod.rs | 37 ++++- src/merkle/smt/full/tests.rs | 75 ++++++++- src/merkle/smt/mod.rs | 286 ++++++++++++++++++++++++++++++++++- src/merkle/smt/simple/mod.rs | 39 ++++- 5 files changed, 421 insertions(+), 20 deletions(-) diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 8954d4d..a562aa5 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -22,8 +22,8 @@ pub use path::{MerklePath, RootPath, ValuePath}; mod smt; pub use smt::{ - LeafIndex, SimpleSmt, Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError, SMT_DEPTH, - SMT_MAX_DEPTH, SMT_MIN_DEPTH, + LeafIndex, MutationSet, SimpleSmt, Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError, + SMT_DEPTH, SMT_MAX_DEPTH, SMT_MIN_DEPTH, }; mod mmr; diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 85051bd..e5b1e31 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -5,8 +5,8 @@ use alloc::{ }; use super::{ - EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, MerklePath, - NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD, + DoneMutationSet, EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, + MerklePath, MutationSet, NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD, }; mod error; @@ -167,6 +167,39 @@ impl Smt { >::insert(self, key, value) } + /// Start a prospective mutation transaction, which can be queried, discarded, or applied. + /// + /// This method returns a [`MutationSet`], which you can call [`MutationSet::insert()`] on to + /// perform prospective mutations to this Merkle tree, and check the computed root hash given + /// those mutations with [`MutationSet::root()`]. When you are done, call + /// [`SparseMerkleTree::commit()`] with the return value of [`MutationSet::done()`] to apply the + /// prospective mutations to this tree. Or to discard the changes, simply [`drop`] the + /// [`MutationSet`]. + /// + /// ## Example + /// ``` + /// # use miden_crypto::{hash::rpo::RpoDigest, Felt, Word}; + /// # use miden_crypto::merkle::{Smt, EmptySubtreeRoots, SMT_DEPTH}; + /// let mut smt = Smt::new(); + /// let mut mutations = smt.mutate(); + /// mutations.insert(RpoDigest::default(), Word::default()); + /// assert_eq!(mutations.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0)); + /// smt.commit(mutations.done()); + /// assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0)); + /// ``` + pub fn mutate(&self) -> MutationSet { + >::mutate(self) + } + + /// Apply prospective mutations started with [`Smt::mutate()`] to this tree. + /// + /// This method takes the return value of [`MutationSet::done()`], and applies the changes + /// represented in that [`MutationSet`] to this Merkle tree. See [`Smt::mutate()`] for more + /// details. + pub fn commit(&mut self, mutations: DoneMutationSet) { + >::commit(self, mutations) + } + // HELPERS // -------------------------------------------------------------------------------------------- diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 4a95a9e..3e69196 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -259,7 +259,7 @@ fn test_smt_removal() { } #[test] -fn test_prospective_hash() { +fn test_prospective_insertion() { let mut smt = Smt::default(); let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64; @@ -275,8 +275,10 @@ fn test_prospective_hash() { let value_2 = [2_u32.into(); WORD_SIZE]; let value_3: [Felt; 4] = [3_u32.into(); WORD_SIZE]; + let root_empty = smt.root(); + // insert key-value 1 - { + let root_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); @@ -284,10 +286,12 @@ fn test_prospective_hash() { assert_eq!(prospective, smt.get_leaf(&key_1).hash()); assert_eq!(smt.get_leaf(&key_1), SmtLeaf::Single((key_1, value_1))); - } + + smt.root() + }; // insert key-value 2 - { + let root_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); @@ -298,10 +302,12 @@ fn test_prospective_hash() { smt.get_leaf(&key_2), SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)]), ); - } + + smt.root() + }; // insert key-value 3 - { + let root_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); @@ -312,7 +318,9 @@ fn test_prospective_hash() { smt.get_leaf(&key_3), SmtLeaf::Multiple(vec![(key_3, value_3), (key_1, value_1), (key_2, value_2)]), ); - } + + smt.root() + }; // remove key 3 { @@ -355,6 +363,59 @@ fn test_prospective_hash() { assert_eq!(smt.get_leaf(&key_1), SmtLeaf::new_empty(key_1.into())); } + + // Test incremental updates. + let mut mutations = smt.mutate(); + mutations.insert(key_1, value_1); + assert_eq!(mutations.root(), root_1); + smt.commit(mutations.done()); + assert_eq!(smt.root(), root_1); + + let mut mutations = smt.mutate(); + mutations.insert(key_2, value_2); + assert_eq!(mutations.root(), root_2); + smt.commit(mutations.done()); + assert_eq!(smt.root(), root_2); + + let mut mutations = smt.mutate(); + mutations.insert(key_3, value_3); + assert_eq!(mutations.root(), root_3); + smt.commit(mutations.done()); + assert_eq!(smt.root(), root_3); + + let mut mutations = smt.mutate(); + mutations.insert(key_3, EMPTY_WORD); + assert_eq!(mutations.root(), root_2); + smt.commit(mutations.done()); + assert_eq!(smt.root(), root_2); + + // Test batch updates. + let mut smt = Smt::default(); + + let mut mutations = smt.mutate(); + mutations.insert(key_1, value_1); + assert_eq!(mutations.root(), root_1); + mutations.insert(key_2, value_2); + assert_eq!(mutations.root(), root_2); + mutations.insert(key_3, value_2); + assert_ne!(mutations.root(), root_2); + assert_ne!(mutations.root(), root_3); + mutations.insert(key_3, value_3); + assert_eq!(mutations.root(), root_3); + + smt.commit(mutations.done()); + assert_eq!(smt.root(), root_3); + + let mut mutations = smt.mutate(); + mutations.insert(key_3, EMPTY_WORD); + assert_eq!(mutations.root(), root_2); + // Ensure the order we remove these doesn't matter. + mutations.insert(key_1, EMPTY_WORD); + mutations.insert(key_2, EMPTY_WORD); + assert_eq!(mutations.root(), root_empty); + + smt.commit(mutations.done()); + assert_eq!(smt.root(), root_empty); } /// Tests that 2 key-value pairs stored in the same leaf have the same path diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 4c9a175..a53cc92 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -1,4 +1,4 @@ -use alloc::vec::Vec; +use alloc::{collections::BTreeMap, vec::Vec}; use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex}; use crate::{ @@ -43,13 +43,13 @@ pub const SMT_MAX_DEPTH: u8 = 64; /// must accomodate all keys that map to the same leaf. /// /// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs. -pub(crate) trait SparseMerkleTree { +pub trait SparseMerkleTree { /// The type for a key - type Key: Clone; + type Key: Clone + Ord; /// The type for a value type Value: Clone + PartialEq; /// The type for a leaf - type Leaf; + type Leaf: Clone; /// The type for an opening (i.e. a "proof") of a leaf type Opening; @@ -140,6 +140,64 @@ pub(crate) trait SparseMerkleTree { self.set_root(node_hash); } + /// Start a prospective mutation transaction, which can be queried, discarded, or applied. + /// + /// This method returns a [`MutationSet`], which you can call [`MutationSet::insert()`] on to + /// perform prospective mutations to this Merkle tree, and check the computed root hash given + /// those mutations with [`MutationSet::root()`]. When you are done, call + /// [`SparseMerkleTree::commit()`] with the return value of [`MutationSet::done()`] to apply the + /// prospective mutations to this tree. Or to discard the changes, simply [`drop`] the + /// [`MutationSet`]. + /// + /// ## Example + /// ``` + /// # use miden_crypto::{hash::rpo::RpoDigest, Felt}; + /// # use miden_crypto::merkle::{Smt, EmptySubtreeRoots, SMT_DEPTH}; + /// # let mut smt = Smt::default(); + /// let mut mutations = smt.mutate(); + /// mutations.insert(Default::default(), Default::default()); + /// assert_eq!(mutations.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0)); + /// smt.commit(mutations.done()); + /// assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0)); + /// ``` + fn mutate(&self) -> MutationSet + where + Self: Sized, + { + MutationSet { + tree: self, + node_mutations: Default::default(), + new_pairs: Default::default(), + new_root: self.root(), + } + } + + /// Apply prospective mutations started with [`SparseMerkleTree::mutate()`] to this tree. + /// + /// This method takes the return value of [`MutationSet::done()`], and applies the changes + /// represented in that [`MutationSet`] to this Merkle tree. See [`SparseMerkleTree::mutate()`] + /// for more details. + fn commit(&mut self, mutations: DoneMutationSet) + where + Self: Sized, + { + use NodeMutation::*; + let DoneMutationSet { node_mutations, new_pairs, new_root } = mutations; + + for (index, mutation) in node_mutations { + match mutation { + Removal => self.remove_inner_node(index), + Addition(node) => self.insert_inner_node(index, node), + } + } + + for (key, value) in new_pairs { + self.insert_value(key, value); + } + + self.set_root(new_root); + } + // REQUIRED METHODS // --------------------------------------------------------------------------------------------- @@ -182,7 +240,6 @@ pub(crate) trait SparseMerkleTree { /// 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, @@ -204,7 +261,7 @@ pub(crate) trait SparseMerkleTree { #[derive(Debug, Default, Clone, PartialEq, Eq)] #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] -pub(crate) struct InnerNode { +pub struct InnerNode { pub left: RpoDigest, pub right: RpoDigest, } @@ -267,3 +324,220 @@ impl TryFrom for LeafIndex { Self::new(node_index.value()) } } + +// MUTATIONS +// ================================================================================================ + +/// A change to an inner node of a [`SparseMerkleTree`] that hasn't been applied yet. +/// [`MutationSet`] stores this type in relation to a [`NodeIndex`] to keep track of what changes +/// need to occur at what node indices. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum NodeMutation { + Removal, + Addition(InnerNode), +} + +/// Represents prospective mutations to a [`SparseMerkleTree`], created by +/// [`SparseMerkleTree::mutate()`], and applied with [`SparseMerkleTree::commit()`]. +/// +/// `T` is the Merkle tree type this is prospectively mutating. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MutationSet<'t, const DEPTH: u8, T> +where + T: SparseMerkleTree, +{ + /// A reference to the tree we're mutating. We need to be able to query the tree's existing + /// nodes and leaves to calculate prospective nodes, leaves, and roots. + tree: &'t T, + /// The set of nodes that need to be removed or added. The "effective" node at an index is the + /// Merkle tree's existing node at that index, with the [`NodeMutation`] in this map at that + /// index overlayed, if any. + node_mutations: BTreeMap, + /// The set of top-level key-value pairs we're prospectively adding to the tree, including + /// adding empty values. The "effective" value for a key is the value in this BTreeMap, falling + /// back to the existing value in the Merkle tree. + new_pairs: BTreeMap, + /// The currently-calculated root for the Merkle tree, given these mutations. When this struct + /// is constructed, `new_root` is set to the existing root, since there aren't any changes yet. + new_root: RpoDigest, +} + +impl<'t, const DEPTH: u8, T> MutationSet<'t, DEPTH, T> +where + T: SparseMerkleTree, +{ + /// Pass the return value of this method to [`SparseMerkleTree::commit()`] to apply these + /// mutations to the Merkle tree. + pub fn done(self) -> DoneMutationSet { + // Get rid of the `tree` reference, so we can pass the mutation info to a `&mut self` method + // on `SparseMerkleTree`. + let Self { node_mutations, new_pairs, new_root, .. } = self; + DoneMutationSet { node_mutations, new_pairs, new_root } + } + + /// Get the prospective root hash for if these mutations were applied to the Merkle tree. + pub fn root(&self) -> RpoDigest { + self.new_root + } + + /// Add a new key-value pair to the set of prospective changes to the Merkle tree. + /// + /// After `insert()`, you may call [`MutationSet::root()`] to query what the new calculated root + /// hash of the Merkle tree would be after all requested insertions. Recall that semantically + /// removing a value is done by inserting [`SparseMerkleTree::EMPTY_VALUE`]. + pub fn insert(&mut self, key: T::Key, value: T::Value) { + // This function's calculations are eager, so multiple inserts that affect the same nodes + // will calculate those node's hashes each time. Future work could lazily calculate node + // hashes all at once instead. + + // If the old value and the new value are the same, there is nothing to update. + let old_value = self.get_effective_value(&key); + if value == old_value { + return; + } + + let mut node_index: NodeIndex = T::key_to_leaf_index(&key).into(); + + let old_leaf = self.get_effective_leaf(&key); + let new_leaf = self.tree.get_prospective_leaf(old_leaf, &key, &value); + + let mut new_child_hash = T::hash_leaf(&new_leaf); + + for node_depth in (0..node_index.depth()).rev() { + // Whether the node we're replacing is the right child or left child. + let is_right = node_index.is_value_odd(); + node_index.move_up(); + + let old_node = self.get_effective_node(node_index); + + let new_node = if is_right { + InnerNode { + left: old_node.left, + right: new_child_hash, + } + } else { + InnerNode { + left: new_child_hash, + right: old_node.right, + } + }; + + // The next iteration will operate on this new node's hash. + new_child_hash = new_node.hash(); + + let &equivalent_empty_hash = EmptySubtreeRoots::entry(DEPTH, node_depth); + if new_child_hash == equivalent_empty_hash { + self.remove_inner_node(node_index); + } else { + self.insert_inner_node(node_index, new_node); + } + } + + // Once we're at depth 0, the last node we made is the new root. + self.new_root = new_child_hash; + self.new_pairs.insert(key, value); + } + + // Implementation details. + + fn get_effective_value(&self, key: &T::Key) -> T::Value { + self.new_pairs.get(key).cloned().unwrap_or_else(|| self.tree.get_value(key)) + } + + fn get_effective_leaf(&self, key: &T::Key) -> T::Leaf { + let leaf_key = key; + let leaf_index = T::key_to_leaf_index(leaf_key); + let pairs_at_index = self + .new_pairs + .iter() + .filter(|&(new_key, _)| T::key_to_leaf_index(new_key) == leaf_index); + + let initial_acc = self.tree.get_leaf(key); + pairs_at_index.fold(initial_acc, |acc, (k, v)| { + // In practice this should only run once (if that), most of the time. + let existing_leaf = acc.clone(); + self.tree.get_prospective_leaf(existing_leaf, k, v) + }) + } + + fn get_effective_node(&self, index: NodeIndex) -> InnerNode { + use NodeMutation::*; + + self.node_mutations + .get(&index) + .map(|mutation| match mutation { + Addition(node) => node.clone(), + Removal => { + let &child = EmptySubtreeRoots::entry(DEPTH, index.depth() + 1); + InnerNode { left: child, right: child } + }, + }) + .unwrap_or_else(|| self.tree.get_inner_node(index)) + } + + fn remove_inner_node(&mut self, index: NodeIndex) { + use std::collections::btree_map::Entry::*; + use NodeMutation::*; + + match self.node_mutations.entry(index) { + Vacant(entry) => { + entry.insert(Removal); + }, + Occupied(mut entry) => match entry.get_mut() { + // If we have an addition with this index, we don't care about what value it has, we + // just need to convert it to a removal instead. + Addition(_) => { + entry.insert(Removal); + }, + Removal => (), + }, + } + } + + fn insert_inner_node(&mut self, index: NodeIndex, new_node: InnerNode) { + use std::collections::btree_map::Entry::*; + use NodeMutation::*; + + match self.node_mutations.entry(index) { + Vacant(entry) => { + entry.insert(Addition(new_node)); + }, + Occupied(mut entry) => match entry.get_mut() { + Addition(existing) => { + // If there's an existing addition with this key, then overwrite it to be an + // addition of this new node instead. + *existing = new_node; + }, + Removal => { + // Likewise a removal of this key gets overwritten with an addition. + entry.insert(Addition(new_node)); + }, + }, + } + } +} + +/// Helper type that represents a [`MutationSet`] that's ready to be applied to a [`SparseMerkleTree`]. +/// +/// It is created by [`MutationSet::done()`], which should be directly passed to +/// [`SparseMerkleTree::commit()`]: +/// ``` +/// # use miden_crypto::{hash::rpo::RpoDigest, merkle::Smt, Felt}; +/// # let mut smt = Smt::default(); +/// # let key = RpoDigest::default(); +/// # let value: [Felt; 4] = Default::default(); +/// let mut mutations = smt.mutate(); +/// mutations.insert(key, value); +/// smt.commit(mutations.done()); +/// ``` +// This type exists for the sake of the borrow checker -- SparseMerkleTree::commit() needs a +// mutable reference to `self`, but MutationSet stores a shared reference to the SparseMerkleTree, +// and those can't both exist at the same time. By going through this type first, which stores all +// the same data as MutationSet except for the shared reference to SparseMerkleTree, the shared +// reference is dropped just before SparseMerkleTree::commit() is called, and the borrow checker is +// happy. Interior mutability would also work, but then we would lose static lifetime checking. +pub struct DoneMutationSet { + node_mutations: BTreeMap, + new_pairs: BTreeMap, + new_root: RpoDigest, +} diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 4d27b5c..9809ad3 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -1,9 +1,9 @@ use alloc::collections::{BTreeMap, BTreeSet}; use super::{ - super::ValuePath, EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, - MerklePath, NodeIndex, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD, SMT_MAX_DEPTH, - SMT_MIN_DEPTH, + super::ValuePath, DoneMutationSet, EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, + MerkleError, MerklePath, MutationSet, NodeIndex, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD, + SMT_MAX_DEPTH, SMT_MIN_DEPTH, }; #[cfg(test)] @@ -188,6 +188,39 @@ impl SimpleSmt { >::insert(self, key, value) } + /// Start a prospective mutation transaction, which can be queried, discarded, or applied. + /// + /// This method returns a [`MutationSet`], which you can call [`MutationSet::insert()`] on to + /// perform prospective mutations to this Merkle tree, and check the computed root hash given + /// those mutations with [`MutationSet::root()`]. When you are done, call + /// [`SparseMerkleTree::commit()`] with the return value of [`MutationSet::done()`] to apply the + /// prospective mutations to this tree. Or to discard the changes, simply [`drop`] the + /// [`MutationSet`]. + /// + /// ## Example + /// ``` + /// # use miden_crypto::{hash::rpo::RpoDigest, Felt, Word}; + /// # use miden_crypto::merkle::{LeafIndex, SimpleSmt, EmptySubtreeRoots, SMT_DEPTH}; + /// let mut smt: SimpleSmt<3> = SimpleSmt::new().unwrap(); + /// let mut mutations = smt.mutate(); + /// mutations.insert(LeafIndex::default(), Word::default()); + /// assert_eq!(mutations.root(), *EmptySubtreeRoots::entry(3, 0)); + /// smt.commit(mutations.done()); + /// assert_eq!(smt.root(), *EmptySubtreeRoots::entry(3, 0)); + /// ``` + pub fn mutate(&self) -> MutationSet { + >::mutate(self) + } + + /// Apply prospective mutations started with [`SimpleSmt::mutate()`] to this tree. + /// + /// This method takes the return value of [`MutationSet::done()`], and applies the changes + /// represented in that [`MutationSet`] to this Merkle tree. See [`SimpleSmt::mutate()`] for + /// more details. + pub fn commit(&mut self, mutations: DoneMutationSet, Word>) { + >::commit(self, mutations) + } + /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is /// computed as `DEPTH - SUBTREE_DEPTH`. ///