From 52ab7bca97bb78e2c22d52c5edf71220f49b92e4 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Mon, 12 Aug 2024 15:56:07 -0600 Subject: [PATCH 1/3] feat: impl hashing Merkle leaves that don't yet exist This commit implements 'prospective leaf hashing' -- computing what the hash of a sparse Merkle tree leaf *would* be for a key-value insertion without actually performing that insertion. For SimpleSmt, this is trivial, since the leaf hash and its payload are the same. For the full Smt, the new leaf payload (and thus, its hash) depend on the existing payload in that leaf, making the prospective hash logic a combination of the normal insertion logic and the normal hash logic. But because we're only interested in the hash and not the intermediate value, we can skip allocations and sorts for the payload itself. --- src/merkle/smt/full/leaf.rs | 4 +- src/merkle/smt/full/mod.rs | 87 ++++++++++++++++++++++++++++++++++ src/merkle/smt/full/tests.rs | 92 +++++++++++++++++++++++++++++++++++- src/merkle/smt/mod.rs | 8 ++++ src/merkle/smt/simple/mod.rs | 4 ++ 5 files changed, 192 insertions(+), 3 deletions(-) diff --git a/src/merkle/smt/full/leaf.rs b/src/merkle/smt/full/leaf.rs index 095a4fb..585fc40 100644 --- a/src/merkle/smt/full/leaf.rs +++ b/src/merkle/smt/full/leaf.rs @@ -350,7 +350,7 @@ impl Deserializable for SmtLeaf { // ================================================================================================ /// Converts a key-value tuple to an iterator of `Felt`s -fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator { +pub(crate) fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator { let key_elements = key.into_iter(); let value_elements = value.into_iter(); @@ -359,7 +359,7 @@ fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator /// Compares two keys, compared element-by-element using their integer representations starting with /// the most significant element. -fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering { +pub(crate) fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering { for (v1, v2) in key_1.iter().zip(key_2.iter()).rev() { let v1 = v1.as_int(); let v2 = v2.as_int(); diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 5cd510b..626be36 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -3,6 +3,7 @@ use alloc::{ string::ToString, vec::Vec, }; +use core::iter; use super::{ EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, MerklePath, @@ -263,6 +264,92 @@ impl SparseMerkleTree for Smt { leaf.hash() } + fn hash_prospective_leaf(&self, key: &RpoDigest, value: &Word) -> RpoDigest { + // This function combines logic from SmtLeaf::insert() and SmtLeaf::hash() to determine what + // the hash of a leaf would be with the `(key, value)` pair inserted into it, without simply + // cloning the leaf which could be expensive for some leaves, and is easily avoidable when + // we can combine the insertion and hashing operations. + let new_pair = (*key, *value); + let is_removal: bool = *value == EMPTY_WORD; + + let leaf_index: LeafIndex = Self::key_to_leaf_index(key); + match self.leaves.get(&leaf_index.value()) { + // If this key doesn't have a value, our job is very simple. + None => SmtLeaf::Single(new_pair).hash(), + + // If this key already has a value, then the hash will be based off a prospective + // mutation on the leaf. + Some(existing_leaf) => match existing_leaf { + // Inserting an empty value into an empty leaf or a single leaf both do the same + // thing. + SmtLeaf::Empty(_) | SmtLeaf::Single(_) if is_removal => { + SmtLeaf::new_empty(key.into()).hash() + }, + + SmtLeaf::Empty(_) => SmtLeaf::Single(new_pair).hash(), + + SmtLeaf::Single(pair) => { + if pair.0 == *key { + SmtLeaf::Single(new_pair).hash() + } else { + // Inserting a non-empty value into a new key would change this to a + // multi-leaf. + // TODO: mini-optimization: use an array with each key's and value's Felts + // flattened inline to avoid the Vec allocation. + let elements: Vec = + [*pair, new_pair].into_iter().flat_map(leaf::kv_to_elements).collect(); + + Rpo256::hash_elements(&elements) + } + }, + + SmtLeaf::Multiple(pairs) => { + match pairs.binary_search_by(|(cur_key, _)| leaf::cmp_keys(*cur_key, *key)) { + Ok(pos) => { + if is_removal && pairs.len() == 2 { + // This removal would convert this Multi into a Single, so we can + // just stop here. + return SmtLeaf::Single(pairs[0]).hash(); + } + + let (before_pos, rest) = pairs.split_at(pos); + let with_pos_removed = rest.iter().copied().skip(1); + let middle = if !is_removal { Some(new_pair) } else { None }; + let elements: Vec = before_pos + .iter() + .copied() + .chain(middle) + .chain(with_pos_removed) + .flat_map(leaf::kv_to_elements) + .collect(); + + Rpo256::hash_elements(&elements) + }, + Err(pos_for_insert) => { + if is_removal { + // The only values are at other keys, so we just hash the leaf + // as-is. + return existing_leaf.hash(); + } + + let (before_pos, rest) = pairs.split_at(pos_for_insert); + let middle = iter::once(new_pair); + let elements: Vec = before_pos + .iter() + .copied() + .chain(middle) + .chain(rest.iter().copied()) + .flat_map(leaf::kv_to_elements) + .collect(); + + Rpo256::hash_elements(&elements) + }, + } + }, + }, + } + } + fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex { let most_significant_felt = key[3]; LeafIndex::new_max_depth(most_significant_felt.as_int()) diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 66cd203..1ee7f9d 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::{EmptySubtreeRoots, MerkleStore}, + merkle::{smt::SparseMerkleTree, EmptySubtreeRoots, MerkleStore}, utils::{Deserializable, Serializable}, Word, ONE, WORD_SIZE, }; @@ -258,6 +258,96 @@ fn test_smt_removal() { } } +#[test] +fn test_prospective_hash() { + let mut smt = Smt::default(); + + 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)]); + let key_3: RpoDigest = + RpoDigest::from([3_u32.into(), 3_u32.into(), 3_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]; + + // insert key-value 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!(prospective, smt.get_leaf(&key_1).hash()); + + assert_eq!(smt.get_leaf(&key_1), SmtLeaf::Single((key_1, value_1))); + + } + + // insert key-value 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!(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)]) + ); + } + + // insert key-value 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!(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)]) + ); + } + + // remove key 3 + { + let old_hash = smt.get_leaf(&key_3).hash(); + let old_value_3 = smt.insert(key_3, EMPTY_WORD); + assert_eq!(old_value_3, value_3); + assert_eq!(old_hash, smt.hash_prospective_leaf(&key_3, &old_value_3)); + + assert_eq!( + smt.get_leaf(&key_3), + SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)]) + ); + } + + // remove key 2 + { + let old_hash = smt.get_leaf(&key_2).hash(); + let old_value_2 = smt.insert(key_2, EMPTY_WORD); + 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))); + } + + // remove key 1 + { + let old_hash = smt.get_leaf(&key_1).hash(); + let old_value_1 = smt.insert(key_1, EMPTY_WORD); + 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())); + } +} + /// Tests that 2 key-value pairs stored in the same leaf have the same path #[test] fn test_smt_path_to_keys_in_same_leaf_are_equal() { diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index d7d42da..1c71204 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -167,6 +167,14 @@ pub(crate) trait SparseMerkleTree { /// Returns the hash of a leaf fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest; + /// Returns the hash of a leaf if the leaf WERE inserted into the tree, + /// without performing any insertion or other mutation. + /// + /// 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. + fn hash_prospective_leaf(&self, key: &Self::Key, value: &Self::Value) -> RpoDigest; + /// 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 f1ff0dc..53bc797 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 hash_prospective_leaf(&self, _key: &LeafIndex, value: &Word) -> RpoDigest { + Self::hash_leaf(value) + } + fn key_to_leaf_index(key: &LeafIndex) -> LeafIndex { *key } From 00e2be79a40d625b3bbb8ed34488c0fb87ce25ee Mon Sep 17 00:00:00 2001 From: Qyriad Date: Wed, 21 Aug 2024 13:22:11 -0600 Subject: [PATCH 2/3] feat: export Merkle get_value() in the trait --- src/merkle/smt/full/mod.rs | 16 ++++++++++------ src/merkle/smt/mod.rs | 4 ++++ src/merkle/smt/simple/mod.rs | 8 ++++++++ 3 files changed, 22 insertions(+), 6 deletions(-) diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 626be36..396ee42 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -122,12 +122,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 @@ -251,6 +246,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 1c71204..9ed0155 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 53bc797..a6cc7eb 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -289,6 +289,14 @@ impl SparseMerkleTree for SimpleSmt { } } + fn get_value(&self, key: &Self::Key) -> Self::Value { + let leaf_pos = key.value(); + match self.leaves.get(&leaf_pos) { + Some(word) => *word, + None => Self::EMPTY_VALUE, + } + } + fn get_leaf(&self, key: &LeafIndex) -> Word { let leaf_pos = key.value(); match self.leaves.get(&leaf_pos) { From 1795ef4ba39ffbb99de94020c4bb24ec20a81ebd Mon Sep 17 00:00:00 2001 From: Qyriad Date: Wed, 21 Aug 2024 14:49:47 -0600 Subject: [PATCH 3/3] WIP: smt: implement root-checked insertion --- src/merkle/smt/full/mod.rs | 18 +++++++ src/merkle/smt/full/tests.rs | 81 +++++++++++++++++++++++++---- src/merkle/smt/mod.rs | 98 ++++++++++++++++++++++++++++++++++++ src/merkle/smt/simple/mod.rs | 19 +++++++ 4 files changed, 206 insertions(+), 10 deletions(-) diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 396ee42..40819e9 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -168,6 +168,24 @@ 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 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 9ed0155..ccddb36 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -110,6 +110,104 @@ 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. fn recompute_nodes_from_index_to_root( diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index a6cc7eb..f44f106 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -188,6 +188,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`. ///