From 4f17c1deb15be244d4279fce4d812f9f71ba100d Mon Sep 17 00:00:00 2001 From: Qyriad Date: Thu, 29 Aug 2024 16:30:37 -0600 Subject: [PATCH 1/3] feat(smt): make leaf-emptiness knowable at the trait level --- src/merkle/smt/full/mod.rs | 7 +++++++ src/merkle/smt/mod.rs | 4 ++++ src/merkle/smt/simple/mod.rs | 4 ++++ 3 files changed, 15 insertions(+) diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 2fc774a..e39f3d2 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -263,6 +263,13 @@ impl SparseMerkleTree for Smt { leaf.hash() } + fn is_leaf_empty(leaf: &SmtLeaf) -> bool { + match leaf { + SmtLeaf::Empty(_) => true, + _ => false, + } + } + 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, diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 95a2457..739f4a0 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -167,6 +167,10 @@ pub(crate) trait SparseMerkleTree { /// Returns the hash of a leaf fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest; + /// Returns true if a leaf is completely sparse and does not hold any key-value pairs. + #[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. /// diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 53bc797..dd12566 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -302,6 +302,10 @@ impl SparseMerkleTree for SimpleSmt { leaf.into() } + fn is_leaf_empty(leaf: &Word) -> bool { + *leaf == Self::EMPTY_VALUE + } + fn hash_prospective_leaf(&self, _key: &LeafIndex, value: &Word) -> RpoDigest { Self::hash_leaf(value) } From 4369743b630e03712c6e1df73f6bfcd316877568 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Wed, 21 Aug 2024 13:22:11 -0600 Subject: [PATCH 2/3] feat(smt): export get_value() in the trait --- src/merkle/smt/full/mod.rs | 16 ++++++++++------ src/merkle/smt/mod.rs | 4 ++++ src/merkle/smt/simple/mod.rs | 4 ++++ 3 files changed, 18 insertions(+), 6 deletions(-) diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index e39f3d2..205a443 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -121,12 +121,7 @@ impl Smt { /// Returns the value associated with `key` pub fn get_value(&self, key: &RpoDigest) -> Word { - let leaf_pos = LeafIndex::::from(*key).value(); - - match self.leaves.get(&leaf_pos) { - Some(leaf) => leaf.get_value(key).unwrap_or_default(), - None => EMPTY_WORD, - } + >::get_value(self, key) } /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle @@ -250,6 +245,15 @@ impl SparseMerkleTree for Smt { } } + fn get_value(&self, key: &Self::Key) -> Self::Value { + let leaf_pos = LeafIndex::::from(*key).value(); + + match self.leaves.get(&leaf_pos) { + Some(leaf) => leaf.get_value(key).unwrap_or_default(), + None => EMPTY_WORD, + } + } + fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf { let leaf_pos = LeafIndex::::from(*key).value(); diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 739f4a0..962c893 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -161,6 +161,10 @@ pub(crate) trait SparseMerkleTree { /// 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; + /// Returns the value at the specified key. Recall that by definition, any key that hasn't been + /// updated is associated with [`Self::EMPTY_VALUE`]. + fn get_value(&self, key: &Self::Key) -> Self::Value; + /// Returns the leaf at the specified index. fn get_leaf(&self, key: &Self::Key) -> Self::Leaf; diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index dd12566..bcc4099 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -289,6 +289,10 @@ impl SparseMerkleTree for SimpleSmt { } } + fn get_value(&self, key: &LeafIndex) -> Word { + self.get_leaf(key) + } + fn get_leaf(&self, key: &LeafIndex) -> Word { let leaf_pos = key.value(); match self.leaves.get(&leaf_pos) { From c16b747a73bc271d548b0e1604488983d12dd541 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Thu, 22 Aug 2024 19:00:36 -0600 Subject: [PATCH 3/3] feat(smt): implement root-checked insertion This commit implements SparseMerkleTree::insert_ensure_root(), a version of SparseMerkleTree::insert() that prevents modification of the tree if the insert turns out to not result in the correct tree root (i.e., the transaction is not valid). This is an initial step towards issue 222. As a logical next step, generalizing this to validation of inserting multiple values at once will directly enable removing Merkle tree clones in miden-node InnerState::apply_block(). For further work, as we generalize pre-validation for mutations we will want a separate type to represent an arbitrary set of prospective mutations on the Merkle tree and their validity, like the `ChangeSet` type suggested in the issue. --- CHANGELOG.md | 1 + src/merkle/smt/full/mod.rs | 18 ++++++++ src/merkle/smt/full/tests.rs | 81 ++++++++++++++++++++++++++++----- src/merkle/smt/mod.rs | 86 +++++++++++++++++++++++++++++++++++- src/merkle/smt/simple/mod.rs | 18 ++++++++ 5 files changed, 193 insertions(+), 11 deletions(-) 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..9dd5d62 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 // -------------------------------------------------------------------------------------------- diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 1ee7f9d..dc3faf3 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,23 +259,29 @@ 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 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); @@ -284,10 +290,29 @@ 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 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); @@ -298,10 +323,25 @@ 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 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); @@ -310,14 +350,29 @@ 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)); @@ -325,26 +380,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 962c893..6ebd895 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -110,6 +110,91 @@ 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_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 = 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( @@ -181,7 +266,6 @@ pub(crate) trait SparseMerkleTree { /// 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; /// Maps a key to a leaf index diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index bcc4099..5463fe5 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`. ///