Qyriad
862d995b11
This commit adds a type, MutationSet, which represents a set of changes to a SparseMerkleTree that haven't happened yet, and can be queried on to ensure a set of insertions result in the correct tree root before finalizing and committing the mutation. This is a direct step towards issue 222, and will directly enable removing Merkle tree clones in miden-node InnerState::apply_block(). As part of this change, SparseMerkleTree now requires its Key to be Ord and its Leaf to be Clone (both bounds which were already met by existing implementations). The Ord bound could instead be changed to Eq + Hash, if MutationSet were changed to use a HashMap instead of a BTreeMap. Additionally, as MutationSet is a generic type which works on any type that implements SparseMerkleTree, but is intended for public API use, the SparseMerkleTree trait and InnerNode type have been made public so MutationSet can be used outside of this crate.
539 lines
20 KiB
Rust
539 lines
20 KiB
Rust
use alloc::{collections::BTreeMap, vec::Vec};
|
|
|
|
use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex};
|
|
use crate::{
|
|
hash::rpo::{Rpo256, RpoDigest},
|
|
Felt, Word, EMPTY_WORD,
|
|
};
|
|
|
|
mod full;
|
|
pub use full::{Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError, SMT_DEPTH};
|
|
|
|
mod simple;
|
|
pub use simple::SimpleSmt;
|
|
|
|
// CONSTANTS
|
|
// ================================================================================================
|
|
|
|
/// Minimum supported depth.
|
|
pub const SMT_MIN_DEPTH: u8 = 1;
|
|
|
|
/// Maximum supported depth.
|
|
pub const SMT_MAX_DEPTH: u8 = 64;
|
|
|
|
// SPARSE MERKLE TREE
|
|
// ================================================================================================
|
|
|
|
/// An abstract description of a sparse Merkle tree.
|
|
///
|
|
/// A sparse Merkle tree is a key-value map which also supports proving that a given value is indeed
|
|
/// stored at a given key in the tree. It is viewed as always being fully populated. If a leaf's
|
|
/// value was not explicitly set, then its value is the default value. Typically, the vast majority
|
|
/// of leaves will store the default value (hence it is "sparse"), and therefore the internal
|
|
/// representation of the tree will only keep track of the leaves that have a different value from
|
|
/// the default.
|
|
///
|
|
/// All leaves sit at the same depth. The deeper the tree, the more leaves it has; but also the
|
|
/// longer its proofs are - of exactly `log(depth)` size. A tree cannot have depth 0, since such a
|
|
/// tree is just a single value, and is probably a programming mistake.
|
|
///
|
|
/// Every key maps to one leaf. If there are as many keys as there are leaves, then
|
|
/// [Self::Leaf] should be the same type as [Self::Value], as is the case with
|
|
/// [crate::merkle::SimpleSmt]. However, if there are more keys than leaves, then [`Self::Leaf`]
|
|
/// must accomodate all keys that map to the same leaf.
|
|
///
|
|
/// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs.
|
|
pub trait SparseMerkleTree<const DEPTH: u8> {
|
|
/// The type for a key
|
|
type Key: Clone + Ord;
|
|
/// The type for a value
|
|
type Value: Clone + PartialEq;
|
|
/// The type for a leaf
|
|
type Leaf: Clone;
|
|
/// The type for an opening (i.e. a "proof") of a leaf
|
|
type Opening;
|
|
|
|
/// The default value used to compute the hash of empty leaves
|
|
const EMPTY_VALUE: Self::Value;
|
|
|
|
// PROVIDED METHODS
|
|
// ---------------------------------------------------------------------------------------------
|
|
|
|
/// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
|
|
/// path to the leaf, as well as the leaf itself.
|
|
fn open(&self, key: &Self::Key) -> Self::Opening {
|
|
let leaf = self.get_leaf(key);
|
|
|
|
let mut index: NodeIndex = {
|
|
let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(key);
|
|
leaf_index.into()
|
|
};
|
|
|
|
let merkle_path = {
|
|
let mut path = Vec::with_capacity(index.depth() as usize);
|
|
for _ in 0..index.depth() {
|
|
let is_right = index.is_value_odd();
|
|
index.move_up();
|
|
let InnerNode { left, right } = self.get_inner_node(index);
|
|
let value = if is_right { left } else { right };
|
|
path.push(value);
|
|
}
|
|
|
|
MerklePath::new(path)
|
|
};
|
|
|
|
Self::path_and_leaf_to_opening(merkle_path, leaf)
|
|
}
|
|
|
|
/// Inserts a value at the specified key, returning the previous value associated with that key.
|
|
/// Recall that by definition, any key that hasn't been updated is associated with
|
|
/// [`Self::EMPTY_VALUE`].
|
|
///
|
|
/// This also recomputes all hashes between the leaf (associated with the key) and the root,
|
|
/// updating the root itself.
|
|
fn insert(&mut self, key: Self::Key, value: Self::Value) -> Self::Value {
|
|
let old_value = self.insert_value(key.clone(), value.clone()).unwrap_or(Self::EMPTY_VALUE);
|
|
|
|
// if the old value and new value are the same, there is nothing to update
|
|
if value == old_value {
|
|
return value;
|
|
}
|
|
|
|
let leaf = self.get_leaf(&key);
|
|
let node_index = {
|
|
let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(&key);
|
|
leaf_index.into()
|
|
};
|
|
|
|
self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf));
|
|
|
|
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(
|
|
&mut self,
|
|
mut index: NodeIndex,
|
|
node_hash_at_index: RpoDigest,
|
|
) {
|
|
let mut node_hash = node_hash_at_index;
|
|
for node_depth in (0..index.depth()).rev() {
|
|
let is_right = index.is_value_odd();
|
|
index.move_up();
|
|
let InnerNode { left, right } = self.get_inner_node(index);
|
|
let (left, right) = if is_right {
|
|
(left, node_hash)
|
|
} else {
|
|
(node_hash, right)
|
|
};
|
|
node_hash = Rpo256::merge(&[left, right]);
|
|
|
|
if node_hash == *EmptySubtreeRoots::entry(DEPTH, node_depth) {
|
|
// If a subtree is empty, when can remove the inner node, since it's equal to the
|
|
// default value
|
|
self.remove_inner_node(index)
|
|
} else {
|
|
self.insert_inner_node(index, InnerNode { left, right });
|
|
}
|
|
}
|
|
self.set_root(node_hash);
|
|
}
|
|
|
|
/// Start a prospective mutation transaction, which can be queried, discarded, or applied.
|
|
///
|
|
/// This method returns a [`MutationSet`], which you can call [`MutationSet::insert()`] on to
|
|
/// perform prospective mutations to this Merkle tree, and check the computed root hash given
|
|
/// those mutations with [`MutationSet::root()`]. When you are done, call
|
|
/// [`SparseMerkleTree::commit()`] with the return value of [`MutationSet::done()`] to apply the
|
|
/// prospective mutations to this tree. Or to discard the changes, simply [`drop`] the
|
|
/// [`MutationSet`].
|
|
///
|
|
/// ## Example
|
|
/// ```
|
|
/// # use miden_crypto::{hash::rpo::RpoDigest, Felt};
|
|
/// # use miden_crypto::merkle::{Smt, EmptySubtreeRoots, SMT_DEPTH};
|
|
/// # let mut smt = Smt::default();
|
|
/// let mut mutations = smt.mutate();
|
|
/// mutations.insert(Default::default(), Default::default());
|
|
/// assert_eq!(mutations.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
|
|
/// smt.commit(mutations.done());
|
|
/// assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
|
|
/// ```
|
|
fn mutate(&self) -> MutationSet<DEPTH, Self>
|
|
where
|
|
Self: Sized,
|
|
{
|
|
MutationSet {
|
|
tree: self,
|
|
node_mutations: Default::default(),
|
|
new_pairs: Default::default(),
|
|
new_root: self.root(),
|
|
}
|
|
}
|
|
|
|
/// Apply prospective mutations started with [`SparseMerkleTree::mutate()`] to this tree.
|
|
///
|
|
/// This method takes the return value of [`MutationSet::done()`], and applies the changes
|
|
/// represented in that [`MutationSet`] to this Merkle tree. See [`SparseMerkleTree::mutate()`]
|
|
/// for more details.
|
|
fn commit(&mut self, mutations: DoneMutationSet<DEPTH, Self::Key, Self::Value>)
|
|
where
|
|
Self: Sized,
|
|
{
|
|
use NodeMutation::*;
|
|
let DoneMutationSet { node_mutations, new_pairs, new_root } = mutations;
|
|
|
|
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);
|
|
}
|
|
|
|
// REQUIRED METHODS
|
|
// ---------------------------------------------------------------------------------------------
|
|
|
|
/// The root of the tree
|
|
fn root(&self) -> RpoDigest;
|
|
|
|
/// Sets the root of the tree
|
|
fn set_root(&mut self, root: RpoDigest);
|
|
|
|
/// Retrieves an inner node at the given index
|
|
fn get_inner_node(&self, index: NodeIndex) -> InnerNode;
|
|
|
|
/// Inserts an inner node at the given index
|
|
fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode);
|
|
|
|
/// Removes an inner node at the given index
|
|
fn remove_inner_node(&mut self, index: NodeIndex);
|
|
|
|
/// 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 `existing_leaf` would look like if `key` and `value` WERE inserted into the
|
|
/// tree, without mutating the tree itself.
|
|
///
|
|
/// `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<DEPTH>;
|
|
|
|
/// Maps a (MerklePath, Self::Leaf) to an opening.
|
|
///
|
|
/// The length `path` is guaranteed to be equal to `DEPTH`
|
|
fn path_and_leaf_to_opening(path: MerklePath, leaf: Self::Leaf) -> Self::Opening;
|
|
}
|
|
|
|
// INNER NODE
|
|
// ================================================================================================
|
|
|
|
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
|
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
|
|
pub struct InnerNode {
|
|
pub left: RpoDigest,
|
|
pub right: RpoDigest,
|
|
}
|
|
|
|
impl InnerNode {
|
|
pub fn hash(&self) -> RpoDigest {
|
|
Rpo256::merge(&[self.left, self.right])
|
|
}
|
|
}
|
|
|
|
// LEAF INDEX
|
|
// ================================================================================================
|
|
|
|
/// The index of a leaf, at a depth known at compile-time.
|
|
#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
|
|
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
|
|
pub struct LeafIndex<const DEPTH: u8> {
|
|
index: NodeIndex,
|
|
}
|
|
|
|
impl<const DEPTH: u8> LeafIndex<DEPTH> {
|
|
pub fn new(value: u64) -> Result<Self, MerkleError> {
|
|
if DEPTH < SMT_MIN_DEPTH {
|
|
return Err(MerkleError::DepthTooSmall(DEPTH));
|
|
}
|
|
|
|
Ok(LeafIndex { index: NodeIndex::new(DEPTH, value)? })
|
|
}
|
|
|
|
pub fn value(&self) -> u64 {
|
|
self.index.value()
|
|
}
|
|
}
|
|
|
|
impl LeafIndex<SMT_MAX_DEPTH> {
|
|
pub const fn new_max_depth(value: u64) -> Self {
|
|
LeafIndex {
|
|
index: NodeIndex::new_unchecked(SMT_MAX_DEPTH, value),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<const DEPTH: u8> From<LeafIndex<DEPTH>> for NodeIndex {
|
|
fn from(value: LeafIndex<DEPTH>) -> Self {
|
|
value.index
|
|
}
|
|
}
|
|
|
|
impl<const DEPTH: u8> TryFrom<NodeIndex> for LeafIndex<DEPTH> {
|
|
type Error = MerkleError;
|
|
|
|
fn try_from(node_index: NodeIndex) -> Result<Self, Self::Error> {
|
|
if node_index.depth() != DEPTH {
|
|
return Err(MerkleError::InvalidDepth {
|
|
expected: DEPTH,
|
|
provided: node_index.depth(),
|
|
});
|
|
}
|
|
|
|
Self::new(node_index.value())
|
|
}
|
|
}
|
|
|
|
// MUTATIONS
|
|
// ================================================================================================
|
|
|
|
/// A change to an inner node of a [`SparseMerkleTree`] that hasn't been applied yet.
|
|
/// [`MutationSet`] stores this type in relation to a [`NodeIndex`] to keep track of what changes
|
|
/// need to occur at what node indices.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub(crate) enum NodeMutation {
|
|
Removal,
|
|
Addition(InnerNode),
|
|
}
|
|
|
|
/// Represents prospective mutations to a [`SparseMerkleTree`], created by
|
|
/// [`SparseMerkleTree::mutate()`], and applied with [`SparseMerkleTree::commit()`].
|
|
///
|
|
/// `T` is the Merkle tree type this is prospectively mutating.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub struct MutationSet<'t, const DEPTH: u8, T>
|
|
where
|
|
T: SparseMerkleTree<DEPTH>,
|
|
{
|
|
/// A reference to the tree we're mutating. We need to be able to query the tree's existing
|
|
/// nodes and leaves to calculate prospective nodes, leaves, and roots.
|
|
tree: &'t T,
|
|
/// 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.
|
|
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.
|
|
new_pairs: BTreeMap<T::Key, T::Value>,
|
|
/// The currently-calculated root for the Merkle tree, given these mutations. When this struct
|
|
/// is constructed, `new_root` is set to the existing root, since there aren't any changes yet.
|
|
new_root: RpoDigest,
|
|
}
|
|
|
|
impl<'t, const DEPTH: u8, T> MutationSet<'t, DEPTH, T>
|
|
where
|
|
T: SparseMerkleTree<DEPTH>,
|
|
{
|
|
/// Pass the return value of this method to [`SparseMerkleTree::commit()`] to apply these
|
|
/// mutations to the Merkle tree.
|
|
pub fn done(self) -> DoneMutationSet<DEPTH, T::Key, T::Value> {
|
|
// Get rid of the `tree` reference, so we can pass the mutation info to a `&mut self` method
|
|
// on `SparseMerkleTree`.
|
|
let Self { node_mutations, new_pairs, new_root, .. } = self;
|
|
DoneMutationSet { node_mutations, new_pairs, new_root }
|
|
}
|
|
|
|
/// Get the prospective root hash for if these mutations were applied to the Merkle tree.
|
|
pub fn root(&self) -> RpoDigest {
|
|
self.new_root
|
|
}
|
|
|
|
/// Add a new key-value pair to the set of prospective changes to the Merkle tree.
|
|
///
|
|
/// After `insert()`, you may call [`MutationSet::root()`] to query what the new calculated root
|
|
/// hash of the Merkle tree would be after all requested insertions. Recall that semantically
|
|
/// removing a value is done by inserting [`SparseMerkleTree::EMPTY_VALUE`].
|
|
pub fn insert(&mut self, key: T::Key, value: T::Value) {
|
|
// This function's calculations are eager, so multiple inserts that affect the same nodes
|
|
// will calculate those node's hashes each time. Future work could lazily calculate node
|
|
// hashes all at once instead.
|
|
|
|
// If the old value and the new value are the same, there is nothing to update.
|
|
let old_value = self.get_effective_value(&key);
|
|
if value == old_value {
|
|
return;
|
|
}
|
|
|
|
let mut node_index: NodeIndex = T::key_to_leaf_index(&key).into();
|
|
|
|
let old_leaf = self.get_effective_leaf(&key);
|
|
let new_leaf = self.tree.get_prospective_leaf(old_leaf, &key, &value);
|
|
|
|
let mut new_child_hash = T::hash_leaf(&new_leaf);
|
|
|
|
for node_depth in (0..node_index.depth()).rev() {
|
|
// Whether the node we're replacing is the right child or left child.
|
|
let is_right = node_index.is_value_odd();
|
|
node_index.move_up();
|
|
|
|
let old_node = self.get_effective_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);
|
|
if new_child_hash == equivalent_empty_hash {
|
|
self.remove_inner_node(node_index);
|
|
} else {
|
|
self.insert_inner_node(node_index, new_node);
|
|
}
|
|
}
|
|
|
|
// Once we're at depth 0, the last node we made is the new root.
|
|
self.new_root = new_child_hash;
|
|
self.new_pairs.insert(key, value);
|
|
}
|
|
|
|
// Implementation details.
|
|
|
|
fn get_effective_value(&self, key: &T::Key) -> T::Value {
|
|
self.new_pairs.get(key).cloned().unwrap_or_else(|| self.tree.get_value(key))
|
|
}
|
|
|
|
fn get_effective_leaf(&self, key: &T::Key) -> T::Leaf {
|
|
let leaf_key = key;
|
|
let leaf_index = T::key_to_leaf_index(leaf_key);
|
|
let pairs_at_index = self
|
|
.new_pairs
|
|
.iter()
|
|
.filter(|&(new_key, _)| T::key_to_leaf_index(new_key) == leaf_index);
|
|
|
|
let initial_acc = self.tree.get_leaf(key);
|
|
pairs_at_index.fold(initial_acc, |acc, (k, v)| {
|
|
// In practice this should only run once (if that), most of the time.
|
|
let existing_leaf = acc.clone();
|
|
self.tree.get_prospective_leaf(existing_leaf, k, v)
|
|
})
|
|
}
|
|
|
|
fn get_effective_node(&self, index: NodeIndex) -> InnerNode {
|
|
use NodeMutation::*;
|
|
|
|
self.node_mutations
|
|
.get(&index)
|
|
.map(|mutation| match mutation {
|
|
Addition(node) => node.clone(),
|
|
Removal => {
|
|
let &child = EmptySubtreeRoots::entry(DEPTH, index.depth() + 1);
|
|
InnerNode { left: child, right: child }
|
|
},
|
|
})
|
|
.unwrap_or_else(|| self.tree.get_inner_node(index))
|
|
}
|
|
|
|
fn remove_inner_node(&mut self, index: NodeIndex) {
|
|
use alloc::collections::btree_map::Entry::*;
|
|
use NodeMutation::*;
|
|
|
|
match self.node_mutations.entry(index) {
|
|
Vacant(entry) => {
|
|
entry.insert(Removal);
|
|
},
|
|
Occupied(mut entry) => match entry.get_mut() {
|
|
// If we have an addition with this index, we don't care about what value it has, we
|
|
// just need to convert it to a removal instead.
|
|
Addition(_) => {
|
|
entry.insert(Removal);
|
|
},
|
|
Removal => (),
|
|
},
|
|
}
|
|
}
|
|
|
|
fn insert_inner_node(&mut self, index: NodeIndex, new_node: InnerNode) {
|
|
use alloc::collections::btree_map::Entry::*;
|
|
use NodeMutation::*;
|
|
|
|
match self.node_mutations.entry(index) {
|
|
Vacant(entry) => {
|
|
entry.insert(Addition(new_node));
|
|
},
|
|
Occupied(mut entry) => match entry.get_mut() {
|
|
Addition(existing) => {
|
|
// If there's an existing addition with this key, then overwrite it to be an
|
|
// addition of this new node instead.
|
|
*existing = new_node;
|
|
},
|
|
Removal => {
|
|
// Likewise a removal of this key gets overwritten with an addition.
|
|
entry.insert(Addition(new_node));
|
|
},
|
|
},
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Helper type that represents a [`MutationSet`] that's ready to be applied to a [`SparseMerkleTree`].
|
|
///
|
|
/// It is created by [`MutationSet::done()`], which should be directly passed to
|
|
/// [`SparseMerkleTree::commit()`]:
|
|
/// ```
|
|
/// # use miden_crypto::{hash::rpo::RpoDigest, merkle::Smt, Felt};
|
|
/// # let mut smt = Smt::default();
|
|
/// # let key = RpoDigest::default();
|
|
/// # let value: [Felt; 4] = Default::default();
|
|
/// let mut mutations = smt.mutate();
|
|
/// mutations.insert(key, value);
|
|
/// smt.commit(mutations.done());
|
|
/// ```
|
|
// This type exists for the sake of the borrow checker -- SparseMerkleTree::commit() needs a
|
|
// mutable reference to `self`, but MutationSet stores a shared reference to the SparseMerkleTree,
|
|
// and those can't both exist at the same time. By going through this type first, which stores all
|
|
// the same data as MutationSet except for the shared reference to SparseMerkleTree, the shared
|
|
// reference is dropped just before SparseMerkleTree::commit() is called, and the borrow checker is
|
|
// happy. Interior mutability would also work, but then we would lose static lifetime checking.
|
|
pub struct DoneMutationSet<const DEPTH: u8, K, V> {
|
|
node_mutations: BTreeMap<NodeIndex, NodeMutation>,
|
|
new_pairs: BTreeMap<K, V>,
|
|
new_root: RpoDigest,
|
|
}
|