diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 74ed99a..7b85a11 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -166,6 +166,25 @@ impl Smt { >::insert(self, key, value) } + /// Like [`Self::insert()`], but only performs the insert if the the new tree's root + /// hash would be equal to the hash given in `expected_root`. + /// + /// # Errors + /// Returns [`MerkleError::ConflictingRoots`] with a two-item [Vec] if the new root of the tree is + /// different from the expected root. The first item of the vector is the expected root, and the + /// second is actual root. + /// + /// No mutations are performed if the roots do no match. + pub fn insert_ensure_root( + &mut self, + key: RpoDigest, + value: Word, + expected_root: RpoDigest, + ) -> Result + { + >::insert_ensure_root(self, key, value, expected_root) + } + // HELPERS // -------------------------------------------------------------------------------------------- diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 27d24bd..c4e9caa 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -1,6 +1,6 @@ use super::{Felt, LeafIndex, NodeIndex, Rpo256, RpoDigest, Smt, SmtLeaf, EMPTY_WORD, SMT_DEPTH}; use crate::{ - merkle::{smt::SparseMerkleTree, EmptySubtreeRoots, MerkleStore}, + merkle::{smt::SparseMerkleTree, EmptySubtreeRoots, MerkleError, MerkleStore}, utils::{Deserializable, Serializable}, Word, ONE, WORD_SIZE, }; @@ -258,8 +258,11 @@ fn test_smt_removal() { } #[test] -fn test_prospective_hash() { +fn test_checked_insertion() { + use MerkleError::ConflictingRoots; + let mut smt = Smt::default(); + let smt_empty = smt.clone(); let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64; @@ -273,50 +276,91 @@ fn test_prospective_hash() { let value_2 = [2_u32.into(); WORD_SIZE]; let value_3: [Felt; 4] = [3_u32.into(); WORD_SIZE]; + let root_empty = smt.root(); + // insert key-value 1 - { + let root_1 = { let prospective = smt.hash_prospective_leaf(&key_1, &value_1); let old_value_1 = smt.insert(key_1, value_1); assert_eq!(old_value_1, EMPTY_WORD); - assert_eq!(smt.get_leaf(&key_1).hash(), prospective); - + assert_eq!(prospective, smt.get_leaf(&key_1).hash()); assert_eq!(smt.get_leaf(&key_1), SmtLeaf::Single((key_1, value_1))); + smt.root() + }; + + { + // Trying to insert something else into key_1 with the existing root should fail, and + // should not modify the tree at all. + let smt_before = smt.clone(); + assert!(matches!(smt.insert_ensure_root(key_1, value_2, root_1), Err(ConflictingRoots(_)))); + assert_eq!(smt, smt_before); + + // And inserting an empty word should bring us back to where we were. + assert_eq!(smt.insert_ensure_root(key_1, EMPTY_WORD, root_empty), Ok(value_1)); + assert_eq!(smt, smt_empty); + + smt.insert_ensure_root(key_1, value_1, root_1).unwrap(); + assert_eq!(smt, smt_before); } // insert key-value 2 - { + let root_2 = { let prospective = smt.hash_prospective_leaf(&key_2, &value_2); let old_value_2 = smt.insert(key_2, value_2); assert_eq!(old_value_2, EMPTY_WORD); - assert_eq!(smt.get_leaf(&key_2).hash(), prospective); + assert_eq!(prospective, smt.get_leaf(&key_2).hash()); assert_eq!( smt.get_leaf(&key_2), SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)]) ); + + smt.root() + }; + + { + let smt_before = smt.clone(); + assert!(matches!(smt.insert_ensure_root(key_2, value_1, root_2), Err(ConflictingRoots(_)))); + assert_eq!(smt, smt_before); + + assert_eq!(smt.insert_ensure_root(key_2, EMPTY_WORD, root_1), Ok(value_2)); + smt.insert_ensure_root(key_2, value_2, root_2).unwrap(); + assert_eq!(smt, smt_before); } // insert key-value 3 - { - let prospective_hash = smt.hash_prospective_leaf(&key_3, &value_3); + let root_3 = { + let prospective = smt.hash_prospective_leaf(&key_3, &value_3); let old_value_3 = smt.insert(key_3, value_3); assert_eq!(old_value_3, EMPTY_WORD); - assert_eq!(smt.get_leaf(&key_3).hash(), prospective_hash); + assert_eq!(prospective, smt.get_leaf(&key_3).hash()); assert_eq!( smt.get_leaf(&key_3), SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2), (key_3, value_3)]) ); + + smt.root() + }; + + { + let smt_before = smt.clone(); + assert!(matches!(smt.insert_ensure_root(key_3, value_1, root_3), Err(ConflictingRoots(_)))); + assert_eq!(smt, smt_before); + + assert_eq!(smt.insert_ensure_root(key_3, EMPTY_WORD, root_2), Ok(value_3)); + smt.insert_ensure_root(key_3, value_3, root_3).unwrap(); + assert_eq!(smt, smt_before); } // remove key 3 { let old_hash = smt.get_leaf(&key_3).hash(); - let old_value_3 = smt.insert(key_3, EMPTY_WORD); + let old_value_3 = smt.insert_ensure_root(key_3, EMPTY_WORD, root_2).unwrap(); assert_eq!(old_value_3, value_3); assert_eq!(old_hash, smt.hash_prospective_leaf(&key_3, &old_value_3)); @@ -324,26 +368,32 @@ fn test_prospective_hash() { smt.get_leaf(&key_3), SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)]) ); + + assert_eq!(smt.root(), root_2); } // remove key 2 { let old_hash = smt.get_leaf(&key_2).hash(); - let old_value_2 = smt.insert(key_2, EMPTY_WORD); + let old_value_2 = smt.insert_ensure_root(key_2, EMPTY_WORD, root_1).unwrap(); assert_eq!(old_value_2, value_2); assert_eq!(old_hash, smt.hash_prospective_leaf(&key_2, &old_value_2)); assert_eq!(smt.get_leaf(&key_2), SmtLeaf::Single((key_1, value_1))); + + assert_eq!(smt.root(), root_1); } // remove key 1 { let old_hash = smt.get_leaf(&key_1).hash(); - let old_value_1 = smt.insert(key_1, EMPTY_WORD); + let old_value_1 = smt.insert_ensure_root(key_1, EMPTY_WORD, root_empty).unwrap(); assert_eq!(old_value_1, value_1); assert_eq!(old_hash, smt.hash_prospective_leaf(&key_1, &old_value_1)); assert_eq!(smt.get_leaf(&key_1), SmtLeaf::new_empty(key_1.into())); + + assert_eq!(smt.root(), root_empty); } } diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 9330605..ebc2f5d 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -119,6 +119,106 @@ pub(crate) trait SparseMerkleTree { old_value } + /// Like [`Self::insert()`], but only performs the insert if the the new tree's root + /// hash would be equal to the hash given in `expected_root`. + /// + /// # Errors + /// Returns [`MerkleError::ConflictingRoots`] with a two-item [Vec] if the new root of the tree is + /// different from the expected root. The first item of the vector is the expected root, and the + /// second is actual root. + /// + /// No mutations are performed if the roots do no match. + fn insert_ensure_root( + &mut self, + key: Self::Key, + value: Self::Value, + expected_root: RpoDigest, + ) -> Result + { + + let old_value = self.get_value(&key); + // if the old value and new value are the same, there is nothing to update + if value == old_value { + return Ok(value); + } + + // Compute the nodes we'll need to make and remove. + let mut removals: Vec = Vec::with_capacity(DEPTH as usize); + let mut additions: Vec<(NodeIndex, InnerNode)> = Vec::with_capacity(DEPTH as usize); + + let (mut node_index, mut parent_node) = { + let leaf_index: LeafIndex = Self::key_to_leaf_index(&key); + let node_index = NodeIndex::from(leaf_index); + + let mut parent_index = node_index.clone(); + parent_index.move_up(); + + (node_index, Some(self.get_inner_node(parent_index))) + }; + + let mut new_child_hash = self.hash_prospective_leaf(&key, &value); + for node_depth in (0..node_index.depth()).rev() { + let is_right = node_index.is_value_odd(); + node_index.move_up(); + + let old_node = match parent_node.take() { + // On the first iteration, the 'old node' is the parent of the + // perspective leaf. + Some(parent_node) => parent_node, + // Otherwise it's a regular existing node. + None => self.get_inner_node(node_index), + }; + + //let new_node = new_node_from(is_right, old_node, new_child_hash); + let new_node = if is_right { + InnerNode { + left: old_node.left, + right: new_child_hash, + } + } else { + InnerNode { + left: new_child_hash, + right: old_node.right, + } + }; + + // The next iteration will operate on this node's new hash. + new_child_hash = new_node.hash(); + + let &equivalent_empty_hash = EmptySubtreeRoots::entry(DEPTH, node_depth); + if new_child_hash == equivalent_empty_hash { + // If a subtree is empty, we can remove the inner node, since it's equal to the + // default value. + removals.push(node_index); + } else { + additions.push((node_index, new_node)); + } + } + + // Once we're at depth 0, the last node we made is the new root. + let new_root = new_child_hash; + + if expected_root != new_root { + return Err(MerkleError::ConflictingRoots(vec![expected_root, new_root])); + } + + // Actual mutations start here. + + self.insert_value(key, value); + + for index in removals.drain(..) { + self.remove_inner_node(index); + } + + for (index, new_node) in additions.drain(..) { + self.insert_inner_node(index, new_node); + } + + self.set_root(new_root); + + Ok(old_value) + } + /// 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. /// diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index a267045..14307a9 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -187,6 +187,25 @@ impl SimpleSmt { >::insert(self, key, value) } + /// Like [`Self::insert()`], but only performs the insert if the the new tree's root + /// hash would be equal to the hash given in `expected_root`. + /// + /// # Errors + /// Returns [`MerkleError::ConflictingRoots`] with a two-item [Vec] if the new root of the tree is + /// different from the expected root. The first item of the vector is the expected root, and the + /// second is actual root. + /// + /// No mutations are performed if the roots do no match. + pub fn insert_ensure_root( + &mut self, + key: LeafIndex, + value: Word, + expected_root: RpoDigest, + ) -> Result + { + >::insert_ensure_root(self, key, value, expected_root) + } + /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is /// computed as `DEPTH - SUBTREE_DEPTH`. ///