966 lines
38 KiB
Rust
966 lines
38 KiB
Rust
use alloc::{collections::BTreeMap, vec::Vec};
|
|
use core::mem;
|
|
|
|
use num::Integer;
|
|
|
|
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(crate) 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;
|
|
|
|
/// The root of the empty tree with provided DEPTH
|
|
const EMPTY_ROOT: RpoDigest;
|
|
|
|
// 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);
|
|
}
|
|
|
|
/// 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
|
|
// ---------------------------------------------------------------------------------------------
|
|
|
|
/// 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 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>;
|
|
|
|
/// Constructs a single leaf from an arbitrary amount of key-value pairs.
|
|
/// Those pairs must all have the same leaf index.
|
|
fn pairs_to_leaf(pairs: Vec<(Self::Key, Self::Value)>) -> Self::Leaf;
|
|
|
|
/// 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;
|
|
|
|
#[cfg_attr(not(test), allow(dead_code))]
|
|
fn sorted_pairs_to_leaves(
|
|
pairs: Vec<(Self::Key, Self::Value)>,
|
|
) -> PairComputations<u64, Self::Leaf> {
|
|
let mut accumulator: PairComputations<u64, Self::Leaf> = Default::default();
|
|
|
|
// The kv-pairs we've seen so far that correspond to a single leaf.
|
|
let mut current_leaf_buffer: Vec<(Self::Key, Self::Value)> = Default::default();
|
|
|
|
let mut iter = pairs.into_iter().peekable();
|
|
while let Some((key, value)) = iter.next() {
|
|
let col = Self::key_to_leaf_index(&key).index.value();
|
|
let peeked_col = iter.peek().map(|(key, _v)| {
|
|
let index = Self::key_to_leaf_index(key);
|
|
let next_col = index.index.value();
|
|
// We panic if `pairs` is not sorted by column.
|
|
debug_assert!(next_col >= col);
|
|
next_col
|
|
});
|
|
current_leaf_buffer.push((key, value));
|
|
|
|
// If the next pair is the same column as this one, then we're done after adding this
|
|
// pair to the buffer.
|
|
if peeked_col == Some(col) {
|
|
continue;
|
|
}
|
|
|
|
// Otherwise, the next pair is a different column, or there is no next pair. Either way
|
|
// it's time to swap out our buffer.
|
|
let leaf_pairs = mem::take(&mut current_leaf_buffer);
|
|
let leaf = Self::pairs_to_leaf(leaf_pairs);
|
|
let hash = Self::hash_leaf(&leaf);
|
|
|
|
accumulator.nodes.insert(col, leaf);
|
|
add_subtree_leaf(&mut accumulator.leaves, SubtreeLeaf { col, hash });
|
|
|
|
debug_assert!(current_leaf_buffer.is_empty());
|
|
}
|
|
|
|
accumulator
|
|
}
|
|
|
|
/// Builds Merkle nodes from a bottom layer of tuples of horizontal indices and their hashes,
|
|
/// sorted by their position.
|
|
///
|
|
/// The leaves are 'conceptual' leaves, simply being entities at the bottom of some subtree, not
|
|
/// [`Self::Leaf`].
|
|
///
|
|
/// # Panics
|
|
/// With debug assertions on, this function panics under invalid inputs: if `leaves` contains
|
|
/// more entries than can fit in a depth-8 subtree (more than 256), if `bottom_depth` is
|
|
/// lower in the tree than the specified maximum depth (`DEPTH`), or if `leaves` is not sorted.
|
|
// FIXME: more complete docstring.
|
|
fn build_subtree(
|
|
mut leaves: Vec<SubtreeLeaf>,
|
|
bottom_depth: u8,
|
|
) -> (BTreeMap<NodeIndex, InnerNode>, Vec<SubtreeLeaf>) {
|
|
debug_assert!(bottom_depth <= DEPTH);
|
|
debug_assert!(Integer::is_multiple_of(&bottom_depth, &8));
|
|
debug_assert!(leaves.len() <= usize::pow(2, 8));
|
|
|
|
let subtree_root = bottom_depth - 8;
|
|
|
|
let mut inner_nodes: BTreeMap<NodeIndex, InnerNode> = Default::default();
|
|
|
|
let mut next_leaves: Vec<SubtreeLeaf> = Vec::with_capacity(leaves.len() / 2);
|
|
|
|
for next_depth in (subtree_root..bottom_depth).rev() {
|
|
debug_assert!(next_depth <= bottom_depth);
|
|
|
|
// `next_depth` is the stuff we're making.
|
|
// `current_depth` is the stuff we have.
|
|
let current_depth = next_depth + 1;
|
|
|
|
let mut iter = leaves.drain(..).peekable();
|
|
while let Some(first) = iter.next() {
|
|
// On non-continuous iterations, including the first iteration, `first_column` may
|
|
// be a left or right node. On subsequent continuous iterations, we will always call
|
|
// `iter.next()` twice.
|
|
|
|
// On non-continuous iterations (including the very first iteration), this column
|
|
// could be either on the left or the right. If the next iteration is not
|
|
// discontinuous with our right node, then the next iteration's
|
|
|
|
let is_right = first.col.is_odd();
|
|
let (left, right) = if is_right {
|
|
// Discontinuous iteration: we have no left node, so it must be empty.
|
|
|
|
let left = SubtreeLeaf {
|
|
col: first.col - 1,
|
|
hash: *EmptySubtreeRoots::entry(DEPTH, current_depth),
|
|
};
|
|
let right = first;
|
|
|
|
(left, right)
|
|
} else {
|
|
let left = first;
|
|
|
|
let right_col = first.col + 1;
|
|
let right = match iter.peek().copied() {
|
|
Some(SubtreeLeaf { col, .. }) if col == right_col => {
|
|
// Our inputs must be sorted.
|
|
debug_assert!(left.col <= col);
|
|
// The next leaf in the iterator is our sibling. Use it and consume it!
|
|
iter.next().unwrap()
|
|
},
|
|
// Otherwise, the leaves don't contain our sibling, so our sibling must be
|
|
// empty.
|
|
_ => SubtreeLeaf {
|
|
col: right_col,
|
|
hash: *EmptySubtreeRoots::entry(DEPTH, current_depth),
|
|
},
|
|
};
|
|
|
|
(left, right)
|
|
};
|
|
|
|
let index = NodeIndex::new_unchecked(current_depth, left.col).parent();
|
|
let node = InnerNode { left: left.hash, right: right.hash };
|
|
let hash = node.hash();
|
|
|
|
let &equivalent_empty_hash = EmptySubtreeRoots::entry(DEPTH, next_depth);
|
|
// If this hash is empty, then it doesn't become a new inner node, nor does it count
|
|
// as a leaf for the next depth.
|
|
if hash != equivalent_empty_hash {
|
|
inner_nodes.insert(index, node);
|
|
// FIXME: is it possible for this to end up not being sorted? I don't think so.
|
|
next_leaves.push(SubtreeLeaf { col: index.value(), hash });
|
|
}
|
|
}
|
|
|
|
// Stop borrowing `leaves`, so we can swap it.
|
|
// The iterator is empty at this point anyway.
|
|
drop(iter);
|
|
|
|
// After each depth, consider the stuff we just made the new "leaves", and empty the
|
|
// other collection.
|
|
mem::swap(&mut leaves, &mut next_leaves);
|
|
}
|
|
|
|
(inner_nodes, leaves)
|
|
}
|
|
}
|
|
|
|
// 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 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
|
|
}
|
|
}
|
|
|
|
// HELPERS
|
|
// ================================================================================================
|
|
|
|
/// A depth-8 subtree contains 256 "columns" that can possibly be occupied.
|
|
const COLS_PER_SUBTREE: u64 = u64::pow(2, 8);
|
|
|
|
/// Helper struct for organizing the data we care about when computing Merkle subtrees.
|
|
///
|
|
/// Note that these represet "conceptual" leaves of some subtree, not necessarily
|
|
/// [`SparseMerkleTree::Leaf`].
|
|
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Default)]
|
|
pub struct SubtreeLeaf {
|
|
/// The 'value' field of [`NodeIndex`]. When computing a subtree, the depth is already known.
|
|
pub col: u64,
|
|
/// The hash of the node this `SubtreeLeaf` represents.
|
|
pub hash: RpoDigest,
|
|
}
|
|
|
|
impl SubtreeLeaf {
|
|
#[cfg_attr(not(test), allow(dead_code))]
|
|
fn from_smt_leaf(leaf: &crate::merkle::SmtLeaf) -> Self {
|
|
Self {
|
|
col: leaf.index().index.value(),
|
|
hash: leaf.hash(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Helper struct to organize the return value of [`SparseMerkleTree::sorted_pairs_to_leaves()`].
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub(crate) struct PairComputations<K, L> {
|
|
/// Literal leaves to be added to the sparse Merkle tree's internal mapping.
|
|
pub nodes: BTreeMap<K, L>,
|
|
/// "Conceptual" leaves that will be used for computations.
|
|
pub leaves: Vec<Vec<SubtreeLeaf>>,
|
|
}
|
|
|
|
// Derive requires `L` to impl Default, even though we don't actually need that.
|
|
impl<K, L> Default for PairComputations<K, L> {
|
|
fn default() -> Self {
|
|
Self {
|
|
nodes: Default::default(),
|
|
leaves: Default::default(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Handles the logic for figuring out whether the new leaf starts a new subtree or not.
|
|
fn add_subtree_leaf(subtrees: &mut Vec<Vec<SubtreeLeaf>>, leaf: SubtreeLeaf) {
|
|
let last_subtree = match subtrees.last_mut() {
|
|
// Base case.
|
|
None => {
|
|
subtrees.push(vec![leaf]);
|
|
return;
|
|
},
|
|
Some(last_subtree) => last_subtree,
|
|
};
|
|
|
|
debug_assert!(!last_subtree.is_empty());
|
|
debug_assert!(last_subtree.len() <= COLS_PER_SUBTREE as usize);
|
|
|
|
// The multiple of 256 after 0 is 1, but 0 and 1 do not belong to different subtrees.
|
|
let last_subtree_col = u64::max(1, last_subtree.last().unwrap().col);
|
|
let next_subtree_col = if Integer::is_multiple_of(&last_subtree_col, &COLS_PER_SUBTREE) {
|
|
u64::next_multiple_of(last_subtree_col + 1, COLS_PER_SUBTREE)
|
|
} else {
|
|
last_subtree_col.next_multiple_of(COLS_PER_SUBTREE)
|
|
};
|
|
|
|
if leaf.col < next_subtree_col {
|
|
last_subtree.push(leaf);
|
|
} else {
|
|
let next_subtree = vec![leaf];
|
|
subtrees.push(next_subtree);
|
|
}
|
|
}
|
|
|
|
// TESTS
|
|
// ================================================================================================
|
|
#[cfg(test)]
|
|
mod test {
|
|
use core::mem;
|
|
|
|
use alloc::{collections::BTreeMap, vec::Vec};
|
|
|
|
use super::{SparseMerkleTree, SubtreeLeaf};
|
|
use crate::{
|
|
hash::rpo::RpoDigest,
|
|
merkle::{
|
|
smt::{InnerNode, PairComputations},
|
|
LeafIndex, NodeIndex, Smt, SmtLeaf, SMT_DEPTH,
|
|
},
|
|
Felt, Word, ONE,
|
|
};
|
|
|
|
#[test]
|
|
fn test_sorted_pairs_to_leaves() {
|
|
let entries: Vec<(RpoDigest, Word)> = vec![
|
|
// Subtree 0.
|
|
(RpoDigest::new([ONE, ONE, ONE, Felt::new(16)]), [ONE; 4]),
|
|
(RpoDigest::new([ONE, ONE, ONE, Felt::new(17)]), [ONE; 4]),
|
|
// Leaf index collision.
|
|
(RpoDigest::new([ONE, ONE, Felt::new(10), Felt::new(20)]), [ONE; 4]),
|
|
(RpoDigest::new([ONE, ONE, Felt::new(20), Felt::new(20)]), [ONE; 4]),
|
|
// Subtree 1. Normal single leaf again.
|
|
(RpoDigest::new([ONE, ONE, ONE, Felt::new(400)]), [ONE; 4]), // Subtree boundary.
|
|
(RpoDigest::new([ONE, ONE, ONE, Felt::new(401)]), [ONE; 4]),
|
|
// Subtree 2. Another normal leaf.
|
|
(RpoDigest::new([ONE, ONE, ONE, Felt::new(1024)]), [ONE; 4]),
|
|
];
|
|
|
|
let control = Smt::with_entries(entries.clone()).unwrap();
|
|
|
|
let control_leaves: Vec<SmtLeaf> = {
|
|
let mut entries_iter = entries.iter().cloned();
|
|
let mut next_entry = || entries_iter.next().unwrap();
|
|
let control_leaves = vec![
|
|
// Subtree 0.
|
|
SmtLeaf::Single(next_entry()),
|
|
SmtLeaf::Single(next_entry()),
|
|
SmtLeaf::new_multiple(vec![next_entry(), next_entry()]).unwrap(),
|
|
// Subtree 1.
|
|
SmtLeaf::Single(next_entry()),
|
|
SmtLeaf::Single(next_entry()),
|
|
// Subtree 2.
|
|
SmtLeaf::Single(next_entry()),
|
|
];
|
|
assert_eq!(entries_iter.next(), None);
|
|
control_leaves
|
|
};
|
|
|
|
let control_subtree_leaves: Vec<Vec<SubtreeLeaf>> = {
|
|
let mut control_leaves_iter = control_leaves.iter();
|
|
let mut next_leaf = || control_leaves_iter.next().unwrap();
|
|
|
|
let control_subtree_leaves: Vec<Vec<SubtreeLeaf>> = [
|
|
// Subtree 0.
|
|
vec![next_leaf(), next_leaf(), next_leaf()],
|
|
// Subtree 1.
|
|
vec![next_leaf(), next_leaf()],
|
|
// Subtree 2.
|
|
vec![next_leaf()],
|
|
]
|
|
.map(|subtree| subtree.into_iter().map(SubtreeLeaf::from_smt_leaf).collect())
|
|
.to_vec();
|
|
assert_eq!(control_leaves_iter.next(), None);
|
|
control_subtree_leaves
|
|
};
|
|
|
|
let subtrees = Smt::sorted_pairs_to_leaves(entries);
|
|
// This will check that the hashes, columns, and subtree assignments all match.
|
|
assert_eq!(subtrees.leaves, control_subtree_leaves);
|
|
|
|
// Then finally we might as well check the computed leaf nodes too.
|
|
let control_leaves: BTreeMap<u64, SmtLeaf> = control
|
|
.leaves()
|
|
.map(|(index, value)| (index.index.value(), value.clone()))
|
|
.collect();
|
|
|
|
for (column, test_leaf) in subtrees.nodes {
|
|
if test_leaf.is_empty() {
|
|
continue;
|
|
}
|
|
let control_leaf = control_leaves
|
|
.get(&column)
|
|
.unwrap_or_else(|| panic!("no leaf node found for column {column}"));
|
|
assert_eq!(control_leaf, &test_leaf);
|
|
}
|
|
}
|
|
|
|
// Helper for the below tests.
|
|
fn generate_entries(pair_count: u64) -> Vec<(RpoDigest, Word)> {
|
|
(0..pair_count)
|
|
.map(|i| {
|
|
let leaf_index = ((i as f64 / pair_count as f64) * (pair_count as f64)) as u64;
|
|
let key = RpoDigest::new([ONE, ONE, Felt::new(i), Felt::new(leaf_index)]);
|
|
let value = [ONE, ONE, ONE, Felt::new(i)];
|
|
(key, value)
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
#[test]
|
|
fn test_single_subtree() {
|
|
// A single subtree's worth of leaves.
|
|
const PAIR_COUNT: u64 = 256;
|
|
|
|
let entries = generate_entries(PAIR_COUNT);
|
|
|
|
let control = Smt::with_entries(entries.clone()).unwrap();
|
|
|
|
// `entries` should already be sorted by nature of how we constructed it.
|
|
let leaves = Smt::sorted_pairs_to_leaves(entries).leaves;
|
|
let leaves = leaves.into_iter().next().unwrap();
|
|
|
|
let (first_subtree, _) = Smt::build_subtree(leaves, SMT_DEPTH);
|
|
assert!(!first_subtree.is_empty());
|
|
|
|
for (index, node) in first_subtree.into_iter() {
|
|
let control = control.get_inner_node(index);
|
|
assert_eq!(
|
|
control, node,
|
|
"subtree-computed node at index {index:?} does not match control",
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_two_subtrees() {
|
|
// Two subtrees' worth of leaves.
|
|
const PAIR_COUNT: u64 = 512;
|
|
|
|
let entries = generate_entries(PAIR_COUNT);
|
|
|
|
let control = Smt::with_entries(entries.clone()).unwrap();
|
|
|
|
let PairComputations { leaves, .. } = Smt::sorted_pairs_to_leaves(entries);
|
|
// With two subtrees' worth of leaves, we should have exactly two subtrees.
|
|
let [first, second]: [_; 2] = leaves.try_into().unwrap();
|
|
assert_eq!(first.len() as u64, PAIR_COUNT / 2);
|
|
assert_eq!(first.len(), second.len());
|
|
|
|
let mut current_depth = SMT_DEPTH;
|
|
let mut next_leaves: Vec<SubtreeLeaf> = Default::default();
|
|
|
|
let (first_nodes, leaves) = Smt::build_subtree(first, current_depth);
|
|
next_leaves.extend(leaves);
|
|
|
|
let (second_nodes, leaves) = Smt::build_subtree(second, current_depth);
|
|
next_leaves.extend(leaves);
|
|
|
|
// All new inner nodes + the new subtree-leaves should be 512, for one depth-cycle.
|
|
let total_computed = first_nodes.len() + second_nodes.len() + next_leaves.len();
|
|
assert_eq!(total_computed as u64, PAIR_COUNT);
|
|
|
|
// Verify the computed nodes of both subtrees.
|
|
let computed_nodes = first_nodes.clone().into_iter().chain(second_nodes);
|
|
for (index, test_node) in computed_nodes {
|
|
let control_node = control.get_inner_node(index);
|
|
assert_eq!(
|
|
control_node, test_node,
|
|
"subtree-computed node at index {index:?} does not match control",
|
|
);
|
|
}
|
|
|
|
current_depth -= 8;
|
|
|
|
let (nodes, next_leaves) = Smt::build_subtree(next_leaves, current_depth);
|
|
assert_eq!(nodes.len(), 8);
|
|
assert_eq!(next_leaves.len(), 1);
|
|
|
|
for (index, test_node) in nodes {
|
|
let control_node = control.get_inner_node(index);
|
|
assert_eq!(
|
|
control_node, test_node,
|
|
"subtree-computed node at index {index:?} does not match control",
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_singlethreaded_subtrees() {
|
|
const PAIR_COUNT: u64 = 4096 * 4;
|
|
|
|
let entries = generate_entries(PAIR_COUNT);
|
|
|
|
let control = Smt::with_entries(entries.clone()).unwrap();
|
|
|
|
let mut accumulated_nodes: BTreeMap<NodeIndex, InnerNode> = Default::default();
|
|
|
|
let PairComputations {
|
|
leaves: mut leaf_subtrees,
|
|
nodes: test_leaves,
|
|
} = Smt::sorted_pairs_to_leaves(entries);
|
|
|
|
for current_depth in (8..=SMT_DEPTH).step_by(8).rev() {
|
|
for (i, subtree) in mem::take(&mut leaf_subtrees).into_iter().enumerate() {
|
|
// Pre-assertions.
|
|
assert!(
|
|
subtree.is_sorted(),
|
|
"subtree {i} at bottom-depth {current_depth} is not sorted",
|
|
);
|
|
assert!(
|
|
!subtree.is_empty(),
|
|
"subtree {i} at bottom-depth {current_depth} is empty!",
|
|
);
|
|
|
|
// Do actual things.
|
|
let (nodes, next_leaves) = Smt::build_subtree(subtree, current_depth);
|
|
|
|
// Post-assertions.
|
|
assert!(next_leaves.is_sorted());
|
|
for (&index, test_node) in nodes.iter() {
|
|
let control_node = control.get_inner_node(index);
|
|
assert_eq!(
|
|
test_node, &control_node,
|
|
"depth {} subtree {}: test node does not match control at index {:?}",
|
|
current_depth, i, index,
|
|
);
|
|
}
|
|
|
|
// Update state.
|
|
accumulated_nodes.extend(nodes);
|
|
|
|
for subtree_leaf in next_leaves {
|
|
super::add_subtree_leaf(&mut leaf_subtrees, subtree_leaf);
|
|
}
|
|
}
|
|
|
|
assert!(!leaf_subtrees.is_empty(), "on depth {current_depth}");
|
|
}
|
|
|
|
// Make sure the true leaves match, checking length first and then each individual leaf.
|
|
let control_leaves: BTreeMap<_, _> = control.leaves().collect();
|
|
let control_leaves_len = control_leaves.len();
|
|
let test_leaves_len = test_leaves.len();
|
|
assert_eq!(test_leaves_len, control_leaves_len);
|
|
for (col, ref test_leaf) in test_leaves {
|
|
let index = LeafIndex::new_max_depth(col);
|
|
let &control_leaf = control_leaves.get(&index).unwrap();
|
|
assert_eq!(test_leaf, control_leaf);
|
|
}
|
|
|
|
// Make sure the inner nodes match, checking length first and then each individual leaf.
|
|
let control_nodes_len = control.inner_nodes().count();
|
|
let test_nodes_len = accumulated_nodes.len();
|
|
assert_eq!(test_nodes_len, control_nodes_len);
|
|
for (index, test_node) in accumulated_nodes.clone() {
|
|
let control_node = control.get_inner_node(index);
|
|
assert_eq!(test_node, control_node, "test node does not match control at {index:?}");
|
|
}
|
|
|
|
// After the last iteration of the above for loop, we should have the new root node actually
|
|
// in two places: one in `accumulated_nodes`, and the other as the "next leaves" return from
|
|
// `build_subtree()`. So let's check both!
|
|
|
|
let control_root = control.get_inner_node(NodeIndex::root());
|
|
|
|
// That for loop should have left us with only one leaf subtree...
|
|
let [leaf_subtree]: [_; 1] = leaf_subtrees.try_into().unwrap();
|
|
// which itself contains only one 'leaf'...
|
|
let [root_leaf]: [_; 1] = leaf_subtree.try_into().unwrap();
|
|
// which matches the expected root.
|
|
assert_eq!(control.root(), root_leaf.hash);
|
|
|
|
// Likewise `accumulated_nodes` should contain a node at the root index...
|
|
assert!(accumulated_nodes.contains_key(&NodeIndex::root()));
|
|
// and it should match our actual root.
|
|
let test_root = accumulated_nodes.get(&NodeIndex::root()).unwrap();
|
|
assert_eq!(control_root, *test_root);
|
|
// And of course the root we got from each place should match.
|
|
assert_eq!(control.root(), root_leaf.hash);
|
|
}
|
|
}
|