From affbcf30ed8a12e3fae5519eff082f05380ba8f7 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Mon, 12 Aug 2024 15:56:07 -0600 Subject: [PATCH 1/3] feat(smt): impl constructing leaves that don't yet exist This commit implements 'prospective leaf construction' -- computing sparse Merkle tree leaves for a key-value insertion without actually performing that insertion. For SimpleSmt, this is trivial, since the leaf type is simply the value being inserted. For the full Smt, the new leaf payload depends on the existing payload in that leaf. Since almost all leaves are very small, we can just clone the leaf and modify a copy. This will allow us to perform more general prospective changes on Merkle trees. --- src/merkle/smt/full/leaf.rs | 4 +- src/merkle/smt/full/mod.rs | 22 +++++++ src/merkle/smt/full/tests.rs | 109 ++++++++++++++++++++++++++++++++++- src/merkle/smt/mod.rs | 15 +++++ src/merkle/smt/simple/mod.rs | 9 +++ 5 files changed, 156 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..70b7944 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,113 @@ 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.get_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.get_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.get_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.get_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.get_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.get_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..fd9e660 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -167,6 +167,21 @@ 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. + #[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 e430c3096e28a703c61b550a96ce4732d62cbd5d Mon Sep 17 00:00:00 2001 From: Qyriad Date: Wed, 21 Aug 2024 13:22:11 -0600 Subject: [PATCH 2/3] 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 2c38256..74810c3 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 fd9e660..dce9ef4 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 66c2a62..c0b2ae1 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 6f240076570398a497051022fcfd8ae689cca514 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Wed, 4 Sep 2024 15:03:24 -0600 Subject: [PATCH 3/3] 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, 428 insertions(+), 13 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 74810c3..09e4b07 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 70b7944..ebfe45a 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -365,6 +365,81 @@ fn test_prospective_hash() { } } +/// This tests that we can perform prospective changes correctly. +#[test] +fn test_prospective_insertion() { + 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]; + + let root_empty = smt.root(); + + let root_1 = { + smt.insert(key_1, value_1); + smt.root() + }; + + let root_2 = { + smt.insert(key_2, value_2); + smt.root() + }; + + let root_3 = { + smt.insert(key_3, value_3); + smt.root() + }; + + // Test incremental updates. + + let mut smt = Smt::default(); + + assert_eq!(smt.root(), root_empty); + 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); + mutations.insert(key_3, value_3); + assert_eq!(mutations.root(), root_3); + // Inserting an empty value should bring us back to root_2 again. + mutations.insert(key_3, EMPTY_WORD); + 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); + + // Test batch updates. + 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 #[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 dce9ef4..5f035ef 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 // --------------------------------------------------------------------------------------------- @@ -178,7 +236,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, @@ -200,7 +257,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, } @@ -263,3 +320,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 alloc::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 alloc::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 c0b2ae1..1869bf2 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`. ///