diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 86ebaaf7..0dfd66cf 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -232,6 +232,10 @@ impl SparseMerkleTree for Smt { self.inner_nodes.insert(index, inner_node); } + fn remove_inner_node(&mut self, index: NodeIndex) { + let _ = self.inner_nodes.remove(&index); + } + fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option { // inserting an `EMPTY_VALUE` is equivalent to removing any value associated with `key` if value != Self::EMPTY_VALUE { diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 4a6d3e31..25672366 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -25,7 +25,7 @@ fn test_smt_insert_at_same_key() { // Insert value 1 and ensure root is as expected { - let leaf_node = build_single_leaf_node(key_1, value_1); + let leaf_node = build_empty_or_single_leaf_node(key_1, value_1); let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root; let old_value_1 = smt.insert(key_1, value_1); @@ -36,7 +36,7 @@ fn test_smt_insert_at_same_key() { // Insert value 2 and ensure root is as expected { - let leaf_node = build_single_leaf_node(key_1, value_2); + let leaf_node = build_empty_or_single_leaf_node(key_1, value_2); let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root; let old_value_2 = smt.insert(key_1, value_2); @@ -64,7 +64,7 @@ fn test_smt_insert_at_same_key_2() { let mut store: MerkleStore = { let mut store = MerkleStore::default(); - let leaf_node = build_single_leaf_node(key_already_present, value_already_present); + let leaf_node = build_empty_or_single_leaf_node(key_already_present, value_already_present); store .set_node(*EmptySubtreeRoots::entry(SMT_DEPTH, 0), key_already_present_index, leaf_node) .unwrap(); @@ -109,10 +109,29 @@ fn test_smt_insert_at_same_key_2() { } } -/// This test ensures that the root of the tree is as expected when we add 3 items at 3 different -/// keys. This also tests that the merkle paths produced are as expected. +/// This test ensures that the root of the tree is as expected when we add/remove 3 items at 3 +/// different keys. This also tests that the merkle paths produced are as expected. #[test] -fn test_smt_insert_multiple_values() { +fn test_smt_insert_and_remove_multiple_values() { + fn insert_values_and_assert_path( + smt: &mut Smt, + store: &mut MerkleStore, + key_values: &[(RpoDigest, Word)], + ) { + for &(key, value) in key_values { + let key_index: NodeIndex = LeafIndex::::from(key).into(); + + let leaf_node = build_empty_or_single_leaf_node(key, value); + let tree_root = store.set_node(smt.root(), key_index, leaf_node).unwrap().root; + + let _ = smt.insert(key, value); + + assert_eq!(smt.root(), tree_root); + + let expected_path = store.get_path(tree_root, key_index).unwrap(); + assert_eq!(smt.open(&key).0, expected_path.path); + } + } let mut smt = Smt::default(); let mut store: MerkleStore = MerkleStore::default(); @@ -140,22 +159,20 @@ fn test_smt_insert_multiple_values() { let value_2 = [ONE + ONE; WORD_SIZE]; let value_3 = [ONE + ONE + ONE; WORD_SIZE]; + // Insert values in the tree let key_values = [(key_1, value_1), (key_2, value_2), (key_3, value_3)]; + insert_values_and_assert_path(&mut smt, &mut store, &key_values); - for (key, value) in key_values { - let key_index: NodeIndex = LeafIndex::::from(key).into(); + // Remove values from the tree + let key_empty_values = [(key_1, EMPTY_WORD), (key_2, EMPTY_WORD), (key_3, EMPTY_WORD)]; + insert_values_and_assert_path(&mut smt, &mut store, &key_empty_values); - let leaf_node = build_single_leaf_node(key, value); - let tree_root = store.set_node(smt.root(), key_index, leaf_node).unwrap().root; + let empty_root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0); + assert_eq!(smt.root(), empty_root); - let old_value = smt.insert(key, value); - assert_eq!(old_value, EMPTY_WORD); - - assert_eq!(smt.root(), tree_root); - - let expected_path = store.get_path(tree_root, key_index).unwrap(); - assert_eq!(smt.open(&key).0, expected_path.path); - } + // an empty tree should have no leaves or inner nodes + assert!(smt.leaves.is_empty()); + assert!(smt.inner_nodes.is_empty()); } /// This tests that inserting the empty value does indeed remove the key-value contained at the @@ -309,8 +326,12 @@ fn test_smt_entries() { // HELPERS // -------------------------------------------------------------------------------------------- -fn build_single_leaf_node(key: RpoDigest, value: Word) -> RpoDigest { - SmtLeaf::Single((key, value)).hash() +fn build_empty_or_single_leaf_node(key: RpoDigest, value: Word) -> RpoDigest { + if value == EMPTY_WORD { + SmtLeaf::Empty.hash() + } else { + SmtLeaf::Single((key, value)).hash() + } } fn build_multiple_leaf_node(kv_pairs: &[(RpoDigest, Word)]) -> RpoDigest { diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 63e6ac89..c8f2ba1f 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -5,7 +5,7 @@ use crate::{ Word, }; -use super::{MerkleError, MerklePath, NodeIndex, Vec}; +use super::{EmptySubtreeRoots, MerkleError, MerklePath, NodeIndex, Vec}; mod full; pub use full::{Smt, SmtLeaf, SMT_DEPTH}; @@ -118,16 +118,27 @@ pub(crate) trait SparseMerkleTree { mut index: NodeIndex, node_hash_at_index: RpoDigest, ) { - let mut value = node_hash_at_index; - for _ in 0..index.depth() { + let mut node_hash = node_hash_at_index; + for node_depth in (0..index.depth()).rev() { let is_right = index.is_value_odd(); index.move_up(); let InnerNode { left, right } = self.get_inner_node(index); - let (left, right) = if is_right { (left, value) } else { (value, right) }; - self.insert_inner_node(index, InnerNode { left, right }); - value = Rpo256::merge(&[left, right]); + let (left, right) = if is_right { + (left, node_hash) + } else { + (node_hash, right) + }; + node_hash = Rpo256::merge(&[left, right]); + + if node_hash == *EmptySubtreeRoots::entry(DEPTH, node_depth) { + // If a subtree is empty, when can remove the inner node, since it's equal to the + // default value + self.remove_inner_node(index) + } else { + self.insert_inner_node(index, InnerNode { left, right }); + } } - self.set_root(value); + self.set_root(node_hash); } // REQUIRED METHODS @@ -145,6 +156,9 @@ pub(crate) trait SparseMerkleTree { /// Inserts an inner node at the given index fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode); + /// Removes an inner node at the given index + fn remove_inner_node(&mut self, index: NodeIndex); + /// Inserts a leaf node, and returns the value at the key if already exists fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option; diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 02b3b343..9ccc0f63 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -275,6 +275,10 @@ impl SparseMerkleTree for SimpleSmt { self.inner_nodes.insert(index, inner_node); } + fn remove_inner_node(&mut self, index: NodeIndex) { + let _ = self.inner_nodes.remove(&index); + } + fn insert_value(&mut self, key: LeafIndex, value: Word) -> Option { self.leaves.insert(key.value(), value) }