feat(smt): implement generic prospective insertions

This commit adds two methods to SparseMerkleTree: compute_mutations()
and apply_mutations(), which respectively create and consume a new
MutationSet type. This type represents as 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.

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::compute_mutations()` and `Smt::apply_mutations()` for validation-checked insertions (#327).
 
 ## 0.10.0 (2024-08-06)
 

src/main.rs

@@ -35,6 +35,7 @@ pub fn benchmark_smt() {
 
     let mut tree = construction(entries, tree_size).unwrap();
     insertion(&mut tree, tree_size).unwrap();
+    batched_insertion(&mut tree, tree_size).unwrap();
     proof_generation(&mut tree, tree_size).unwrap();
 }
 
@@ -82,6 +83,34 @@ pub fn insertion(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
     Ok(())
 }
 
+pub fn batched_insertion(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
+    println!("Running a batched insertion benchmark:");
+
+    let new_pairs: Vec<(RpoDigest, Word)> = (0..20)
+        .map(|i| {
+            let key = Rpo256::hash(&rand_value::<u64>().to_be_bytes());
+            let value = [ONE, ONE, ONE, Felt::new(size + i)];
+            (key, value)
+        })
+        .collect();
+
+    let now = Instant::now();
+
+    let mutations = tree.compute_mutations(new_pairs);
+    tree.apply_mutations(mutations).unwrap();
+
+    let elapsed = now.elapsed();
+
+    println!(
+        "An average insertion time measured by a 20-batch insert into an SMT with {} key-value pairs over {:.3} milliseconds is {:.3} milliseconds",
+        size,
+        elapsed.as_secs_f32() * 1000f32,
+        elapsed.as_secs_f32() * 50f32,
+    );
+
+    Ok(())
+}
+
 /// Runs the proof generation benchmark for the [`Smt`].
 pub fn proof_generation(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
     println!("Running a proof generation benchmark:");

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;

src/merkle/smt/full/mod.rs

@@ -6,7 +6,7 @@ use alloc::{
 
 use super::{
     EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, MerklePath,
-    NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD,
+    MutationSet, NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD,
 };
 
 mod error;
@@ -167,6 +167,47 @@ impl Smt {
         <Self as SparseMerkleTree<SMT_DEPTH>>::insert(self, key, value)
     }
 
+    /// Compute what changes are necessary to insert the specified key-value pairs into this Merkle
+    /// tree, allowing for validation before applying those changes.
+    ///
+    /// This method returns a [`MutationSet`], which contains all the information for inserting
+    /// `kv_pairs` into this Merkle tree already calculated, including the new root hash, which can
+    /// be queried with [`MutationSet::root()`]. When you are done, call [`Smt::apply_mutations()`]
+    /// to commit these changes to the Merkle tree, dor [`drop()`] the [`MutationSet`] to discard
+    /// them.
+    ///
+    /// # Example
+    /// ```
+    /// # use miden_crypto::{hash::rpo::RpoDigest, Felt, Word};
+    /// # use miden_crypto::merkle::{Smt, EmptySubtreeRoots, SMT_DEPTH};
+    /// let mut smt = Smt::new();
+    /// let pair = (RpoDigest::default(), Word::default());
+    /// let mutations = smt.compute_mutations(vec![pair]);
+    /// assert_eq!(mutations.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
+    /// smt.apply_mutations(mutations);
+    /// assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
+    /// ```
+    pub fn compute_mutations(
+        &self,
+        kv_pairs: impl IntoIterator<Item = (RpoDigest, Word)>,
+    ) -> MutationSet<SMT_DEPTH, RpoDigest, Word> {
+        <Self as SparseMerkleTree<SMT_DEPTH>>::compute_mutations(self, kv_pairs)
+    }
+
+    /// Apply the prospective mutations began with [`Smt::compute_mutations()`] to this tree.
+    ///
+    /// # Errors
+    /// If `mutations` was computed on a tree with a different root than this one, returns
+    /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
+    /// the `mutations` were computed against, and the second item is the actual current root of
+    /// this tree.
+    pub fn apply_mutations(
+        &mut self,
+        mutations: MutationSet<SMT_DEPTH, RpoDigest, Word>,
+    ) -> Result<(), MerkleError> {
+        <Self as SparseMerkleTree<SMT_DEPTH>>::apply_mutations(self, mutations)
+    }
+
     // HELPERS
     // --------------------------------------------------------------------------------------------
 

src/merkle/smt/full/tests.rs

@@ -371,6 +371,82 @@ 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 mutations = smt.compute_mutations(vec![(key_1, value_1)]);
+    assert_eq!(mutations.root(), root_1, "prospective root 1 did not match actual root 1");
+    smt.apply_mutations(mutations).unwrap();
+    assert_eq!(smt.root(), root_1, "mutations before and after apply did not match");
+
+    let mutations = smt.compute_mutations(vec![(key_2, value_2)]);
+    assert_eq!(mutations.root(), root_2, "prospective root 2 did not match actual root 2");
+    let mutations =
+        smt.compute_mutations(vec![(key_3, EMPTY_WORD), (key_2, value_2), (key_3, value_3)]);
+    assert_eq!(mutations.root(), root_3, "mutations before and after apply did not match");
+    smt.apply_mutations(mutations).unwrap();
+
+    // Edge case: multiple values at the same key, where a later pair restores the original value.
+    let mutations = smt.compute_mutations(vec![(key_3, EMPTY_WORD), (key_3, value_3)]);
+    assert_eq!(mutations.root(), root_3);
+    smt.apply_mutations(mutations).unwrap();
+    assert_eq!(smt.root(), root_3);
+
+    // Test batch updates, and that the order doesn't matter.
+    let pairs =
+        vec![(key_3, value_2), (key_2, EMPTY_WORD), (key_1, EMPTY_WORD), (key_3, EMPTY_WORD)];
+    let mutations = smt.compute_mutations(pairs);
+    assert_eq!(
+        mutations.root(),
+        root_empty,
+        "prospective root for batch removal did not match actual root",
+    );
+    smt.apply_mutations(mutations).unwrap();
+    assert_eq!(smt.root(), root_empty, "mutations before and after apply did not match");
+
+    let pairs = vec![(key_3, value_3), (key_1, value_1), (key_2, value_2)];
+    let mutations = smt.compute_mutations(pairs);
+    assert_eq!(mutations.root(), root_3);
+    smt.apply_mutations(mutations).unwrap();
+    assert_eq!(smt.root(), root_3);
+}
+
 /// 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() {

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::{
@@ -45,11 +45,11 @@ pub const SMT_MAX_DEPTH: u8 = 64;
 /// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs.
 pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     /// 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,152 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
         self.set_root(node_hash);
     }
 
+    /// Compute what changes are necessary to insert the specified key-value pairs into this Merkle
+    /// tree, allowing for validation before applying those changes.
+    ///
+    /// This method returns a [`MutationSet`], which contains all the information for inserting
+    /// `kv_pairs` into this Merkle tree already calculated, including the new root hash, which can
+    /// be queried with [`MutationSet::root()`]. When you are done, call
+    /// [`SparseMerkleTree::apply_mutations()`] to commit these changes to the Merkle tree, or
+    /// [`drop()`] the [`MutationSet`] to discard them.
+    fn compute_mutations(
+        &self,
+        kv_pairs: impl IntoIterator<Item = (Self::Key, Self::Value)>,
+    ) -> MutationSet<DEPTH, Self::Key, Self::Value> {
+        use NodeMutation::*;
+
+        let mut new_root = self.root();
+        let mut new_pairs: BTreeMap<Self::Key, Self::Value> = Default::default();
+        let mut node_mutations: BTreeMap<NodeIndex, NodeMutation> = Default::default();
+
+        for (key, value) in kv_pairs {
+            // If the old value and the new value are the same, there is nothing to update.
+            // For the unusual case that kv_pairs has multiple values at the same key, we'll have
+            // to check the key-value pairs we've already seen to get the "effective" old value.
+            let old_value = new_pairs.get(&key).cloned().unwrap_or_else(|| self.get_value(&key));
+            if value == old_value {
+                continue;
+            }
+
+            let leaf_index = Self::key_to_leaf_index(&key);
+            let mut node_index = NodeIndex::from(leaf_index);
+
+            // We need the current leaf's hash to calculate the new leaf, but in the rare case that
+            // `kv_pairs` has multiple pairs that go into the same leaf, then those pairs are also
+            // part of the "current leaf".
+            let old_leaf = {
+                let pairs_at_index = new_pairs
+                    .iter()
+                    .filter(|&(new_key, _)| Self::key_to_leaf_index(new_key) == leaf_index);
+
+                pairs_at_index.fold(self.get_leaf(&key), |acc, (k, v)| {
+                    // Most of the time `pairs_at_index` should only contain a single entry (or
+                    // none at all), as multi-leaves should be really rare.
+                    let existing_leaf = acc.clone();
+                    self.construct_prospective_leaf(existing_leaf, k, v)
+                })
+            };
+
+            let new_leaf = self.construct_prospective_leaf(old_leaf, &key, &value);
+
+            let mut new_child_hash = Self::hash_leaf(&new_leaf);
+
+            for node_depth in (0..node_index.depth()).rev() {
+                // Whether the node we're replacing is the right child or the left child.
+                let is_right = node_index.is_value_odd();
+                node_index.move_up();
+
+                let old_node = node_mutations
+                    .get(&node_index)
+                    .map(|mutation| match mutation {
+                        Addition(node) => node.clone(),
+                        Removal => {
+                            let &child = EmptySubtreeRoots::entry(DEPTH, node_depth + 1);
+                            InnerNode { left: child, right: child }
+                        },
+                    })
+                    .unwrap_or_else(|| self.get_inner_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);
+                let is_removal = new_child_hash == equivalent_empty_hash;
+                let new_entry = if is_removal { Removal } else { Addition(new_node) };
+                node_mutations.insert(node_index, new_entry);
+            }
+
+            // Once we're at depth 0, the last node we made is the new root.
+            new_root = new_child_hash;
+            // And then we're done with this pair; on to the next one.
+            new_pairs.insert(key, value);
+        }
+
+        MutationSet {
+            old_root: self.root(),
+            new_root,
+            node_mutations,
+            new_pairs,
+        }
+    }
+
+    /// Apply the prospective mutations began with [`SparseMerkleTree::compute_mutations()`] to
+    /// this tree.
+    ///
+    /// # Errors
+    /// If `mutations` was computed on a tree with a different root than this one, returns
+    /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
+    /// the `mutations` were computed against, and the second item is the actual current root of
+    /// this tree.
+    fn apply_mutations(
+        &mut self,
+        mutations: MutationSet<DEPTH, Self::Key, Self::Value>,
+    ) -> Result<(), MerkleError>
+    where
+        Self: Sized,
+    {
+        use NodeMutation::*;
+        let MutationSet {
+            old_root,
+            node_mutations,
+            new_pairs,
+            new_root,
+        } = mutations;
+
+        // Guard against accidentally trying to apply mutations that were computed against a
+        // different tree, including a stale version of this tree.
+        if old_root != self.root() {
+            return Err(MerkleError::ConflictingRoots(vec![old_root, self.root()]));
+        }
+
+        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);
+
+        Ok(())
+    }
+
     // REQUIRED METHODS
     // ---------------------------------------------------------------------------------------------
 
@@ -266,3 +412,51 @@ impl<const DEPTH: u8> TryFrom<NodeIndex> for LeafIndex<DEPTH> {
         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 {
+    /// Corresponds to [`SparseMerkleTree::remove_inner_node()`].
+    Removal,
+    /// Corresponds to [`SparseMerkleTree::insert_inner_node()`].
+    Addition(InnerNode),
+}
+
+/// Represents a group of prospective mutations to a [`SparseMerkleTree`], created by
+/// [`SparseMerkleTree::compute_mutations()`], and that can be applied with
+/// [`SparseMerkleTree::apply_mutations()`].
+#[derive(Debug, Clone, PartialEq, Eq, Default)]
+pub struct MutationSet<const DEPTH: u8, K, V> {
+    /// The root of the Merkle tree this MutationSet is for, recorded at the time
+    /// [`SparseMerkleTree::compute_mutations()`] was called. Exists to guard against applying
+    /// mutations to the wrong tree or applying stale mutations to a tree that has since changed.
+    old_root: RpoDigest,
+    /// 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. Each [`NodeMutation::Addition`] corresponds to a
+    /// [`SparseMerkleTree::insert_inner_node()`] call, and each [`NodeMutation::Removal`] corresponds
+    /// to a [`SparseMerkleTree::remove_inner_node()`] call.
+    node_mutations: BTreeMap<NodeIndex, NodeMutation>,
+    /// 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. Each entry corresponds to a
+    /// [`SparseMerkleTree::insert_value()`] call.
+    new_pairs: BTreeMap<K, V>,
+    /// The calculated root for the Merkle tree, given these mutations. Set on computation, and
+    /// publcily retrievable with with [`MutationSet::root()`]. Corresponds to a
+    /// [`SparseMerkleTree::set_root()`] call.
+    new_root: RpoDigest,
+}
+
+impl<const DEPTH: u8, K, V> MutationSet<DEPTH, K, V> {
+    /// Queries the root that was calculated during [`SparseMerkleTree::compute_mutations()`]. See
+    /// that method for more information.
+    pub fn root(&self) -> RpoDigest {
+        self.new_root
+    }
+}

src/merkle/smt/simple/mod.rs

@@ -2,8 +2,8 @@ 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,
+    MerklePath, MutationSet, NodeIndex, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD,
+    SMT_MAX_DEPTH, SMT_MIN_DEPTH,
 };
 
 #[cfg(test)]
@@ -188,6 +188,47 @@ impl<const DEPTH: u8> SimpleSmt<DEPTH> {
         <Self as SparseMerkleTree<DEPTH>>::insert(self, key, value)
     }
 
+    /// Compute what changes are necessary to insert the specified key-value pairs into this Merkle
+    /// tree, allowing for validation before applying those changes.
+    ///
+    /// This method returns a [`MutationSet`], which contains all the information for inserting
+    /// `kv_pairs` into this Merkle tree already calculated, including the new root hash, which can
+    /// be queried with [`MutationSet::root()`]. When you are done, call
+    /// [`SimpleSmt::apply_mutations()`] to commit these changes to the Merkle tree, or [`drop()`]
+    /// the [`MutationSet`] to discard them.
+    ///
+    /// # 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 pair = (LeafIndex::default(), Word::default());
+    /// let mutations = smt.compute_mutations(vec![pair]);
+    /// assert_eq!(mutations.root(), *EmptySubtreeRoots::entry(3, 0));
+    /// smt.apply_mutations(mutations);
+    /// assert_eq!(smt.root(), *EmptySubtreeRoots::entry(3, 0));
+    /// ```
+    pub fn compute_mutations(
+        &self,
+        kv_pairs: impl IntoIterator<Item = (LeafIndex<DEPTH>, Word)>,
+    ) -> MutationSet<DEPTH, LeafIndex<DEPTH>, Word> {
+        <Self as SparseMerkleTree<DEPTH>>::compute_mutations(self, kv_pairs)
+    }
+
+    /// Apply the prospective mutations began with [`SimpleSmt::compute_mutations()`] to this tree.
+    ///
+    /// # Errors
+    /// If `mutations` was computed on a tree with a different root than this one, returns
+    /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
+    /// the `mutations` were computed against, and the second item is the actual current root of
+    /// this tree.
+    pub fn apply_mutations(
+        &mut self,
+        mutations: MutationSet<DEPTH, LeafIndex<DEPTH>, Word>,
+    ) -> Result<(), MerkleError> {
+        <Self as SparseMerkleTree<DEPTH>>::apply_mutations(self, mutations)
+    }
+
     /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is
     /// computed as `DEPTH - SUBTREE_DEPTH`.
     ///