From 9a18ed6749cb82a865351fffe9cfa52b097d7b3c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philippe=20Laferri=C3=A8re?= Date: Tue, 5 Dec 2023 09:25:21 -0500 Subject: [PATCH] Implement `SimpleSmt::set_subtree` (#232) * recompute_nodes_from_indeX_to_root * MerkleError variant * set_subtree * test_simplesmt_set_subtree * test_simplesmt_set_subtree_entire_tree * test * set_subtree: return root --- src/merkle/error.rs | 4 + src/merkle/simple_smt/mod.rs | 93 +++++++++++++++++++--- src/merkle/simple_smt/tests.rs | 137 ++++++++++++++++++++++++++++++++- 3 files changed, 222 insertions(+), 12 deletions(-) diff --git a/src/merkle/error.rs b/src/merkle/error.rs index ba236bf9b..b513212f5 100644 --- a/src/merkle/error.rs +++ b/src/merkle/error.rs @@ -13,6 +13,7 @@ pub enum MerkleError { DuplicateValuesForKey(RpoDigest), InvalidIndex { depth: u8, value: u64 }, InvalidDepth { expected: u8, provided: u8 }, + InvalidSubtreeDepth { subtree_depth: u8, tree_depth: u8 }, InvalidPath(MerklePath), InvalidNumEntries(usize), NodeNotInSet(NodeIndex), @@ -36,6 +37,9 @@ impl fmt::Display for MerkleError { InvalidDepth { expected, provided } => { write!(f, "the provided depth {provided} is not valid for {expected}") } + InvalidSubtreeDepth { subtree_depth, tree_depth } => { + write!(f, "tried inserting a subtree of depth {subtree_depth} into a tree of depth {tree_depth}") + } InvalidPath(_path) => write!(f, "the provided path is not valid"), InvalidNumEntries(max) => write!(f, "number of entries exceeded the maximum: {max}"), NodeNotInSet(index) => write!(f, "the node with index ({index}) is not in the set"), diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index 9c8049f9f..4ce94a414 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -229,7 +229,7 @@ impl SimpleSmt { /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}. pub fn update_leaf(&mut self, index: u64, value: Word) -> Result { // validate the index before modifying the structure - let mut idx = NodeIndex::new(self.depth(), index)?; + let idx = NodeIndex::new(self.depth(), index)?; let old_value = self.insert_leaf_node(index, value).unwrap_or(Self::EMPTY_VALUE); @@ -238,22 +238,93 @@ impl SimpleSmt { return Ok(value); } - let mut value = RpoDigest::from(value); - for _ in 0..idx.depth() { - let is_right = idx.is_value_odd(); - idx.move_up(); - let BranchNode { left, right } = self.get_branch_node(&idx); - let (left, right) = if is_right { (left, value) } else { (value, right) }; - self.insert_branch_node(idx, left, right); - value = Rpo256::merge(&[left, right]); - } - self.root = value; + self.recompute_nodes_from_index_to_root(idx, RpoDigest::from(value)); + Ok(old_value) } + /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is + /// computed as `self.depth() - subtree.depth()`. + /// + /// Returns the new root. + pub fn set_subtree( + &mut self, + subtree_insertion_index: u64, + subtree: SimpleSmt, + ) -> Result { + if subtree.depth() > self.depth() { + return Err(MerkleError::InvalidSubtreeDepth { + subtree_depth: subtree.depth(), + tree_depth: self.depth(), + }); + } + + // Verify that `subtree_insertion_index` is valid. + let subtree_root_insertion_depth = self.depth() - subtree.depth(); + let subtree_root_index = + NodeIndex::new(subtree_root_insertion_depth, subtree_insertion_index)?; + + // add leaves + // -------------- + + // The subtree's leaf indices live in their own context - i.e. a subtree of depth `d`. If we + // insert the subtree at `subtree_insertion_index = 0`, then the subtree leaf indices are + // valid as they are. However, consider what happens when we insert at + // `subtree_insertion_index = 1`. The first leaf of our subtree now will have index `2^d`; + // you can see it as there's a full subtree sitting on its left. In general, for + // `subtree_insertion_index = i`, there are `i` subtrees sitting before the subtree we want + // to insert, so we need to adjust all its leaves by `i * 2^d`. + let leaf_index_shift: u64 = subtree_insertion_index * 2_u64.pow(subtree.depth().into()); + for (subtree_leaf_idx, leaf_value) in subtree.leaves() { + let new_leaf_idx = leaf_index_shift + subtree_leaf_idx; + debug_assert!(new_leaf_idx < 2_u64.pow(self.depth().into())); + + self.insert_leaf_node(new_leaf_idx, *leaf_value); + } + + // add subtree's branch nodes (which includes the root) + // -------------- + for (branch_idx, branch_node) in subtree.branches { + let new_branch_idx = { + let new_depth = subtree_root_insertion_depth + branch_idx.depth(); + let new_value = subtree_insertion_index * 2_u64.pow(branch_idx.depth().into()) + + branch_idx.value(); + + NodeIndex::new(new_depth, new_value).expect("index guaranteed to be valid") + }; + + self.branches.insert(new_branch_idx, branch_node); + } + + // recompute nodes starting from subtree root + // -------------- + self.recompute_nodes_from_index_to_root(subtree_root_index, subtree.root); + + Ok(self.root) + } + // HELPER METHODS // -------------------------------------------------------------------------------------------- + /// Recomputes the branch nodes (including the root) from `index` all the way to the root. + /// `node_hash_at_index` is the hash of the node stored at index. + fn recompute_nodes_from_index_to_root( + &mut self, + mut index: NodeIndex, + node_hash_at_index: RpoDigest, + ) { + let mut value = node_hash_at_index; + for _ in 0..index.depth() { + let is_right = index.is_value_odd(); + index.move_up(); + let BranchNode { left, right } = self.get_branch_node(&index); + let (left, right) = if is_right { (left, value) } else { (value, right) }; + self.insert_branch_node(index, left, right); + value = Rpo256::merge(&[left, right]); + } + self.root = value; + } + fn get_leaf_node(&self, key: u64) -> Option { self.leaves.get(&key).copied() } diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs index 3fc93865f..3d270be7a 100644 --- a/src/merkle/simple_smt/tests.rs +++ b/src/merkle/simple_smt/tests.rs @@ -3,7 +3,7 @@ use super::{ NodeIndex, Rpo256, Vec, }; use crate::{ - merkle::{digests_to_words, int_to_leaf, int_to_node}, + merkle::{digests_to_words, int_to_leaf, int_to_node, EmptySubtreeRoots}, Word, }; @@ -349,6 +349,141 @@ fn test_simplesmt_with_leaves_nonexisting_leaf() { assert!(result.is_err()); } +#[test] +fn test_simplesmt_set_subtree() { + // Final Tree: + // + // ____k____ + // / \ + // _i_ _j_ + // / \ / \ + // e f g h + // / \ / \ / \ / \ + // a b 0 0 c 0 0 d + + let z = EMPTY_WORD; + + let a = Word::from(Rpo256::merge(&[z.into(); 2])); + let b = Word::from(Rpo256::merge(&[a.into(); 2])); + let c = Word::from(Rpo256::merge(&[b.into(); 2])); + let d = Word::from(Rpo256::merge(&[c.into(); 2])); + + let e = Rpo256::merge(&[a.into(), b.into()]); + let f = Rpo256::merge(&[z.into(), z.into()]); + let g = Rpo256::merge(&[c.into(), z.into()]); + let h = Rpo256::merge(&[z.into(), d.into()]); + + let i = Rpo256::merge(&[e, f]); + let j = Rpo256::merge(&[g, h]); + + let k = Rpo256::merge(&[i, j]); + + // subtree: + // g + // / \ + // c 0 + let subtree = { + let depth = 1; + let entries = vec![(0, c)]; + SimpleSmt::with_leaves(depth, entries).unwrap() + }; + + // insert subtree + let tree = { + let depth = 3; + let entries = vec![(0, a), (1, b), (7, d)]; + let mut tree = SimpleSmt::with_leaves(depth, entries).unwrap(); + + tree.set_subtree(2, subtree).unwrap(); + + tree + }; + + assert_eq!(tree.root(), k); + assert_eq!(tree.get_leaf(4).unwrap(), c); + assert_eq!(tree.get_branch_node(&NodeIndex::new_unchecked(2, 2)).parent(), g); +} + +/// Ensures that an invalid input node index into `set_subtree()` incurs no mutation of the tree +#[test] +fn test_simplesmt_set_subtree_unchanged_for_wrong_index() { + // Final Tree: + // + // ____k____ + // / \ + // _i_ _j_ + // / \ / \ + // e f g h + // / \ / \ / \ / \ + // a b 0 0 c 0 0 d + + let z = EMPTY_WORD; + + let a = Word::from(Rpo256::merge(&[z.into(); 2])); + let b = Word::from(Rpo256::merge(&[a.into(); 2])); + let c = Word::from(Rpo256::merge(&[b.into(); 2])); + let d = Word::from(Rpo256::merge(&[c.into(); 2])); + + // subtree: + // g + // / \ + // c 0 + let subtree = { + let depth = 1; + let entries = vec![(0, c)]; + SimpleSmt::with_leaves(depth, entries).unwrap() + }; + + let mut tree = { + let depth = 3; + let entries = vec![(0, a), (1, b), (7, d)]; + SimpleSmt::with_leaves(depth, entries).unwrap() + }; + let tree_root_before_insertion = tree.root(); + + // insert subtree + assert!(tree.set_subtree(500, subtree).is_err()); + + assert_eq!(tree.root(), tree_root_before_insertion); +} + +/// We insert an empty subtree that has the same depth as the original tree +#[test] +fn test_simplesmt_set_subtree_entire_tree() { + // Initial Tree: + // + // ____k____ + // / \ + // _i_ _j_ + // / \ / \ + // e f g h + // / \ / \ / \ / \ + // a b 0 0 c 0 0 d + + let z = EMPTY_WORD; + + let a = Word::from(Rpo256::merge(&[z.into(); 2])); + let b = Word::from(Rpo256::merge(&[a.into(); 2])); + let c = Word::from(Rpo256::merge(&[b.into(); 2])); + let d = Word::from(Rpo256::merge(&[c.into(); 2])); + + let depth = 3; + + // subtree: E3 + let subtree = { SimpleSmt::with_leaves(depth, Vec::new()).unwrap() }; + assert_eq!(subtree.root(), *EmptySubtreeRoots::entry(depth, 0)); + + // insert subtree + let mut tree = { + let entries = vec![(0, a), (1, b), (4, c), (7, d)]; + SimpleSmt::with_leaves(depth, entries).unwrap() + }; + + tree.set_subtree(0, subtree).unwrap(); + + assert_eq!(tree.root(), *EmptySubtreeRoots::entry(depth, 0)); +} + // HELPER FUNCTIONS // --------------------------------------------------------------------------------------------