@ -1,4 +1,4 @@
use alloc ::vec ::Vec ;
use alloc ::{ collections ::BTreeMap , vec ::Vec } ;
use super ::{ EmptySubtreeRoots , InnerNodeInfo , MerkleError , MerklePath , NodeIndex } ;
use crate ::{
@ -45,11 +45,11 @@ pub const SMT_MAX_DEPTH: u8 = 64;
/// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs.
pub ( crate ) trait SparseMerkleTree < const DEPTH : u8 > {
/// The type for a key
type Key : Clone ;
type Key : Clone + Ord ;
/// The type for a value
type Value : Clone + PartialEq ;
/// The type for a leaf
type Leaf ;
type Leaf : Clone ;
/// The type for an opening (i.e. a "proof") of a leaf
type Opening ;
@ -140,6 +140,149 @@ pub(crate) trait SparseMerkleTree {
self . set_root ( node_hash ) ;
}
/// Computes what changes are necessary to insert the specified key-value pairs into this Merkle
/// tree, allowing for validation before applying those changes.
///
/// This method returns a [`MutationSet`], which contains all the information for inserting
/// `kv_pairs` into this Merkle tree already calculated, including the new root hash, which can
/// be queried with [`MutationSet::root()`]. Once a mutation set is returned,
/// [`SparseMerkleTree::apply_mutations()`] can be called in order to commit these changes to
/// the Merkle tree, or [`drop()`] to discard them.
fn compute_mutations (
& self ,
kv_pairs : impl IntoIterator < Item = ( Self ::Key , Self ::Value ) > ,
) -> MutationSet < DEPTH , Self ::Key , Self ::Value > {
use NodeMutation ::* ;
let mut new_root = self . root ( ) ;
let mut new_pairs : BTreeMap < Self ::Key , Self ::Value > = Default ::default ( ) ;
let mut node_mutations : BTreeMap < NodeIndex , NodeMutation > = Default ::default ( ) ;
for ( key , value ) in kv_pairs {
// If the old value and the new value are the same, there is nothing to update.
// For the unusual case that kv_pairs has multiple values at the same key, we'll have
// to check the key-value pairs we've already seen to get the "effective" old value.
let old_value = new_pairs . get ( & key ) . cloned ( ) . unwrap_or_else ( | | self . get_value ( & key ) ) ;
if value = = old_value {
continue ;
}
let leaf_index = Self ::key_to_leaf_index ( & key ) ;
let mut node_index = NodeIndex ::from ( leaf_index ) ;
// We need the current leaf's hash to calculate the new leaf, but in the rare case that
// `kv_pairs` has multiple pairs that go into the same leaf, then those pairs are also
// part of the "current leaf".
let old_leaf = {
let pairs_at_index = new_pairs
. iter ( )
. filter ( | & ( new_key , _ ) | Self ::key_to_leaf_index ( new_key ) = = leaf_index ) ;
pairs_at_index . fold ( self . get_leaf ( & key ) , | acc , ( k , v ) | {
// Most of the time `pairs_at_index` should only contain a single entry (or
// none at all), as multi-leaves should be really rare.
let existing_leaf = acc . clone ( ) ;
self . construct_prospective_leaf ( existing_leaf , k , v )
} )
} ;
let new_leaf = self . construct_prospective_leaf ( old_leaf , & key , & value ) ;
let mut new_child_hash = Self ::hash_leaf ( & new_leaf ) ;
for node_depth in ( 0 . . node_index . depth ( ) ) . rev ( ) {
// Whether the node we're replacing is the right child or the left child.
let is_right = node_index . is_value_odd ( ) ;
node_index . move_up ( ) ;
let old_node = node_mutations
. get ( & node_index )
. map ( | mutation | match mutation {
Addition ( node ) = > node . clone ( ) ,
Removal = > EmptySubtreeRoots ::get_inner_node ( DEPTH , node_depth ) ,
} )
. unwrap_or_else ( | | 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 new node's hash.
new_child_hash = new_node . hash ( ) ;
let & equivalent_empty_hash = EmptySubtreeRoots ::entry ( DEPTH , node_depth ) ;
let is_removal = new_child_hash = = equivalent_empty_hash ;
let new_entry = if is_removal { Removal } else { Addition ( new_node ) } ;
node_mutations . insert ( node_index , new_entry ) ;
}
// Once we're at depth 0, the last node we made is the new root.
new_root = new_child_hash ;
// And then we're done with this pair; on to the next one.
new_pairs . insert ( key , value ) ;
}
MutationSet {
old_root : self . root ( ) ,
new_root ,
node_mutations ,
new_pairs ,
}
}
/// Apply the prospective mutations computed with [`SparseMerkleTree::compute_mutations()`] to
/// this tree.
///
/// # Errors
/// If `mutations` was computed on a tree with a different root than this one, returns
/// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
/// the `mutations` were computed against, and the second item is the actual current root of
/// this tree.
fn apply_mutations (
& mut self ,
mutations : MutationSet < DEPTH , Self ::Key , Self ::Value > ,
) -> Result < ( ) , MerkleError >
where
Self : Sized ,
{
use NodeMutation ::* ;
let MutationSet {
old_root ,
node_mutations ,
new_pairs ,
new_root ,
} = mutations ;
// Guard against accidentally trying to apply mutations that were computed against a
// different tree, including a stale version of this tree.
if old_root ! = self . root ( ) {
return Err ( MerkleError ::ConflictingRoots ( vec ! [ old_root , self . root ( ) ] ) ) ;
}
for ( index , mutation ) in node_mutations {
match mutation {
Removal = > self . remove_inner_node ( index ) ,
Addition ( node ) = > self . insert_inner_node ( index , node ) ,
}
}
for ( key , value ) in new_pairs {
self . insert_value ( key , value ) ;
}
self . set_root ( new_root ) ;
Ok ( ( ) )
}
// REQUIRED METHODS
// ---------------------------------------------------------------------------------------------
@ -161,12 +304,34 @@ 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 < Self ::Value > ;
/// 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 ;
/// Returns the hash of a leaf
fn hash_leaf ( leaf : & Self ::Leaf ) -> RpoDigest ;
/// Returns what a leaf would look like if a key-value pair were inserted into the tree, without
/// mutating the tree itself. The existing leaf can be empty.
///
/// 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.
///
/// # Invariants
/// Because this method is for a prospective key-value insertion into a specific leaf,
/// `existing_leaf` must have the same leaf index as `key` (as determined by
/// [`SparseMerkleTree::key_to_leaf_index()`]), or the result will be meaningless.
fn construct_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 < DEPTH > ;
@ -244,3 +409,50 @@ impl TryFrom for LeafIndex {
Self ::new ( node_index . value ( ) )
}
}
// MUTATIONS
// ================================================================================================
/// A change to an inner node of a [`SparseMerkleTree`] that hasn't yet been applied.
/// [`MutationSet`] stores this type in relation to a [`NodeIndex`] to keep track of what changes
/// need to occur at which node indices.
#[ derive(Debug, Clone, PartialEq, Eq) ]
pub ( crate ) enum NodeMutation {
/// Corresponds to [`SparseMerkleTree::remove_inner_node()`].
Removal ,
/// Corresponds to [`SparseMerkleTree::insert_inner_node()`].
Addition ( InnerNode ) ,
}
/// Represents a group of prospective mutations to a `SparseMerkleTree`, created by
/// `SparseMerkleTree::compute_mutations()`, and that can be applied with
/// `SparseMerkleTree::apply_mutations()`.
#[ derive(Debug, Clone, PartialEq, Eq, Default) ]
pub struct MutationSet < const DEPTH : u8 , K , V > {
/// The root of the Merkle tree this MutationSet is for, recorded at the time
/// [`SparseMerkleTree::compute_mutations()`] was called. Exists to guard against applying
/// mutations to the wrong tree or applying stale mutations to a tree that has since changed.
old_root : RpoDigest ,
/// The set of nodes that need to be removed or added. The "effective" node at an index is the
/// Merkle tree's existing node at that index, with the [`NodeMutation`] in this map at that
/// index overlayed, if any. Each [`NodeMutation::Addition`] corresponds to a
/// [`SparseMerkleTree::insert_inner_node()`] call, and each [`NodeMutation::Removal`]
/// corresponds to a [`SparseMerkleTree::remove_inner_node()`] call.
node_mutations : BTreeMap < NodeIndex , NodeMutation > ,
/// The set of top-level key-value pairs we're prospectively adding to the tree, including
/// adding empty values. The "effective" value for a key is the value in this BTreeMap, falling
/// back to the existing value in the Merkle tree. Each entry corresponds to a
/// [`SparseMerkleTree::insert_value()`] call.
new_pairs : BTreeMap < K , V > ,
/// The calculated root for the Merkle tree, given these mutations. Publicly retrievable with
/// [`MutationSet::root()`]. Corresponds to a [`SparseMerkleTree::set_root()`]. call.
new_root : RpoDigest ,
}
impl < const DEPTH : u8 , K , V > MutationSet < DEPTH , K , V > {
/// Queries the root that was calculated during `SparseMerkleTree::compute_mutations()`. See
/// that method for more information.
pub fn root ( & self ) -> RpoDigest {
self . new_root
}
}
Qyriad
7 months ago
GitHub