From fcb38c2000304a4ec97f269682e4cbe4c491e969 Mon Sep 17 00:00:00 2001 From: Philippe Laferriere Date: Sat, 3 Feb 2024 08:37:44 -0500 Subject: [PATCH 1/7] add `remove_inner_node` method --- src/merkle/smt/full/mod.rs | 5 +++++ src/merkle/smt/mod.rs | 3 +++ src/merkle/smt/simple/mod.rs | 4 ++++ 3 files changed, 12 insertions(+) diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 86ebaaf7..518f3188 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 { @@ -258,6 +262,7 @@ impl SparseMerkleTree for Smt { let most_significant_felt = key[3]; LeafIndex::new_max_depth(most_significant_felt.as_int()) } + } impl Default for Smt { diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 63e6ac89..df0ece93 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -145,6 +145,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..f5a22e89 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -274,6 +274,10 @@ impl SparseMerkleTree for SimpleSmt { fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) { 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) From e60044fc66871105f8dce76beb3cdf72a3c96b5e Mon Sep 17 00:00:00 2001 From: Philippe Laferriere Date: Sat, 3 Feb 2024 08:42:22 -0500 Subject: [PATCH 2/7] remove inner node --- src/merkle/smt/mod.rs | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index df0ece93..d6a45480 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}; @@ -119,13 +119,20 @@ pub(crate) trait SparseMerkleTree { node_hash_at_index: RpoDigest, ) { let mut value = node_hash_at_index; - for _ in 0..index.depth() { + for node_depth in 0..index.depth() { 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]); + + if value == *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); } From 36ae9623ae50dc383163d7e925dcba86a8a485df Mon Sep 17 00:00:00 2001 From: Philippe Laferriere Date: Sat, 3 Feb 2024 08:51:20 -0500 Subject: [PATCH 3/7] test: remove values and check path --- src/merkle/smt/full/tests.rs | 45 +++++++++++++++++++++++------------- 1 file changed, 29 insertions(+), 16 deletions(-) diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 4a6d3e31..c8f28b6a 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -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_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,16 @@ 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 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); - } + let empty_root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0); + assert_eq!(smt.root(), empty_root); } /// This tests that inserting the empty value does indeed remove the key-value contained at the From 2f8bccb58c733b3dda288bca4de92bd14141788d Mon Sep 17 00:00:00 2001 From: Philippe Laferriere Date: Sat, 3 Feb 2024 09:14:04 -0500 Subject: [PATCH 4/7] remove inner node when empty --- src/merkle/smt/mod.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index d6a45480..16cc3ca2 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -118,15 +118,15 @@ pub(crate) trait SparseMerkleTree { mut index: NodeIndex, node_hash_at_index: RpoDigest, ) { - let mut value = node_hash_at_index; - for node_depth 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) }; - 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 value == *EmptySubtreeRoots::entry(DEPTH, node_depth) { + 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) @@ -134,7 +134,7 @@ pub(crate) trait SparseMerkleTree { self.insert_inner_node(index, InnerNode { left, right }); } } - self.set_root(value); + self.set_root(node_hash); } // REQUIRED METHODS From 7e6b2989988febb5aef19a329dad6ab0523efdcc Mon Sep 17 00:00:00 2001 From: Philippe Laferriere Date: Sat, 3 Feb 2024 13:18:06 -0500 Subject: [PATCH 5/7] fix test --- src/merkle/smt/full/tests.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index c8f28b6a..94fde6cb 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(); @@ -121,7 +121,7 @@ fn test_smt_insert_and_remove_multiple_values() { for &(key, value) in key_values { let key_index: NodeIndex = LeafIndex::::from(key).into(); - let leaf_node = build_single_leaf_node(key, value); + 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); @@ -322,8 +322,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 { From e599e1703affb7bdc909c4972a106647edc3d9be Mon Sep 17 00:00:00 2001 From: Philippe Laferriere Date: Sat, 3 Feb 2024 13:28:23 -0500 Subject: [PATCH 6/7] fmt --- src/merkle/smt/full/mod.rs | 1 - src/merkle/smt/mod.rs | 6 +++++- src/merkle/smt/simple/mod.rs | 2 +- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 518f3188..0dfd66cf 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -262,7 +262,6 @@ impl SparseMerkleTree for Smt { let most_significant_felt = key[3]; LeafIndex::new_max_depth(most_significant_felt.as_int()) } - } impl Default for Smt { diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 16cc3ca2..c8f2ba1f 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -123,7 +123,11 @@ pub(crate) trait SparseMerkleTree { 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, node_hash) } else { (node_hash, 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) { diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index f5a22e89..9ccc0f63 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -274,7 +274,7 @@ impl SparseMerkleTree for SimpleSmt { fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) { self.inner_nodes.insert(index, inner_node); } - + fn remove_inner_node(&mut self, index: NodeIndex) { let _ = self.inner_nodes.remove(&index); } From b2fc301e989869c2f1c61aa365d34d56153b8ddc Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Sun, 4 Feb 2024 00:32:41 -0800 Subject: [PATCH 7/7] test: add assert for Smt empty leaves and nodes --- src/merkle/smt/full/tests.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 94fde6cb..25672366 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -169,6 +169,10 @@ fn test_smt_insert_and_remove_multiple_values() { let empty_root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0); assert_eq!(smt.root(), empty_root); + + // 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