CMT: Subtree Append

lib/concurrent-merkle-tree/src/error.rs

@@ -35,4 +35,12 @@ pub enum CMTError {
         "Valid proof was passed to a leaf, but it's value has changed since the proof was issued"
     )]
     LeafContentsModified,
+
+    /// Attempted to append subtree to initialize a tree
+    #[error("Cannot initialize tree with subtree append")]
+    CannotInitializeWithSubtreeAppend,
+
+    /// Attempted to append subtree of invalid size for the current state of the tree
+    #[error("Cannot append subtree with invalid size")]
+    SubtreeInvalidSize,
 }

lib/concurrent-merkle-tree/src/merkle_roll.rs

@@ -1,7 +1,7 @@
 use crate::{
     error::CMTError,
     state::{ChangeLog, Node, Path, EMPTY},
-    utils::{empty_node, empty_node_cached, fill_in_proof, recompute, hash_to_parent},
+    utils::{empty_node, empty_node_cached, fill_in_proof, hash_to_parent, recompute},
 };
 use bytemuck::{Pod, Zeroable};
 pub(crate) use log_compute;
@@ -144,7 +144,7 @@ impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize> MerkleRoll<MAX_DEPTH,
         }
     }
 
-    /// Basic operation that always succeeds
+    /// Append leaf to tree
     pub fn append(&mut self, mut node: Node) -> Result<Node, CMTError> {
         check_bounds(MAX_DEPTH, MAX_BUFFER_SIZE);
         if node == EMPTY {
@@ -196,6 +196,91 @@ impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize> MerkleRoll<MAX_DEPTH,
         Ok(node)
     }
 
+    /// Append subtree to current tree
+    pub fn append_subtree(
+        &mut self,
+        subtree_root: Node,
+        subtree_rightmost_leaf: Node,
+        subtree_rightmost_index: u32,
+        subtree_rightmost_proof: Vec<Node>,
+    ) -> Result<Node, CMTError> {
+        check_bounds(MAX_DEPTH, MAX_BUFFER_SIZE);
+        // TODO: consider adding check that the subtree is not empty (but will omit for now)
+        //       note: it will be more time consuming to check that the subtree_root is not a root to an empty tree, because it requires applying O(depth) hashes of trivial nodes
+        if self.rightmost_proof.index >= 1 << MAX_DEPTH {
+            return Err(CMTError::TreeFull);
+        }
+
+        if self.rightmost_proof.index == 0 {
+            // TODO(sorend): Consider if we want to be able to initialize a tree with a subtree append (I think it would be cool, but is a little complex). write an equivallent function for subtree_append: initialize_tree_from_subtree_append
+            // For now
+            return Err(CMTError::CannotInitializeWithSubtreeAppend);
+        }
+
+        // Confirm that subtree_rightmost_proof is valid
+        if recompute(
+            subtree_rightmost_leaf,
+            &subtree_rightmost_proof[..],
+            subtree_rightmost_index - 1,
+        ) != subtree_root
+        {
+            return Err(CMTError::InvalidProof);
+        }
+
+        let leaf = subtree_rightmost_leaf.clone();
+        let intersection = self.rightmost_proof.index.trailing_zeros() as usize;
+
+        // @dev: At any given time (other than initialization), there is only one valid size of subtree that can be appended
+        if subtree_rightmost_proof.len() != intersection {
+            return Err(CMTError::SubtreeInvalidSize);
+        }
+
+        let mut change_list = [EMPTY; MAX_DEPTH];
+        let mut intersection_node = self.rightmost_proof.leaf;
+
+        // This will be mutated into the new root after the append by gradually hashing this node with the RMP to the subtree, then the critical node, and then the rest of the RMP to this tree.
+        let mut node = subtree_rightmost_leaf;
+
+        for i in 0..MAX_DEPTH {
+            change_list[i] = node;
+            if i < intersection {
+                // Compute proof to the appended node from empty nodes
+                hash_to_parent(
+                    &mut intersection_node,
+                    &self.rightmost_proof.proof[i],
+                    ((self.rightmost_proof.index - 1) >> i) & 1 == 0,
+                );
+                hash_to_parent(
+                    &mut node,
+                    &subtree_rightmost_proof[i],
+                    ((subtree_rightmost_index - 1) >> i) & 1 == 0,
+                );
+                self.rightmost_proof.proof[i] = subtree_rightmost_proof[i];
+            } else if i == intersection {
+                // Compute where the new node intersects the main tree
+                hash_to_parent(&mut node, &intersection_node, false);
+                self.rightmost_proof.proof[intersection] = intersection_node;
+            } else {
+                // Update the change list path up to the root
+                hash_to_parent(
+                    &mut node,
+                    &self.rightmost_proof.proof[i],
+                    ((self.rightmost_proof.index - 1) >> i) & 1 == 0,
+                );
+            }
+        }
+
+        self.update_internal_counters();
+        self.change_logs[self.active_index as usize] = ChangeLog::<MAX_DEPTH>::new(
+            node,
+            change_list,
+            self.rightmost_proof.index + subtree_rightmost_index - 1,
+        );
+        self.rightmost_proof.index = self.rightmost_proof.index + subtree_rightmost_index;
+        self.rightmost_proof.leaf = leaf;
+        Ok(node)
+    }
+
     /// Convenience function for `set_leaf`
     /// On write conflict:
     /// Will append

lib/concurrent-merkle-tree/src/state.rs

@@ -37,7 +37,12 @@ impl<const MAX_DEPTH: usize> ChangeLog<MAX_DEPTH> {
     }
 
     /// Sets all change log values from a leaf and valid proof
-    pub fn replace_and_recompute_path(&mut self, index: u32, mut node: Node, proof: &[Node]) -> Node {
+    pub fn replace_and_recompute_path(
+        &mut self,
+        index: u32,
+        mut node: Node,
+        proof: &[Node],
+    ) -> Node {
         self.index = index;
         for (i, sibling) in proof.iter().enumerate() {
             self.path[i] = node;

lib/concurrent-merkle-tree/tests/tests.rs

@@ -61,6 +61,121 @@ async fn test_append() {
     );
 }
 
+#[tokio::test(threaded_scheduler)]
+async fn test_append_complete_subtree() {
+    let (mut merkle_roll, mut off_chain_tree) = setup();
+    let mut rng = thread_rng();
+    merkle_roll.initialize().unwrap();
+
+    // insert eight leaves into the large tree
+    for i in 0..8 {
+        let leaf = rng.gen::<[u8; 32]>();
+        merkle_roll.append(leaf).unwrap();
+        off_chain_tree.add_leaf(leaf, i);
+        assert_eq!(
+            merkle_roll.get_change_log().root,
+            off_chain_tree.get_root(),
+            "On chain tree failed to update properly on an append",
+        );
+    }
+
+    // create a simple subtree to append of depth three
+    let mut onchain_subtree = MerkleRoll::<3, 8>::new();
+    onchain_subtree.initialize().unwrap();
+
+    // completely fill the subtree with unique leaves, and also append them to the off-chain tree
+    for i in 8..16 {
+        let leaf = rng.gen::<[u8; 32]>();
+        onchain_subtree.append(leaf).unwrap();
+        off_chain_tree.add_leaf(leaf, i);
+    }
+
+    // append the on_chain subtree to the merkle_roll
+    merkle_roll
+        .append_subtree(
+            onchain_subtree.get_change_log().root,
+            onchain_subtree.rightmost_proof.leaf,
+            onchain_subtree.rightmost_proof.index,
+            onchain_subtree.rightmost_proof.proof.to_vec(),
+        )
+        .unwrap();
+
+    // The result should be that the merkle_roll's new root is the same as the root of the off-chain tree which had leaves 0..15 appended one by one
+    assert_eq!(
+        merkle_roll.get_change_log().root,
+        off_chain_tree.get_root(),
+        "On chain tree failed to update properly on an append",
+    );
+
+    // Show that we can still append to the large tree after performing a subtree append
+    let leaf = rng.gen::<[u8; 32]>();
+    merkle_roll.append(leaf).unwrap();
+    off_chain_tree.add_leaf(leaf, 16);
+    assert_eq!(
+        merkle_roll.get_change_log().root,
+        off_chain_tree.get_root(),
+        "Failed to append accurately to merkle roll after subtree append",
+    );
+}
+
+#[tokio::test(threaded_scheduler)]
+async fn test_append_incomplete_subtree() {
+    let (mut merkle_roll, mut off_chain_tree) = setup();
+    let mut rng = thread_rng();
+    merkle_roll.initialize().unwrap();
+
+    // insert four leaves into the large tree
+    for i in 0..4 {
+        let leaf = rng.gen::<[u8; 32]>();
+        merkle_roll.append(leaf).unwrap();
+        off_chain_tree.add_leaf(leaf, i);
+        assert_eq!(
+            merkle_roll.get_change_log().root,
+            off_chain_tree.get_root(),
+            "On chain tree failed to update properly on an append",
+        );
+    }
+
+    // create a simple subtree to append of depth three
+    let mut onchain_subtree = MerkleRoll::<2, 8>::new();
+    onchain_subtree.initialize().unwrap();
+
+    // append leaves to the subtree, and also append them to the off-chain tree
+    // note: this gives us a partially filled tree, the other two leaves are empty nodes
+    for i in 4..6 {
+        let leaf = rng.gen::<[u8; 32]>();
+        onchain_subtree.append(leaf).unwrap();
+        off_chain_tree.add_leaf(leaf, i);
+    }
+
+    // append the on_chain subtree to the merkle_roll
+    merkle_roll
+        .append_subtree(
+            onchain_subtree.get_change_log().root,
+            onchain_subtree.rightmost_proof.leaf,
+            onchain_subtree.rightmost_proof.index,
+            onchain_subtree.rightmost_proof.proof.to_vec(),
+        )
+        .unwrap();
+
+    // The result should be that the merkle_roll's new root is the same as the root of the off-chain tree which had leaves 0..5 appended one by one
+    assert_eq!(
+        merkle_roll.get_change_log().root,
+        off_chain_tree.get_root(),
+        "On chain tree failed to update properly on an append",
+    );
+
+    // Show that we can still append to the large tree after performing a subtree append
+    let leaf = rng.gen::<[u8; 32]>();
+    merkle_roll.append(leaf).unwrap();
+    off_chain_tree.add_leaf(leaf, 6);
+    assert_eq!(
+        merkle_roll.get_change_log().root,
+        off_chain_tree.get_root(),
+        "Failed to append accurately to merkle roll after subtree append",
+    );
+}
+
 #[tokio::test(threaded_scheduler)]
 async fn test_prove_leaf() {
     let (mut merkle_roll, mut off_chain_tree) = setup();