Implement `SimpleSmt::set_subtree` (#232)

* recompute_nodes_from_indeX_to_root

* MerkleError variant

* set_subtree

* test_simplesmt_set_subtree

* test_simplesmt_set_subtree_entire_tree

* test

* set_subtree: return root

src/merkle/error.rs

@@ -13,6 +13,7 @@ pub enum MerkleError {
     DuplicateValuesForKey(RpoDigest),
     InvalidIndex { depth: u8, value: u64 },
     InvalidDepth { expected: u8, provided: u8 },
+    InvalidSubtreeDepth { subtree_depth: u8, tree_depth: u8 },
     InvalidPath(MerklePath),
     InvalidNumEntries(usize),
     NodeNotInSet(NodeIndex),
@@ -36,6 +37,9 @@ impl fmt::Display for MerkleError {
             InvalidDepth { expected, provided } => {
                 write!(f, "the provided depth {provided} is not valid for {expected}")
             }
+            InvalidSubtreeDepth { subtree_depth, tree_depth } => {
+                write!(f, "tried inserting a subtree of depth {subtree_depth} into a tree of depth {tree_depth}")
+            }
             InvalidPath(_path) => write!(f, "the provided path is not valid"),
             InvalidNumEntries(max) => write!(f, "number of entries exceeded the maximum: {max}"),
             NodeNotInSet(index) => write!(f, "the node with index ({index}) is not in the set"),

src/merkle/simple_smt/mod.rs

@@ -229,7 +229,7 @@ impl SimpleSmt {
     /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
     pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<Word, MerkleError> {
         // validate the index before modifying the structure
-        let mut idx = NodeIndex::new(self.depth(), index)?;
+        let idx = NodeIndex::new(self.depth(), index)?;
 
         let old_value = self.insert_leaf_node(index, value).unwrap_or(Self::EMPTY_VALUE);
 
@@ -238,22 +238,93 @@ impl SimpleSmt {
             return Ok(value);
         }
 
-        let mut value = RpoDigest::from(value);
-        for _ in 0..idx.depth() {
-            let is_right = idx.is_value_odd();
-            idx.move_up();
-            let BranchNode { left, right } = self.get_branch_node(&idx);
-            let (left, right) = if is_right { (left, value) } else { (value, right) };
-            self.insert_branch_node(idx, left, right);
-            value = Rpo256::merge(&[left, right]);
-        }
-        self.root = value;
+        self.recompute_nodes_from_index_to_root(idx, RpoDigest::from(value));
+
         Ok(old_value)
     }
 
+    /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is
+    /// computed as `self.depth() - subtree.depth()`.
+    ///
+    /// Returns the new root.
+    pub fn set_subtree(
+        &mut self,
+        subtree_insertion_index: u64,
+        subtree: SimpleSmt,
+    ) -> Result<RpoDigest, MerkleError> {
+        if subtree.depth() > self.depth() {
+            return Err(MerkleError::InvalidSubtreeDepth {
+                subtree_depth: subtree.depth(),
+                tree_depth: self.depth(),
+            });
+        }
+
+        // Verify that `subtree_insertion_index` is valid.
+        let subtree_root_insertion_depth = self.depth() - subtree.depth();
+        let subtree_root_index =
+            NodeIndex::new(subtree_root_insertion_depth, subtree_insertion_index)?;
+
+        // add leaves
+        // --------------
+
+        // The subtree's leaf indices live in their own context - i.e. a subtree of depth `d`. If we
+        // insert the subtree at `subtree_insertion_index = 0`, then the subtree leaf indices are
+        // valid as they are. However, consider what happens when we insert at
+        // `subtree_insertion_index = 1`. The first leaf of our subtree now will have index `2^d`;
+        // you can see it as there's a full subtree sitting on its left. In general, for
+        // `subtree_insertion_index = i`, there are `i` subtrees sitting before the subtree we want
+        // to insert, so we need to adjust all its leaves by `i * 2^d`.
+        let leaf_index_shift: u64 = subtree_insertion_index * 2_u64.pow(subtree.depth().into());
+        for (subtree_leaf_idx, leaf_value) in subtree.leaves() {
+            let new_leaf_idx = leaf_index_shift + subtree_leaf_idx;
+            debug_assert!(new_leaf_idx < 2_u64.pow(self.depth().into()));
+
+            self.insert_leaf_node(new_leaf_idx, *leaf_value);
+        }
+
+        // add subtree's branch nodes (which includes the root)
+        // --------------
+        for (branch_idx, branch_node) in subtree.branches {
+            let new_branch_idx = {
+                let new_depth = subtree_root_insertion_depth + branch_idx.depth();
+                let new_value = subtree_insertion_index * 2_u64.pow(branch_idx.depth().into())
+                    + branch_idx.value();
+
+                NodeIndex::new(new_depth, new_value).expect("index guaranteed to be valid")
+            };
+
+            self.branches.insert(new_branch_idx, branch_node);
+        }
+
+        // recompute nodes starting from subtree root
+        // --------------
+        self.recompute_nodes_from_index_to_root(subtree_root_index, subtree.root);
+
+        Ok(self.root)
+    }
+
     // HELPER METHODS
     // --------------------------------------------------------------------------------------------
 
+    /// 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(
+        &mut self,
+        mut index: NodeIndex,
+        node_hash_at_index: RpoDigest,
+    ) {
+        let mut value = node_hash_at_index;
+        for _ in 0..index.depth() {
+            let is_right = index.is_value_odd();
+            index.move_up();
+            let BranchNode { left, right } = self.get_branch_node(&index);
+            let (left, right) = if is_right { (left, value) } else { (value, right) };
+            self.insert_branch_node(index, left, right);
+            value = Rpo256::merge(&[left, right]);
+        }
+        self.root = value;
+    }
+
     fn get_leaf_node(&self, key: u64) -> Option<Word> {
         self.leaves.get(&key).copied()
     }

src/merkle/simple_smt/tests.rs

@@ -3,7 +3,7 @@ use super::{
     NodeIndex, Rpo256, Vec,
 };
 use crate::{
-    merkle::{digests_to_words, int_to_leaf, int_to_node},
+    merkle::{digests_to_words, int_to_leaf, int_to_node, EmptySubtreeRoots},
     Word,
 };
 
@@ -349,6 +349,141 @@ fn test_simplesmt_with_leaves_nonexisting_leaf() {
     assert!(result.is_err());
 }
 
+#[test]
+fn test_simplesmt_set_subtree() {
+    // Final Tree:
+    //
+    //        ____k____
+    //       /         \
+    //     _i_         _j_
+    //    /   \       /   \
+    //   e     f     g     h
+    //  / \   / \   / \   / \
+    // a   b 0   0 c   0 0   d
+
+    let z = EMPTY_WORD;
+
+    let a = Word::from(Rpo256::merge(&[z.into(); 2]));
+    let b = Word::from(Rpo256::merge(&[a.into(); 2]));
+    let c = Word::from(Rpo256::merge(&[b.into(); 2]));
+    let d = Word::from(Rpo256::merge(&[c.into(); 2]));
+
+    let e = Rpo256::merge(&[a.into(), b.into()]);
+    let f = Rpo256::merge(&[z.into(), z.into()]);
+    let g = Rpo256::merge(&[c.into(), z.into()]);
+    let h = Rpo256::merge(&[z.into(), d.into()]);
+
+    let i = Rpo256::merge(&[e, f]);
+    let j = Rpo256::merge(&[g, h]);
+
+    let k = Rpo256::merge(&[i, j]);
+
+    // subtree:
+    //   g
+    //  / \
+    // c   0
+    let subtree = {
+        let depth = 1;
+        let entries = vec![(0, c)];
+        SimpleSmt::with_leaves(depth, entries).unwrap()
+    };
+
+    // insert subtree
+    let tree = {
+        let depth = 3;
+        let entries = vec![(0, a), (1, b), (7, d)];
+        let mut tree = SimpleSmt::with_leaves(depth, entries).unwrap();
+
+        tree.set_subtree(2, subtree).unwrap();
+
+        tree
+    };
+
+    assert_eq!(tree.root(), k);
+    assert_eq!(tree.get_leaf(4).unwrap(), c);
+    assert_eq!(tree.get_branch_node(&NodeIndex::new_unchecked(2, 2)).parent(), g);
+}
+
+/// Ensures that an invalid input node index into `set_subtree()` incurs no mutation of the tree
+#[test]
+fn test_simplesmt_set_subtree_unchanged_for_wrong_index() {
+    // Final Tree:
+    //
+    //        ____k____
+    //       /         \
+    //     _i_         _j_
+    //    /   \       /   \
+    //   e     f     g     h
+    //  / \   / \   / \   / \
+    // a   b 0   0 c   0 0   d
+
+    let z = EMPTY_WORD;
+
+    let a = Word::from(Rpo256::merge(&[z.into(); 2]));
+    let b = Word::from(Rpo256::merge(&[a.into(); 2]));
+    let c = Word::from(Rpo256::merge(&[b.into(); 2]));
+    let d = Word::from(Rpo256::merge(&[c.into(); 2]));
+
+    // subtree:
+    //   g
+    //  / \
+    // c   0
+    let subtree = {
+        let depth = 1;
+        let entries = vec![(0, c)];
+        SimpleSmt::with_leaves(depth, entries).unwrap()
+    };
+
+    let mut tree = {
+        let depth = 3;
+        let entries = vec![(0, a), (1, b), (7, d)];
+        SimpleSmt::with_leaves(depth, entries).unwrap()
+    };
+    let tree_root_before_insertion = tree.root();
+
+    // insert subtree
+    assert!(tree.set_subtree(500, subtree).is_err());
+
+    assert_eq!(tree.root(), tree_root_before_insertion);
+}
+
+/// We insert an empty subtree that has the same depth as the original tree
+#[test]
+fn test_simplesmt_set_subtree_entire_tree() {
+    // Initial Tree:
+    //
+    //        ____k____
+    //       /         \
+    //     _i_         _j_
+    //    /   \       /   \
+    //   e     f     g     h
+    //  / \   / \   / \   / \
+    // a   b 0   0 c   0 0   d
+
+    let z = EMPTY_WORD;
+
+    let a = Word::from(Rpo256::merge(&[z.into(); 2]));
+    let b = Word::from(Rpo256::merge(&[a.into(); 2]));
+    let c = Word::from(Rpo256::merge(&[b.into(); 2]));
+    let d = Word::from(Rpo256::merge(&[c.into(); 2]));
+
+    let depth = 3;
+
+    // subtree: E3
+    let subtree = { SimpleSmt::with_leaves(depth, Vec::new()).unwrap() };
+    assert_eq!(subtree.root(), *EmptySubtreeRoots::entry(depth, 0));
+
+    // insert subtree
+    let mut tree = {
+        let entries = vec![(0, a), (1, b), (4, c), (7, d)];
+        SimpleSmt::with_leaves(depth, entries).unwrap()
+    };
+
+    tree.set_subtree(0, subtree).unwrap();
+
+    assert_eq!(tree.root(), *EmptySubtreeRoots::entry(depth, 0));
+}
+
 // HELPER FUNCTIONS
 // --------------------------------------------------------------------------------------------