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 74810c3..a1e4d93 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 70b7944..56cd3ff 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -365,6 +365,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 dce9ef4..24967ed 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 // --------------------------------------------------------------------------------------------- @@ -178,7 +235,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 +256,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 +319,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.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 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 c0b2ae1..5fe02e7 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`. ///