feat: implement root-checked Merkle tree insertion

This commit implements SparseMerkleTree::insert_ensure_root(), a version
of SparseMerkleTree::insert() that prevents modification of the tree if
the insert turns out to not result in the correct tree root (i.e., the
transaction is not valid).

This is an initial step towards issue 222.

As a logical next step, generalizing this to validation of inserting
multiple values at once will directly enable removing Merkle tree clones
in miden-node InnerState::apply_block().

For further work, as we generalize pre-validation for mutations we will
want a separate type to represent an arbitrary set of prospective
mutations on the Merkle tree and their validity, like the `ChangeSet`
type suggested in the issue.

src/merkle/smt/full/mod.rs

@@ -168,6 +168,24 @@ impl Smt {
         <Self as SparseMerkleTree<SMT_DEPTH>>::insert(self, key, value)
     }
 
+    /// Like [`Self::insert()`], but only performs the insert if the the new tree's root
+    /// hash would be equal to the hash given in `expected_root`.
+    ///
+    /// # Errors
+    /// Returns [`MerkleError::ConflictingRoots`] with a two-item [Vec] if the new root of the tree
+    /// is different from the expected root. The first item of the vector is the expected root,
+    /// and the second is actual root.
+    ///
+    /// No mutations are performed if the roots do no match.
+    pub fn insert_ensure_root(
+        &mut self,
+        key: RpoDigest,
+        value: Word,
+        expected_root: RpoDigest,
+    ) -> Result<Word, MerkleError> {
+        <Self as SparseMerkleTree<SMT_DEPTH>>::insert_ensure_root(self, key, value, expected_root)
+    }
+
     // HELPERS
     // --------------------------------------------------------------------------------------------
 

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::{smt::SparseMerkleTree, EmptySubtreeRoots, MerkleStore},
+    merkle::{smt::SparseMerkleTree, EmptySubtreeRoots, MerkleError, MerkleStore},
     utils::{Deserializable, Serializable},
     Word, ONE, WORD_SIZE,
 };
@@ -259,23 +259,29 @@ fn test_smt_removal() {
 }
 
 #[test]
-fn test_prospective_hash() {
+fn test_checked_insertion() {
+    use MerkleError::ConflictingRoots;
+
     let mut smt = Smt::default();
+    let smt_empty = smt.clone();
 
     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([3_u32.into(), 3_u32.into(), 3_u32.into(), Felt::new(raw)]);
+        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();
+
     // insert key-value 1
-    {
+    let root_1 = {
         let prospective = smt.hash_prospective_leaf(&key_1, &value_1);
         let old_value_1 = smt.insert(key_1, value_1);
         assert_eq!(old_value_1, EMPTY_WORD);
@@ -284,10 +290,29 @@ fn test_prospective_hash() {
 
         assert_eq!(smt.get_leaf(&key_1), SmtLeaf::Single((key_1, value_1)));
 
+        smt.root()
+    };
+
+    {
+        // Trying to insert something else into key_1 with the existing root should fail, and
+        // should not modify the tree at all.
+        let smt_before = smt.clone();
+        assert!(matches!(
+            smt.insert_ensure_root(key_1, value_2, root_1),
+            Err(ConflictingRoots(_)),
+        ));
+        assert_eq!(smt, smt_before);
+
+        // And inserting an empty word should bring us back to where we were.
+        assert_eq!(smt.insert_ensure_root(key_1, EMPTY_WORD, root_empty), Ok(value_1));
+        assert_eq!(smt, smt_empty);
+
+        smt.insert_ensure_root(key_1, value_1, root_1).unwrap();
+        assert_eq!(smt, smt_before);
     }
 
     // insert key-value 2
-    {
+    let root_2 = {
         let prospective = smt.hash_prospective_leaf(&key_2, &value_2);
         let old_value_2 = smt.insert(key_2, value_2);
         assert_eq!(old_value_2, EMPTY_WORD);
@@ -298,10 +323,25 @@ fn test_prospective_hash() {
             smt.get_leaf(&key_2),
             SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)])
         );
+
+        smt.root()
+    };
+
+    {
+        let smt_before = smt.clone();
+        assert!(matches!(
+            smt.insert_ensure_root(key_2, value_1, root_2),
+            Err(ConflictingRoots(_)),
+        ));
+        assert_eq!(smt, smt_before);
+
+        assert_eq!(smt.insert_ensure_root(key_2, EMPTY_WORD, root_1), Ok(value_2));
+        smt.insert_ensure_root(key_2, value_2, root_2).unwrap();
+        assert_eq!(smt, smt_before);
     }
 
     // insert key-value 3
-    {
+    let root_3 = {
         let prospective = smt.hash_prospective_leaf(&key_3, &value_3);
         let old_value_3 = smt.insert(key_3, value_3);
         assert_eq!(old_value_3, EMPTY_WORD);
@@ -310,14 +350,29 @@ fn test_prospective_hash() {
 
         assert_eq!(
             smt.get_leaf(&key_3),
-            SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2), (key_3, value_3)])
+            SmtLeaf::Multiple(vec![(key_3, value_3), (key_1, value_1), (key_2, value_2)])
         );
+
+        smt.root()
+    };
+
+    {
+        let smt_before = smt.clone();
+        assert!(matches!(
+            smt.insert_ensure_root(key_3, value_1, root_3),
+            Err(ConflictingRoots(_)),
+        ));
+        assert_eq!(smt, smt_before);
+
+        assert_eq!(smt.insert_ensure_root(key_3, EMPTY_WORD, root_2), Ok(value_3));
+        smt.insert_ensure_root(key_3, value_3, root_3).unwrap();
+        assert_eq!(smt, smt_before);
     }
 
     // remove key 3
     {
         let old_hash = smt.get_leaf(&key_3).hash();
-        let old_value_3 = smt.insert(key_3, EMPTY_WORD);
+        let old_value_3 = smt.insert_ensure_root(key_3, EMPTY_WORD, root_2).unwrap();
         assert_eq!(old_value_3, value_3);
         assert_eq!(old_hash, smt.hash_prospective_leaf(&key_3, &old_value_3));
 
@@ -325,26 +380,32 @@ fn test_prospective_hash() {
             smt.get_leaf(&key_3),
             SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)])
         );
+
+        assert_eq!(smt.root(), root_2);
     }
 
     // remove key 2
     {
         let old_hash = smt.get_leaf(&key_2).hash();
-        let old_value_2 = smt.insert(key_2, EMPTY_WORD);
+        let old_value_2 = smt.insert_ensure_root(key_2, EMPTY_WORD, root_1).unwrap();
         assert_eq!(old_value_2, value_2);
         assert_eq!(old_hash, smt.hash_prospective_leaf(&key_2, &old_value_2));
 
         assert_eq!(smt.get_leaf(&key_2), SmtLeaf::Single((key_1, value_1)));
+
+        assert_eq!(smt.root(), root_1);
     }
 
     // remove key 1
     {
         let old_hash = smt.get_leaf(&key_1).hash();
-        let old_value_1 = smt.insert(key_1, EMPTY_WORD);
+        let old_value_1 = smt.insert_ensure_root(key_1, EMPTY_WORD, root_empty).unwrap();
         assert_eq!(old_value_1, value_1);
         assert_eq!(old_hash, smt.hash_prospective_leaf(&key_1, &old_value_1));
 
         assert_eq!(smt.get_leaf(&key_1), SmtLeaf::new_empty(key_1.into()));
+
+        assert_eq!(smt.root(), root_empty);
     }
 }
 

src/merkle/smt/mod.rs

@@ -110,6 +110,104 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
         old_value
     }
 
+    /// Like [`Self::insert()`], but only performs the insert if the the new tree's root
+    /// hash would be equal to the hash given in `expected_root`.
+    ///
+    /// # Errors
+    /// Returns [`MerkleError::ConflictingRoots`] with a two-item [Vec] if the new root of the tree
+    /// is different from the expected root. The first item of the vector is the expected root,
+    /// and the second is actual root.
+    ///
+    /// No mutations are performed if the roots do no match.
+    fn insert_ensure_root(
+        &mut self,
+        key: Self::Key,
+        value: Self::Value,
+        expected_root: RpoDigest,
+    ) -> Result<Self::Value, MerkleError> {
+        let old_value = self.get_value(&key);
+        // if the old value and new value are the same, there is nothing to update
+        if value == old_value {
+            return Ok(value);
+        }
+
+        // Compute the nodes we'll need to make and remove.
+        let mut removals: Vec<NodeIndex> = Vec::with_capacity(DEPTH as usize);
+        let mut additions: Vec<(NodeIndex, InnerNode)> = Vec::with_capacity(DEPTH as usize);
+
+        let (mut node_index, mut parent_node) = {
+            let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(&key);
+            let node_index = NodeIndex::from(leaf_index);
+
+            let mut parent_index = node_index.clone();
+            parent_index.move_up();
+
+            (node_index, Some(self.get_inner_node(parent_index)))
+        };
+
+        let mut new_child_hash = self.hash_prospective_leaf(&key, &value);
+        for node_depth in (0..node_index.depth()).rev() {
+            let is_right = node_index.is_value_odd();
+            node_index.move_up();
+
+            let old_node = match parent_node.take() {
+                // On the first iteration, the 'old node' is the parent of the
+                // perspective leaf.
+                Some(parent_node) => parent_node,
+                // Otherwise it's a regular existing node.
+                None => self.get_inner_node(node_index),
+            };
+
+            //let new_node = new_node_from(is_right, old_node, new_child_hash);
+            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 node's new hash.
+            new_child_hash = new_node.hash();
+
+            let &equivalent_empty_hash = EmptySubtreeRoots::entry(DEPTH, node_depth);
+            if new_child_hash == equivalent_empty_hash {
+                // If a subtree is empty, we can remove the inner node, since it's equal to the
+                // default value.
+                removals.push(node_index);
+            } else {
+                additions.push((node_index, new_node));
+            }
+        }
+
+        // Once we're at depth 0, the last node we made is the new root.
+        let new_root = new_child_hash;
+
+        if expected_root != new_root {
+            return Err(MerkleError::ConflictingRoots(vec![expected_root, new_root]));
+        }
+
+        // Actual mutations start here.
+
+        self.insert_value(key, value);
+
+        for index in removals.drain(..) {
+            self.remove_inner_node(index);
+        }
+
+        for (index, new_node) in additions.drain(..) {
+            self.insert_inner_node(index, new_node);
+        }
+
+        self.set_root(new_root);
+
+        Ok(old_value)
+    }
+
     /// Recomputes the branch nodes (including the root) from `index` all the way to the root.
     /// `node_hash_at_index` is the hash of the node stored at index.
     fn recompute_nodes_from_index_to_root(

src/merkle/smt/simple/mod.rs

@@ -188,6 +188,25 @@ impl<const DEPTH: u8> SimpleSmt<DEPTH> {
         <Self as SparseMerkleTree<DEPTH>>::insert(self, key, value)
     }
 
+    /// Like [`Self::insert()`], but only performs the insert if the the new tree's root
+    /// hash would be equal to the hash given in `expected_root`.
+    ///
+    /// # Errors
+    /// Returns [`MerkleError::ConflictingRoots`] with a two-item [Vec] if the new root of the tree is
+    /// different from the expected root. The first item of the vector is the expected root, and the
+    /// second is actual root.
+    ///
+    /// No mutations are performed if the roots do no match.
+    pub fn insert_ensure_root(
+        &mut self,
+        key: LeafIndex<DEPTH>,
+        value: Word,
+        expected_root: RpoDigest,
+    ) -> Result<Word, MerkleError>
+    {
+        <Self as SparseMerkleTree<DEPTH>>::insert_ensure_root(self, key, value, expected_root)
+    }
+
     /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is
     /// computed as `DEPTH - SUBTREE_DEPTH`.
     ///