cmt: add method to check if the tree is empty (#3640)

Author: ngundotra

libraries/concurrent-merkle-tree/src/concurrent_merkle_tree.rs

@@ -146,6 +146,23 @@ impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize>
         Ok(root)
     }
 
+    /// Errors if one of the leaves of the current merkle tree is non-EMPTY
+    pub fn prove_tree_is_empty(&self) -> Result<(), ConcurrentMerkleTreeError> {
+        let mut empty_node_cache = Box::new([EMPTY; MAX_DEPTH]);
+        if self.get_root()
+            != empty_node_cached::<MAX_DEPTH>(MAX_DEPTH as u32, &mut empty_node_cache)
+        {
+            return Err(ConcurrentMerkleTreeError::TreeNonEmpty);
+        }
+        Ok(())
+    }
+
+    /// Returns the current root of the merkle tree
+    pub fn get_root(&self) -> [u8; 32] {
+        self.get_change_log().root
+    }
+
+    /// Returns the most recent changelog
     pub fn get_change_log(&self) -> Box<ChangeLog<MAX_DEPTH>> {
         Box::new(self.change_logs[self.active_index as usize])
     }
@@ -416,7 +433,7 @@ impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize>
         proof: &[Node; MAX_DEPTH],
         leaf_index: u32,
     ) -> bool {
-        recompute(leaf, proof, leaf_index) == self.get_change_log().root
+        recompute(leaf, proof, leaf_index) == self.get_root()
     }
 
     /// Note: Enabling `allow_inferred_proof` will fast forward the given proof

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

@@ -32,4 +32,8 @@ pub enum ConcurrentMerkleTreeError {
         "Valid proof was passed to a leaf, but its value has changed since the proof was issued"
     )]
     LeafContentsModified,
+
+    /// Tree has at least 1 non-EMTPY leaf
+    #[error("Tree is not empty")]
+    TreeNonEmpty,
 }

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

@@ -389,3 +389,112 @@ async fn test_append_bug_repro_2() {
     }
     assert_eq!(cmt.get_change_log().root, tree.get_root());
 }
+
+#[tokio::test(flavor = "multi_thread")]
+/// Test that empty trees are checked properly by adding & removing leaves one by one
+async fn test_prove_tree_empty_incremental() {
+    let (mut cmt, mut tree) = setup();
+    let mut rng = thread_rng();
+    cmt.initialize().unwrap();
+
+    cmt.prove_tree_is_empty().unwrap();
+
+    // Append a random leaf & remove it, and make sure that the tree is empty at the end
+    let tree_size = 64;
+    for i in 0..tree_size {
+        let leaf = rng.gen::<[u8; 32]>();
+        cmt.append(leaf).unwrap();
+        tree.add_leaf(leaf, i);
+
+        match cmt.prove_tree_is_empty() {
+            Ok(_) => {
+                panic!("Tree has a leaf in it -- should not be possible to prove empty!")
+            }
+            Err(e) => match e {
+                ConcurrentMerkleTreeError::TreeNonEmpty => {}
+                _ => {
+                    panic!("Wrong error thrown. Expected TreeNonEmpty erro")
+                }
+            },
+        }
+
+        cmt.set_leaf(
+            tree.get_root(),
+            tree.get_leaf(i),
+            EMPTY,
+            &tree.get_proof_of_leaf(i),
+            i as u32,
+        )
+        .unwrap();
+        tree.add_leaf(EMPTY, i);
+
+        cmt.prove_tree_is_empty().unwrap();
+    }
+}
+
+#[tokio::test(flavor = "multi_thread")]
+/// Test that empty trees are checked properly by adding & removing leaves in a batch
+async fn test_prove_tree_empty_batched() {
+    let (mut cmt, mut tree) = setup();
+    let mut rng = thread_rng();
+    cmt.initialize().unwrap();
+
+    // Sanity check
+    cmt.prove_tree_is_empty().unwrap();
+
+    // Add random leaves to the tree
+    let tree_size = 64;
+    for i in 0..tree_size {
+        let leaf = rng.gen::<[u8; 32]>();
+        cmt.append(leaf).unwrap();
+        tree.add_leaf(leaf, i);
+
+        match cmt.prove_tree_is_empty() {
+            Ok(_) => {
+                panic!("Tree has a leaf in it -- should not be possible to prove empty!")
+            }
+            Err(e) => match e {
+                ConcurrentMerkleTreeError::TreeNonEmpty => {}
+                _ => {
+                    panic!("Wrong error thrown. Expected TreeNonEmpty erro")
+                }
+            },
+        }
+    }
+    // Remove random leaves
+    for i in 0..tree_size - 1 {
+        cmt.set_leaf(
+            tree.get_root(),
+            tree.get_leaf(i),
+            EMPTY,
+            &tree.get_proof_of_leaf(i),
+            i as u32,
+        )
+        .unwrap();
+        tree.add_leaf(EMPTY, i);
+
+        match cmt.prove_tree_is_empty() {
+            Ok(_) => {
+                panic!("Tree has a leaf in it -- should not be possible to prove empty!")
+            }
+            Err(e) => match e {
+                ConcurrentMerkleTreeError::TreeNonEmpty => {}
+                _ => {
+                    panic!("Wrong error thrown. Expected TreeNonEmpty erro")
+                }
+            },
+        }
+    }
+    cmt.set_leaf(
+        tree.get_root(),
+        tree.get_leaf(tree_size - 1),
+        EMPTY,
+        &tree.get_proof_of_leaf(tree_size - 1),
+        (tree_size - 1) as u32,
+    )
+    .unwrap();
+    tree.add_leaf(EMPTY, tree_size - 1);
+
+    // Check that the last leaf was successfully removed
+    cmt.prove_tree_is_empty().unwrap();
+}