WIP: smt: implement root-checked insertion

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,65 +259,119 @@ 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);
 
-        assert_eq!(smt.get_leaf(&key_1).hash(), prospective);
-
+        assert_eq!(prospective, smt.get_leaf(&key_1).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);
 
-        assert_eq!(smt.get_leaf(&key_2).hash(), prospective);
+        assert_eq!(prospective, smt.get_leaf(&key_2).hash());
 
         assert_eq!(
             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 prospective_hash = smt.hash_prospective_leaf(&key_3, &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);
 
-        assert_eq!(smt.get_leaf(&key_3).hash(), prospective_hash);
+        assert_eq!(prospective, smt.get_leaf(&key_3).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 +379,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`.
     ///