From e7c36fd4fcd9b4eaa9b638c308162d6fe79d93c7 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 | 115 ++++++++++++++++++++++++++++++++++- src/merkle/smt/mod.rs | 18 ++++++ src/merkle/smt/simple/mod.rs | 9 +++ 5 files changed, 165 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..69fa9e6 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 construct_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..1e605a3 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,119 @@ 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.construct_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.construct_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.construct_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.construct_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.construct_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.construct_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..31a675a 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -167,6 +167,24 @@ pub(crate) trait SparseMerkleTree { /// Returns the hash of a leaf fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest; + /// Returns what a leaf would look like if a key-value pair were inserted into the tree, without + /// mutating the tree itself. The existing leaf can be empty. + /// + /// 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. + /// + /// # Invariants + /// Because this method is for a prospective key-value insertion into a specific leaf, + /// `existing_leaf` must have the same leaf index as `key` (as determined by + /// [`SparseMerkleTree::key_to_leaf_index()`]), or the result will be meaningless. + fn construct_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..a168589 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 construct_prospective_leaf( + &self, + _existing_leaf: Word, + _key: &LeafIndex, + value: &Word, + ) -> Word { + *value + } + fn key_to_leaf_index(key: &LeafIndex) -> LeafIndex { *key } From 4dc276a1d20f7e0ab232884b898c3734d6a9a468 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 69fa9e6..a9a0bf1 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 31a675a..c68161d 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 a168589..6922a32 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 f0eaf046e20e7721dad3a75dc54b5df1b9595fab 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. --- CHANGELOG.md | 1 + src/merkle/mod.rs | 4 +- src/merkle/smt/full/mod.rs | 37 ++++- src/merkle/smt/full/tests.rs | 74 +++++++++ src/merkle/smt/mod.rs | 301 ++++++++++++++++++++++++++++++++++- src/merkle/smt/simple/mod.rs | 38 ++++- 6 files changed, 443 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 76c65a7..c6d0b2e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ - [BREAKING]: renamed `Mmr::open()` into `Mmr::open_at()` and `Mmr::peaks()` into `Mmr::peaks_at()` (#234). - Added `Mmr::open()` and `Mmr::peaks()` which rely on `Mmr::open_at()` and `Mmr::peaks()` respectively (#234). - Standardised CI and Makefile across Miden repos (#323). +- Added `Smt::mutate()` and `Smt::commit()` for validation-checked insertions (#327). ## 0.10.0 (2024-08-06) diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 8954d4d..3be0a72 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, + SparseMerkleTree, 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 a9a0bf1..2aaefcd 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -5,8 +5,9 @@ use alloc::{ }; use super::{ - EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, MerklePath, - NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD, + CompletedMutationSet, EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, + MerkleError, MerklePath, MutationSet, NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, + EMPTY_WORD, }; mod error; @@ -167,6 +168,38 @@ impl Smt { >::insert(self, key, value) } + /// Start a prospective mutation transaction, which can be queried, discarded, or applied. + /// + /// This method returns a [`MutationSet`], on which you can call [`MutationSet::insert()`] 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 + /// [`MutationSet::done()`], and pass that to [`Smt::commit()`] 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 the prospective mutations began 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: CompletedMutationSet) { + >::commit(self, mutations) + } + // HELPERS // -------------------------------------------------------------------------------------------- diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 1e605a3..6c45389 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -371,6 +371,80 @@ 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(); + + let mut mutations = smt.mutate(); + mutations.insert(key_1, value_1); + assert_eq!(mutations.root(), root_1, "prospective root 1 did not match actual root 1"); + smt.commit(mutations.done()); + assert_eq!(smt.root(), root_1, "mutations before and after commit did not match"); + + let mut mutations = smt.mutate(); + mutations.insert(key_2, value_2); + assert_eq!(mutations.root(), root_2, "prospective root 2 did not match actual root 2"); + mutations.insert(key_3, value_3); + assert_eq!(mutations.root(), root_3, "mutations before and after commit did not match"); + // Inserting an empty value should bring us back to root_2 again. + mutations.insert(key_3, EMPTY_WORD); + assert_eq!(mutations.root(), root_2, "prospective removal did not undo prospective insert"); + smt.commit(mutations.done()); + assert_eq!(smt.root(), root_2, "mutations before and after commit did not match"); + + let mut mutations = smt.mutate(); + mutations.insert(key_3, value_3); + assert_eq!(mutations.root(), root_3, "prospective root 3 did not match actual root 3"); + smt.commit(mutations.done()); + assert_eq!(smt.root(), root_3, "mutations before and after commit did not match"); + + // Test batch updates. + let mut mutations = smt.mutate(); + mutations.insert(key_3, EMPTY_WORD); + assert_eq!(mutations.root(), root_2, "prospective removal did not undo actual insert"); + // 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, "prospective removals are not order-independent"); + smt.commit(mutations.done()); + assert_eq!(smt.root(), root_empty, "mutations before and after commit did not match"); +} + /// 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 c68161d..ae4a883 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,63 @@ 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`], on which you can call [`MutationSet::insert()`] 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 + /// [`MutationSet::done()`], and pass that to [`Smt::commit()`] to apply the prospective + /// mutations to this tree. Or, to discard the chanes, 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 the prospective mutations began 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: CompletedMutationSet) + where + Self: Sized, + { + use NodeMutation::*; + let CompletedMutationSet { 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 // --------------------------------------------------------------------------------------------- @@ -203,7 +260,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, } @@ -266,3 +323,237 @@ impl TryFrom for LeafIndex { Self::new(node_index.value()) } } + +// MUTATIONS +// ================================================================================================ + +/// A change to an inner node of a [`SparseMerkleTree`] that hasn't yet been applied. +/// [`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 a group of prospective mutations to a [`SparseMerkleTree`], created by +/// [`SparseMerkleTree::mutate()`], and that can be 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) -> CompletedMutationSet { + // 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; + CompletedMutationSet { 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.construct_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.construct_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 CompletedMutationSet { + /// Corresponds to MutationSet's `node_mutations` field, and constructed from it. + /// + /// This is the set of nodes that need to be removed or added. Each [`NodeMutation::Addition`] + /// corresponds to a [`SparseMerkleTree::add_inner_node()`] call, and each + /// [`NodeMutation::Removal`] corresponds to a [`SparseMerkleTree::remove_inner_node()`] call. + node_mutations: BTreeMap, + /// Corresponds to MutationSet's `new_pairs` field, and constructed from it. + /// + /// This is the set of key-value pairs that need to be inserted. Each entry corresponds to a + /// [`SparseMerkleTree::insert_value()`] call. + new_pairs: BTreeMap, + /// Corresponds to MutationSet's `new_root` field, and constructed from it. + /// + /// This is the Merkle tree's root hash that was calculated on the last + /// [`MutationSet::insert()`] call, and is applied to the Merkle tree with + /// [`SparseMerkleTree::set_root()`]. + new_root: RpoDigest, +} diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 6922a32..4f6c45d 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, CompletedMutationSet, EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, + MerkleError, MerklePath, MutationSet, NodeIndex, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD, + SMT_MAX_DEPTH, SMT_MIN_DEPTH, }; #[cfg(test)] @@ -188,6 +188,38 @@ impl SimpleSmt { >::insert(self, key, value) } + /// Start a prospective mutation transaction, which can be queried, discarded, or applied. + /// + /// This method returns a [`MutationSet`], on which you can call [`MutationSet::insert()`] 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 + /// [`MutationSet::done()`], and pass that to [`SimpleSmt::commit()`] 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 the prospective mutations began 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: CompletedMutationSet, 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`. ///