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

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

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;

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

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