From 4bf087daf860b3147c45fc7fa2a725ca66351693 Mon Sep 17 00:00:00 2001
From: Bobbin Threadbare <43513081+bobbinth@users.noreply.github.com>
Date: Tue, 2 Apr 2024 13:01:00 -0700
Subject: [PATCH] fix: decrement leaf count in simple SMT when inserting empty
 value (#303)

---
 CHANGELOG.md                   |  4 ++++
 src/merkle/smt/simple/mod.rs   | 21 +++++++++++----------
 src/merkle/smt/simple/tests.rs | 20 +++++++++++++++++++-
 3 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b413710..66567d2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.9.1 (2024-04-02)
+
+* Added `num_leaves()` method to `SimpleSmt` (#302).
+
 ## 0.9.0 (2024-03-24)
 
 * [BREAKING] Removed deprecated re-exports from liballoc/libstd (#290).
diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs
index 4d14e9c..2fa5ae4 100644
--- a/src/merkle/smt/simple/mod.rs
+++ b/src/merkle/smt/simple/mod.rs
@@ -122,6 +122,11 @@ impl<const DEPTH: u8> SimpleSmt<DEPTH> {
         <Self as SparseMerkleTree<DEPTH>>::root(self)
     }
 
+    /// Returns the number of non-empty leaves in this tree.
+    pub fn num_leaves(&self) -> usize {
+        self.leaves.len()
+    }
+
     /// Returns the leaf at the specified index.
     pub fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
         <Self as SparseMerkleTree<DEPTH>>::get_leaf(self, key)
@@ -152,11 +157,6 @@ impl<const DEPTH: u8> SimpleSmt<DEPTH> {
         <Self as SparseMerkleTree<DEPTH>>::open(self, key)
     }
 
-    /// Returns a count of non-empty leaves.
-    pub fn leaf_count(&self) -> usize {
-        self.leaves.len()
-    }
-
     // ITERATORS
     // --------------------------------------------------------------------------------------------
 
@@ -281,17 +281,18 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
     }
 
     fn insert_value(&mut self, key: LeafIndex<DEPTH>, value: Word) -> Option<Word> {
-        self.leaves.insert(key.value(), value)
+        if value == Self::EMPTY_VALUE {
+            self.leaves.remove(&key.value())
+        } else {
+            self.leaves.insert(key.value(), value)
+        }
     }
 
     fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
-        // the lookup in empty_hashes could fail only if empty_hashes were not built correctly
-        // by the constructor as we check the depth of the lookup above.
         let leaf_pos = key.value();
-
         match self.leaves.get(&leaf_pos) {
             Some(word) => *word,
-            None => Word::from(*EmptySubtreeRoots::entry(DEPTH, DEPTH)),
+            None => Self::EMPTY_VALUE,
         }
     }
 
diff --git a/src/merkle/smt/simple/tests.rs b/src/merkle/smt/simple/tests.rs
index 9a94a65..96ec7df 100644
--- a/src/merkle/smt/simple/tests.rs
+++ b/src/merkle/smt/simple/tests.rs
@@ -50,6 +50,8 @@ fn build_sparse_tree() {
     let mut smt = SimpleSmt::<DEPTH>::new().unwrap();
     let mut values = ZERO_VALUES8.to_vec();
 
+    assert_eq!(smt.num_leaves(), 0);
+
     // insert single value
     let key = 6;
     let new_node = int_to_leaf(7);
@@ -62,6 +64,7 @@ fn build_sparse_tree() {
         smt.open(&LeafIndex::<3>::new(6).unwrap()).path
     );
     assert_eq!(old_value, EMPTY_WORD);
+    assert_eq!(smt.num_leaves(), 1);
 
     // insert second value at distinct leaf branch
     let key = 2;
@@ -75,6 +78,7 @@ fn build_sparse_tree() {
         smt.open(&LeafIndex::<3>::new(2).unwrap()).path
     );
     assert_eq!(old_value, EMPTY_WORD);
+    assert_eq!(smt.num_leaves(), 2);
 }
 
 /// Tests that [`SimpleSmt::with_contiguous_leaves`] works as expected
@@ -146,10 +150,11 @@ fn test_inner_node_iterator() -> Result<(), MerkleError> {
 }
 
 #[test]
-fn update_leaf() {
+fn test_insert() {
     const DEPTH: u8 = 3;
     let mut tree =
         SimpleSmt::<DEPTH>::with_leaves(KEYS8.into_iter().zip(digests_to_words(&VALUES8))).unwrap();
+    assert_eq!(tree.num_leaves(), 8);
 
     // update one value
     let key = 3;
@@ -161,6 +166,7 @@ fn update_leaf() {
     let old_leaf = tree.insert(LeafIndex::<DEPTH>::new(key as u64).unwrap(), new_node);
     assert_eq!(expected_tree.root(), tree.root);
     assert_eq!(old_leaf, *VALUES8[key]);
+    assert_eq!(tree.num_leaves(), 8);
 
     // update another value
     let key = 6;
@@ -171,6 +177,18 @@ fn update_leaf() {
     let old_leaf = tree.insert(LeafIndex::<DEPTH>::new(key as u64).unwrap(), new_node);
     assert_eq!(expected_tree.root(), tree.root);
     assert_eq!(old_leaf, *VALUES8[key]);
+    assert_eq!(tree.num_leaves(), 8);
+
+    // set a leaf to empty value
+    let key = 5;
+    let new_node = EMPTY_WORD;
+    expected_values[key] = new_node;
+    let expected_tree = MerkleTree::new(expected_values.clone()).unwrap();
+
+    let old_leaf = tree.insert(LeafIndex::<DEPTH>::new(key as u64).unwrap(), new_node);
+    assert_eq!(expected_tree.root(), tree.root);
+    assert_eq!(old_leaf, *VALUES8[key]);
+    assert_eq!(tree.num_leaves(), 7);
 }
 
 #[test]