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

@@ -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

@@ -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 {
         <Self as SparseMerkleTree<SMT_DEPTH>>::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<SMT_DEPTH, Self> {
+        <Self as SparseMerkleTree<SMT_DEPTH>>::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<SMT_DEPTH, RpoDigest, Word>) {
+        <Self as SparseMerkleTree<SMT_DEPTH>>::commit(self, mutations)
+    }
+
     // HELPERS
     // --------------------------------------------------------------------------------------------
 

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() {

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<const DEPTH: u8> {
+pub 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,64 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
         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<DEPTH, Self>
+    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<DEPTH, Self::Key, Self::Value>)
+    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<const DEPTH: u8> {
     /// 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<const DEPTH: u8> {
 
 #[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<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 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<DEPTH>,
+{
+    /// 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<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.
+    new_pairs: BTreeMap<T::Key, T::Value>,
+    /// 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<DEPTH>,
+{
+    /// Pass the return value of this method to [`SparseMerkleTree::commit()`] to apply these
+    /// mutations to the Merkle tree.
+    pub fn done(self) -> DoneMutationSet<DEPTH, T::Key, T::Value> {
+        // 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<const DEPTH: u8, K, V> {
+    node_mutations: BTreeMap<NodeIndex, NodeMutation>,
+    new_pairs: BTreeMap<K, V>,
+    new_root: RpoDigest,
+}

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<const DEPTH: u8> SimpleSmt<DEPTH> {
         <Self as SparseMerkleTree<DEPTH>>::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<DEPTH, Self> {
+        <Self as SparseMerkleTree<DEPTH>>::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<DEPTH, LeafIndex<DEPTH>, Word>) {
+        <Self as SparseMerkleTree<DEPTH>>::commit(self, mutations)
+    }
+
     /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is
     /// computed as `DEPTH - SUBTREE_DEPTH`.
     ///