diff --git a/CHANGELOG.md b/CHANGELOG.md index 76c65a7..7d1dd39 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ - [BREAKING]: renamed `Mmr::open()` into `Mmr::open_at()` and `Mmr::peaks()` into `Mmr::peaks_at()` (#234). - Added `Mmr::open()` and `Mmr::peaks()` which rely on `Mmr::open_at()` and `Mmr::peaks()` respectively (#234). - Standardised CI and Makefile across Miden repos (#323). +- Added `Smt::insert_ensure_root()` for root-validated insertion (#327) ## 0.10.0 (2024-08-06) diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 205a443..60b7bd2 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -167,6 +167,24 @@ impl Smt { >::insert(self, key, value) } + /// Like [`Self::insert()`], but only performs the insert if 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 the actual root. + /// + /// No mutations are performed if the roots do not 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 // -------------------------------------------------------------------------------------------- @@ -274,21 +292,23 @@ impl SparseMerkleTree for Smt { } } - fn hash_prospective_leaf(&self, key: &RpoDigest, value: &Word) -> RpoDigest { - // Future work could avoid cloning the leaf by mirroring some of the insertion logic and - // hashing without an intermediate leaf, but cloning is only expensive for multi-leaves, - // which should be really rare. - let leaf_index: LeafIndex = Self::key_to_leaf_index(key); - match self.leaves.get(&leaf_index.value()) { - None => SmtLeaf::new_single(*key, *value).hash(), - Some(existing_leaf) => { - let mut new_leaf = existing_leaf.clone(); + fn get_prospective_leaf( + &self, + mut existing_leaf: SmtLeaf, + key: &RpoDigest, + value: &Word, + ) -> SmtLeaf { + debug_assert_eq!(existing_leaf.index(), Self::key_to_leaf_index(key)); + + match existing_leaf { + SmtLeaf::Empty(_) => SmtLeaf::new_single(*key, *value), + _ => { if *value != EMPTY_WORD { - new_leaf.insert(*key, *value); + existing_leaf.insert(*key, *value); } else { - new_leaf.remove(*key); + existing_leaf.remove(*key); } - new_leaf.hash() + existing_leaf }, } } diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 1ee7f9d..2c4e8c8 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -2,7 +2,7 @@ use alloc::vec::Vec; 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, }; @@ -259,24 +259,30 @@ 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; let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]); let key_2: RpoDigest = RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]); + // Sort key_3 before key_1, to test non-append insertion. let key_3: RpoDigest = - RpoDigest::from([3_u32.into(), 3_u32.into(), 3_u32.into(), Felt::new(raw)]); + RpoDigest::from([0_u32.into(), 0_u32.into(), 0_u32.into(), Felt::new(raw)]); let value_1 = [ONE; WORD_SIZE]; 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 prospective = smt.hash_prospective_leaf(&key_1, &value_1); + let root_1 = { + let prospective = smt.get_prospective_leaf(smt.get_leaf(&key_1), &key_1, &value_1).hash(); let old_value_1 = smt.insert(key_1, value_1); assert_eq!(old_value_1, EMPTY_WORD); @@ -284,11 +290,30 @@ fn test_prospective_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 prospective = smt.hash_prospective_leaf(&key_2, &value_2); + let root_2 = { + let prospective = smt.get_prospective_leaf(smt.get_leaf(&key_2), &key_2, &value_2).hash(); let old_value_2 = smt.insert(key_2, value_2); assert_eq!(old_value_2, EMPTY_WORD); @@ -298,11 +323,26 @@ fn test_prospective_hash() { 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 = smt.hash_prospective_leaf(&key_3, &value_3); + let root_3 = { + let prospective = smt.get_prospective_leaf(smt.get_leaf(&key_3), &key_3, &value_3).hash(); let old_value_3 = smt.insert(key_3, value_3); assert_eq!(old_value_3, EMPTY_WORD); @@ -310,41 +350,71 @@ fn test_prospective_hash() { assert_eq!( smt.get_leaf(&key_3), - SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2), (key_3, value_3)]) + SmtLeaf::Multiple(vec![(key_3, value_3), (key_1, value_1), (key_2, value_2)]) ); + + 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)); + assert_eq!( + old_hash, + smt.get_prospective_leaf(smt.get_leaf(&key_3), &key_3, &old_value_3).hash() + ); assert_eq!( 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!( + old_hash, + smt.get_prospective_leaf(smt.get_leaf(&key_2), &key_2, &old_value_2).hash() + ); 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!( + old_hash, + smt.get_prospective_leaf(smt.get_leaf(&key_1), &key_1, &old_value_1).hash() + ); 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 962c893..a36c314 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -110,6 +110,92 @@ pub(crate) trait SparseMerkleTree { old_value } + /// Like [`Self::insert()`], but only performs the insert if 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 the actual root. + /// + /// No mutations are performed if the roots do not 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 = { + let leaf_index: LeafIndex = Self::key_to_leaf_index(&key); + NodeIndex::from(leaf_index) + }; + + let mut new_child_hash = + Self::hash_leaf(&self.get_prospective_leaf(self.get_leaf(&key), &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 = self.get_inner_node(node_index); + 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. fn recompute_nodes_from_index_to_root( @@ -175,14 +261,19 @@ pub(crate) trait SparseMerkleTree { #[allow(dead_code)] fn is_leaf_empty(leaf: &Self::Leaf) -> bool; - /// Returns the hash of a leaf if the leaf WERE inserted into the tree, - /// without performing any insertion or other mutation. + /// Returns what `existing_leaf` would look like if `key` and `value` WERE inserted into the + /// tree, without mutating the tree itself. /// - /// Note: calling this function after actually performing an insert with the same arguments - /// will *not* return the same result, as inserting multiple times with the same key mutates - /// the leaf each time. - #[cfg_attr(not(test), allow(dead_code))] - fn hash_prospective_leaf(&self, key: &Self::Key, value: &Self::Value) -> RpoDigest; + /// `existing_leaf` must have the same index as the key, or the result will be meaningless. To + /// get a prospective leaf based on the current state of the tree, use `self.get_leaf(key)` as + /// the argument for `existing_leaf`. The return value from this function can be chained back + /// into this function as the first argument to continue making prospective changes. + fn get_prospective_leaf( + &self, + existing_leaf: Self::Leaf, + key: &Self::Key, + value: &Self::Value, + ) -> Self::Leaf; /// Maps a key to a leaf index fn key_to_leaf_index(key: &Self::Key) -> LeafIndex; diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index bcc4099..80bc2f1 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -188,6 +188,24 @@ impl SimpleSmt { >::insert(self, key, value) } + /// Like [`Self::insert()`], but only performs the insert if the new tree's root hash would + /// be equal to the hash given in `expected_root`. + /// + /// # Errors + /// Returns [`MerkleError::ConflictingRoots`] with a two-item [alloc::vec::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 the actual root. + /// + /// No mutations are performed if the roots do not 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`. /// @@ -310,8 +328,13 @@ impl SparseMerkleTree for SimpleSmt { *leaf == Self::EMPTY_VALUE } - fn hash_prospective_leaf(&self, _key: &LeafIndex, value: &Word) -> RpoDigest { - Self::hash_leaf(value) + fn get_prospective_leaf( + &self, + _existing_leaf: Word, + _key: &LeafIndex, + value: &Word, + ) -> Word { + *value } fn key_to_leaf_index(key: &LeafIndex) -> LeafIndex {