Philippe Laferrière
1 year ago
committed by
Bobbin Threadbare
Bobbin Threadbare
13 changed files with 828 additions and 686 deletions
Split View
Diff Options
-
+3 -0.gitignore
-
+11 -3Cargo.toml
-
+43 -44benches/smt.rs
-
+32 -29benches/store.rs
-
+43 -48src/merkle/delta.rs
-
+2 -2src/merkle/mod.rs
-
+10 -2src/merkle/path.rs
-
+0 -389src/merkle/simple_smt/mod.rs
-
+238 -0src/merkle/smt/mod.rs
-
+325 -0src/merkle/smt/simple/mod.rs
-
+58 -112src/merkle/smt/simple/tests.rs
-
+3 -3src/merkle/store/mod.rs
-
+60 -54src/merkle/store/tests.rs
@ -1,389 +0,0 @@ |
|||
use super::{
|
|||
BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerkleTreeDelta,
|
|||
NodeIndex, Rpo256, RpoDigest, StoreNode, TryApplyDiff, Vec, Word,
|
|||
};
|
|||
|
|||
#[cfg(test)]
|
|||
mod tests;
|
|||
|
|||
// SPARSE MERKLE TREE
|
|||
// ================================================================================================
|
|||
|
|||
/// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction.
|
|||
///
|
|||
/// The root of the tree is recomputed on each new leaf update.
|
|||
#[derive(Debug, Clone, PartialEq, Eq)]
|
|||
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
|
|||
pub struct SimpleSmt {
|
|||
depth: u8,
|
|||
root: RpoDigest,
|
|||
leaves: BTreeMap<u64, Word>,
|
|||
branches: BTreeMap<NodeIndex, BranchNode>,
|
|||
}
|
|||
|
|||
impl SimpleSmt {
|
|||
// CONSTANTS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// Minimum supported depth.
|
|||
pub const MIN_DEPTH: u8 = 1;
|
|||
|
|||
/// Maximum supported depth.
|
|||
pub const MAX_DEPTH: u8 = 64;
|
|||
|
|||
/// Value of an empty leaf.
|
|||
pub const EMPTY_VALUE: Word = super::EMPTY_WORD;
|
|||
|
|||
// CONSTRUCTORS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// Returns a new [SimpleSmt] instantiated with the specified depth.
|
|||
///
|
|||
/// All leaves in the returned tree are set to [ZERO; 4].
|
|||
///
|
|||
/// # Errors
|
|||
/// Returns an error if the depth is 0 or is greater than 64.
|
|||
pub fn new(depth: u8) -> Result<Self, MerkleError> {
|
|||
// validate the range of the depth.
|
|||
if depth < Self::MIN_DEPTH {
|
|||
return Err(MerkleError::DepthTooSmall(depth));
|
|||
} else if Self::MAX_DEPTH < depth {
|
|||
return Err(MerkleError::DepthTooBig(depth as u64));
|
|||
}
|
|||
|
|||
let root = *EmptySubtreeRoots::entry(depth, 0);
|
|||
|
|||
Ok(Self {
|
|||
root,
|
|||
depth,
|
|||
leaves: BTreeMap::new(),
|
|||
branches: BTreeMap::new(),
|
|||
})
|
|||
}
|
|||
|
|||
/// Returns a new [SimpleSmt] instantiated with the specified depth and with leaves
|
|||
/// set as specified by the provided entries.
|
|||
///
|
|||
/// All leaves omitted from the entries list are set to [ZERO; 4].
|
|||
///
|
|||
/// # Errors
|
|||
/// Returns an error if:
|
|||
/// - If the depth is 0 or is greater than 64.
|
|||
/// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
|
|||
/// - The provided entries contain multiple values for the same key.
|
|||
pub fn with_leaves(
|
|||
depth: u8,
|
|||
entries: impl IntoIterator<Item = (u64, Word)>,
|
|||
) -> Result<Self, MerkleError> {
|
|||
// create an empty tree
|
|||
let mut tree = Self::new(depth)?;
|
|||
|
|||
// compute the max number of entries. We use an upper bound of depth 63 because we consider
|
|||
// passing in a vector of size 2^64 infeasible.
|
|||
let max_num_entries = 2_usize.pow(tree.depth.min(63).into());
|
|||
|
|||
// This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
|
|||
// entries with the empty value need additional tracking.
|
|||
let mut key_set_to_zero = BTreeSet::new();
|
|||
|
|||
for (idx, (key, value)) in entries.into_iter().enumerate() {
|
|||
if idx >= max_num_entries {
|
|||
return Err(MerkleError::InvalidNumEntries(max_num_entries));
|
|||
}
|
|||
|
|||
let old_value = tree.update_leaf(key, value)?;
|
|||
|
|||
if old_value != Self::EMPTY_VALUE || key_set_to_zero.contains(&key) {
|
|||
return Err(MerkleError::DuplicateValuesForIndex(key));
|
|||
}
|
|||
|
|||
if value == Self::EMPTY_VALUE {
|
|||
key_set_to_zero.insert(key);
|
|||
};
|
|||
}
|
|||
Ok(tree)
|
|||
}
|
|||
|
|||
/// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
|
|||
/// starting at index 0.
|
|||
pub fn with_contiguous_leaves(
|
|||
depth: u8,
|
|||
entries: impl IntoIterator<Item = Word>,
|
|||
) -> Result<Self, MerkleError> {
|
|||
Self::with_leaves(
|
|||
depth,
|
|||
entries
|
|||
.into_iter()
|
|||
.enumerate()
|
|||
.map(|(idx, word)| (idx.try_into().expect("tree max depth is 2^8"), word)),
|
|||
)
|
|||
}
|
|||
|
|||
// PUBLIC ACCESSORS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// Returns the root of this Merkle tree.
|
|||
pub const fn root(&self) -> RpoDigest {
|
|||
self.root
|
|||
}
|
|||
|
|||
/// Returns the depth of this Merkle tree.
|
|||
pub const fn depth(&self) -> u8 {
|
|||
self.depth
|
|||
}
|
|||
|
|||
/// Returns a node at the specified index.
|
|||
///
|
|||
/// # Errors
|
|||
/// Returns an error if the specified index has depth set to 0 or the depth is greater than
|
|||
/// the depth of this Merkle tree.
|
|||
pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
|
|||
if index.is_root() {
|
|||
Err(MerkleError::DepthTooSmall(index.depth()))
|
|||
} else if index.depth() > self.depth() {
|
|||
Err(MerkleError::DepthTooBig(index.depth() as u64))
|
|||
} else if index.depth() == self.depth() {
|
|||
// the lookup in empty_hashes could fail only if empty_hashes were not built correctly
|
|||
// by the constructor as we check the depth of the lookup above.
|
|||
let leaf_pos = index.value();
|
|||
let leaf = match self.get_leaf_node(leaf_pos) {
|
|||
Some(word) => word.into(),
|
|||
None => *EmptySubtreeRoots::entry(self.depth, index.depth()),
|
|||
};
|
|||
Ok(leaf)
|
|||
} else {
|
|||
Ok(self.get_branch_node(&index).parent())
|
|||
}
|
|||
}
|
|||
|
|||
/// Returns a value of the leaf at the specified index.
|
|||
///
|
|||
/// # Errors
|
|||
/// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
|
|||
pub fn get_leaf(&self, index: u64) -> Result<Word, MerkleError> {
|
|||
let index = NodeIndex::new(self.depth, index)?;
|
|||
Ok(self.get_node(index)?.into())
|
|||
}
|
|||
|
|||
/// Returns a Merkle path from the node at the specified index to the root.
|
|||
///
|
|||
/// The node itself is not included in the path.
|
|||
///
|
|||
/// # Errors
|
|||
/// Returns an error if the specified index has depth set to 0 or the depth is greater than
|
|||
/// the depth of this Merkle tree.
|
|||
pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
|
|||
if index.is_root() {
|
|||
return Err(MerkleError::DepthTooSmall(index.depth()));
|
|||
} else if index.depth() > self.depth() {
|
|||
return Err(MerkleError::DepthTooBig(index.depth() as u64));
|
|||
}
|
|||
|
|||
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 BranchNode { left, right } = self.get_branch_node(&index);
|
|||
let value = if is_right { left } else { right };
|
|||
path.push(value);
|
|||
}
|
|||
Ok(MerklePath::new(path))
|
|||
}
|
|||
|
|||
/// Return a Merkle path from the leaf at the specified index to the root.
|
|||
///
|
|||
/// The leaf itself is not included in the path.
|
|||
///
|
|||
/// # Errors
|
|||
/// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
|
|||
pub fn get_leaf_path(&self, index: u64) -> Result<MerklePath, MerkleError> {
|
|||
let index = NodeIndex::new(self.depth(), index)?;
|
|||
self.get_path(index)
|
|||
}
|
|||
|
|||
// ITERATORS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// Returns an iterator over the leaves of this [SimpleSmt].
|
|||
pub fn leaves(&self) -> impl Iterator<Item = (u64, &Word)> {
|
|||
self.leaves.iter().map(|(i, w)| (*i, w))
|
|||
}
|
|||
|
|||
/// Returns an iterator over the inner nodes of this Merkle tree.
|
|||
pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
|
|||
self.branches.values().map(|e| InnerNodeInfo {
|
|||
value: e.parent(),
|
|||
left: e.left,
|
|||
right: e.right,
|
|||
})
|
|||
}
|
|||
|
|||
// STATE MUTATORS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// Updates value of the leaf at the specified index returning the old leaf value.
|
|||
///
|
|||
/// This also recomputes all hashes between the leaf and the root, updating the root itself.
|
|||
///
|
|||
/// # Errors
|
|||
/// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
|
|||
pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<Word, MerkleError> {
|
|||
// validate the index before modifying the structure
|
|||
let idx = NodeIndex::new(self.depth(), index)?;
|
|||
|
|||
let old_value = self.insert_leaf_node(index, value).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 Ok(value);
|
|||
}
|
|||
|
|||
self.recompute_nodes_from_index_to_root(idx, RpoDigest::from(value));
|
|||
|
|||
Ok(old_value)
|
|||
}
|
|||
|
|||
/// Inserts a subtree at the specified index. The depth at which the subtree is inserted is
|
|||
/// computed as `self.depth() - subtree.depth()`.
|
|||
///
|
|||
/// Returns the new root.
|
|||
pub fn set_subtree(
|
|||
&mut self,
|
|||
subtree_insertion_index: u64,
|
|||
subtree: SimpleSmt,
|
|||
) -> Result<RpoDigest, MerkleError> {
|
|||
if subtree.depth() > self.depth() {
|
|||
return Err(MerkleError::InvalidSubtreeDepth {
|
|||
subtree_depth: subtree.depth(),
|
|||
tree_depth: self.depth(),
|
|||
});
|
|||
}
|
|||
|
|||
// Verify that `subtree_insertion_index` is valid.
|
|||
let subtree_root_insertion_depth = self.depth() - subtree.depth();
|
|||
let subtree_root_index =
|
|||
NodeIndex::new(subtree_root_insertion_depth, subtree_insertion_index)?;
|
|||
|
|||
// add leaves
|
|||
// --------------
|
|||
|
|||
// The subtree's leaf indices live in their own context - i.e. a subtree of depth `d`. If we
|
|||
// insert the subtree at `subtree_insertion_index = 0`, then the subtree leaf indices are
|
|||
// valid as they are. However, consider what happens when we insert at
|
|||
// `subtree_insertion_index = 1`. The first leaf of our subtree now will have index `2^d`;
|
|||
// you can see it as there's a full subtree sitting on its left. In general, for
|
|||
// `subtree_insertion_index = i`, there are `i` subtrees sitting before the subtree we want
|
|||
// to insert, so we need to adjust all its leaves by `i * 2^d`.
|
|||
let leaf_index_shift: u64 = subtree_insertion_index * 2_u64.pow(subtree.depth().into());
|
|||
for (subtree_leaf_idx, leaf_value) in subtree.leaves() {
|
|||
let new_leaf_idx = leaf_index_shift + subtree_leaf_idx;
|
|||
debug_assert!(new_leaf_idx < 2_u64.pow(self.depth().into()));
|
|||
|
|||
self.insert_leaf_node(new_leaf_idx, *leaf_value);
|
|||
}
|
|||
|
|||
// add subtree's branch nodes (which includes the root)
|
|||
// --------------
|
|||
for (branch_idx, branch_node) in subtree.branches {
|
|||
let new_branch_idx = {
|
|||
let new_depth = subtree_root_insertion_depth + branch_idx.depth();
|
|||
let new_value = subtree_insertion_index * 2_u64.pow(branch_idx.depth().into())
|
|||
+ branch_idx.value();
|
|||
|
|||
NodeIndex::new(new_depth, new_value).expect("index guaranteed to be valid")
|
|||
};
|
|||
|
|||
self.branches.insert(new_branch_idx, branch_node);
|
|||
}
|
|||
|
|||
// recompute nodes starting from subtree root
|
|||
// --------------
|
|||
self.recompute_nodes_from_index_to_root(subtree_root_index, subtree.root);
|
|||
|
|||
Ok(self.root)
|
|||
}
|
|||
|
|||
// HELPER METHODS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// 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 value = node_hash_at_index;
|
|||
for _ in 0..index.depth() {
|
|||
let is_right = index.is_value_odd();
|
|||
index.move_up();
|
|||
let BranchNode { left, right } = self.get_branch_node(&index);
|
|||
let (left, right) = if is_right { (left, value) } else { (value, right) };
|
|||
self.insert_branch_node(index, left, right);
|
|||
value = Rpo256::merge(&[left, right]);
|
|||
}
|
|||
self.root = value;
|
|||
}
|
|||
|
|||
fn get_leaf_node(&self, key: u64) -> Option<Word> {
|
|||
self.leaves.get(&key).copied()
|
|||
}
|
|||
|
|||
fn insert_leaf_node(&mut self, key: u64, node: Word) -> Option<Word> {
|
|||
self.leaves.insert(key, node)
|
|||
}
|
|||
|
|||
fn get_branch_node(&self, index: &NodeIndex) -> BranchNode {
|
|||
self.branches.get(index).cloned().unwrap_or_else(|| {
|
|||
let node = EmptySubtreeRoots::entry(self.depth, index.depth() + 1);
|
|||
BranchNode { left: *node, right: *node }
|
|||
})
|
|||
}
|
|||
|
|||
fn insert_branch_node(&mut self, index: NodeIndex, left: RpoDigest, right: RpoDigest) {
|
|||
let branch = BranchNode { left, right };
|
|||
self.branches.insert(index, branch);
|
|||
}
|
|||
}
|
|||
|
|||
// BRANCH NODE
|
|||
// ================================================================================================
|
|||
|
|||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
|||
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
|
|||
struct BranchNode {
|
|||
left: RpoDigest,
|
|||
right: RpoDigest,
|
|||
}
|
|||
|
|||
impl BranchNode {
|
|||
fn parent(&self) -> RpoDigest {
|
|||
Rpo256::merge(&[self.left, self.right])
|
|||
}
|
|||
}
|
|||
|
|||
// TRY APPLY DIFF
|
|||
// ================================================================================================
|
|||
impl TryApplyDiff<RpoDigest, StoreNode> for SimpleSmt {
|
|||
type Error = MerkleError;
|
|||
type DiffType = MerkleTreeDelta;
|
|||
|
|||
fn try_apply(&mut self, diff: MerkleTreeDelta) -> Result<(), MerkleError> {
|
|||
if diff.depth() != self.depth() {
|
|||
return Err(MerkleError::InvalidDepth {
|
|||
expected: self.depth(),
|
|||
provided: diff.depth(),
|
|||
});
|
|||
}
|
|||
|
|||
for slot in diff.cleared_slots() {
|
|||
self.update_leaf(*slot, Self::EMPTY_VALUE)?;
|
|||
}
|
|||
|
|||
for (slot, value) in diff.updated_slots() {
|
|||
self.update_leaf(*slot, *value)?;
|
|||
}
|
|||
|
|||
Ok(())
|
|||
}
|
|||
}
|
|||
@ -0,0 +1,238 @@ |
|||
use winter_math::StarkField;
|
|||
|
|||
use crate::{
|
|||
hash::rpo::{Rpo256, RpoDigest},
|
|||
Word,
|
|||
};
|
|||
|
|||
use super::{MerkleError, MerklePath, NodeIndex, Vec};
|
|||
|
|||
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;
|
|||
/// The type for a value
|
|||
type Value: Clone + PartialEq;
|
|||
/// The type for a leaf
|
|||
type Leaf;
|
|||
/// The type for an opening (i.e. a "proof") of a leaf
|
|||
type Opening: From<(MerklePath, Self::Leaf)>;
|
|||
|
|||
/// 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)
|
|||
};
|
|||
|
|||
(merkle_path, leaf).into()
|
|||
}
|
|||
|
|||
/// 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 value = node_hash_at_index;
|
|||
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 (left, right) = if is_right { (left, value) } else { (value, right) };
|
|||
self.insert_inner_node(index, InnerNode { left, right });
|
|||
value = Rpo256::merge(&[left, right]);
|
|||
}
|
|||
self.set_root(value);
|
|||
}
|
|||
|
|||
// 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);
|
|||
|
|||
/// 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 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;
|
|||
|
|||
/// Maps a key to a leaf index
|
|||
fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
|
|||
}
|
|||
|
|||
// INNER NODE
|
|||
// ================================================================================================
|
|||
|
|||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
|||
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
|
|||
pub(crate) 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())
|
|||
}
|
|||
}
|
|||
|
|||
impl From<Word> for LeafIndex<SMT_MAX_DEPTH> {
|
|||
fn from(value: Word) -> Self {
|
|||
// We use the most significant `Felt` of a `Word` as the leaf index.
|
|||
Self::new_max_depth(value[3].as_int())
|
|||
}
|
|||
}
|
|||
|
|||
impl From<RpoDigest> for LeafIndex<SMT_MAX_DEPTH> {
|
|||
fn from(value: RpoDigest) -> Self {
|
|||
Word::from(value).into()
|
|||
}
|
|||
}
|
|||
@ -0,0 +1,325 @@ |
|||
use crate::{
|
|||
merkle::{EmptySubtreeRoots, InnerNodeInfo, MerkleTreeDelta, StoreNode, ValuePath},
|
|||
utils::collections::TryApplyDiff,
|
|||
EMPTY_WORD,
|
|||
};
|
|||
|
|||
use super::{
|
|||
InnerNode, LeafIndex, MerkleError, NodeIndex, RpoDigest, SparseMerkleTree, Word, SMT_MAX_DEPTH,
|
|||
SMT_MIN_DEPTH,
|
|||
};
|
|||
use crate::utils::collections::{BTreeMap, BTreeSet};
|
|||
|
|||
#[cfg(test)]
|
|||
mod tests;
|
|||
|
|||
// SPARSE MERKLE TREE
|
|||
// ================================================================================================
|
|||
|
|||
/// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction.
|
|||
///
|
|||
/// The root of the tree is recomputed on each new leaf update.
|
|||
#[derive(Debug, Clone, PartialEq, Eq)]
|
|||
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
|
|||
pub struct SimpleSmt<const DEPTH: u8> {
|
|||
root: RpoDigest,
|
|||
leaves: BTreeMap<u64, Word>,
|
|||
inner_nodes: BTreeMap<NodeIndex, InnerNode>,
|
|||
}
|
|||
|
|||
impl<const DEPTH: u8> SimpleSmt<DEPTH> {
|
|||
// CONSTANTS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// The default value used to compute the hash of empty leaves
|
|||
pub const EMPTY_VALUE: Word = <Self as SparseMerkleTree<DEPTH>>::EMPTY_VALUE;
|
|||
|
|||
// CONSTRUCTORS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// Returns a new [SimpleSmt].
|
|||
///
|
|||
/// All leaves in the returned tree are set to [ZERO; 4].
|
|||
///
|
|||
/// # Errors
|
|||
/// Returns an error if DEPTH is 0 or is greater than 64.
|
|||
pub fn new() -> Result<Self, MerkleError> {
|
|||
// validate the range of the depth.
|
|||
if DEPTH < SMT_MIN_DEPTH {
|
|||
return Err(MerkleError::DepthTooSmall(DEPTH));
|
|||
} else if SMT_MAX_DEPTH < DEPTH {
|
|||
return Err(MerkleError::DepthTooBig(DEPTH as u64));
|
|||
}
|
|||
|
|||
let root = *EmptySubtreeRoots::entry(DEPTH, 0);
|
|||
|
|||
Ok(Self {
|
|||
root,
|
|||
leaves: BTreeMap::new(),
|
|||
inner_nodes: BTreeMap::new(),
|
|||
})
|
|||
}
|
|||
|
|||
/// Returns a new [SimpleSmt] instantiated with leaves set as specified by the provided entries.
|
|||
///
|
|||
/// All leaves omitted from the entries list are set to [ZERO; 4].
|
|||
///
|
|||
/// # Errors
|
|||
/// Returns an error if:
|
|||
/// - If the depth is 0 or is greater than 64.
|
|||
/// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
|
|||
/// - The provided entries contain multiple values for the same key.
|
|||
pub fn with_leaves(
|
|||
entries: impl IntoIterator<Item = (u64, Word)>,
|
|||
) -> Result<Self, MerkleError> {
|
|||
// create an empty tree
|
|||
let mut tree = Self::new()?;
|
|||
|
|||
// compute the max number of entries. We use an upper bound of depth 63 because we consider
|
|||
// passing in a vector of size 2^64 infeasible.
|
|||
let max_num_entries = 2_usize.pow(DEPTH.min(63).into());
|
|||
|
|||
// This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
|
|||
// entries with the empty value need additional tracking.
|
|||
let mut key_set_to_zero = BTreeSet::new();
|
|||
|
|||
for (idx, (key, value)) in entries.into_iter().enumerate() {
|
|||
if idx >= max_num_entries {
|
|||
return Err(MerkleError::InvalidNumEntries(max_num_entries));
|
|||
}
|
|||
|
|||
let old_value = tree.insert(LeafIndex::<DEPTH>::new(key)?, value);
|
|||
|
|||
if old_value != Self::EMPTY_VALUE || key_set_to_zero.contains(&key) {
|
|||
return Err(MerkleError::DuplicateValuesForIndex(key));
|
|||
}
|
|||
|
|||
if value == Self::EMPTY_VALUE {
|
|||
key_set_to_zero.insert(key);
|
|||
};
|
|||
}
|
|||
Ok(tree)
|
|||
}
|
|||
|
|||
/// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
|
|||
/// starting at index 0.
|
|||
pub fn with_contiguous_leaves(
|
|||
entries: impl IntoIterator<Item = Word>,
|
|||
) -> Result<Self, MerkleError> {
|
|||
Self::with_leaves(
|
|||
entries
|
|||
.into_iter()
|
|||
.enumerate()
|
|||
.map(|(idx, word)| (idx.try_into().expect("tree max depth is 2^8"), word)),
|
|||
)
|
|||
}
|
|||
|
|||
// PUBLIC ACCESSORS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// Returns the depth of the tree
|
|||
pub const fn depth(&self) -> u8 {
|
|||
DEPTH
|
|||
}
|
|||
|
|||
/// Returns the root of the tree
|
|||
pub fn root(&self) -> RpoDigest {
|
|||
<Self as SparseMerkleTree<DEPTH>>::root(self)
|
|||
}
|
|||
|
|||
/// Returns the leaf at the specified index.
|
|||
pub fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
|
|||
<Self as SparseMerkleTree<DEPTH>>::get_leaf(self, key)
|
|||
}
|
|||
|
|||
/// Returns a node at the specified index.
|
|||
///
|
|||
/// # Errors
|
|||
/// Returns an error if the specified index has depth set to 0 or the depth is greater than
|
|||
/// the depth of this Merkle tree.
|
|||
pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
|
|||
if index.is_root() {
|
|||
Err(MerkleError::DepthTooSmall(index.depth()))
|
|||
} else if index.depth() > DEPTH {
|
|||
Err(MerkleError::DepthTooBig(index.depth() as u64))
|
|||
} else if index.depth() == DEPTH {
|
|||
let leaf = self.get_leaf(&LeafIndex::<DEPTH>::try_from(index)?);
|
|||
|
|||
Ok(leaf.into())
|
|||
} else {
|
|||
Ok(self.get_inner_node(index).hash())
|
|||
}
|
|||
}
|
|||
|
|||
/// 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.
|
|||
pub fn open(&self, key: &LeafIndex<DEPTH>) -> ValuePath {
|
|||
<Self as SparseMerkleTree<DEPTH>>::open(self, key)
|
|||
}
|
|||
|
|||
// ITERATORS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// Returns an iterator over the leaves of this [SimpleSmt].
|
|||
pub fn leaves(&self) -> impl Iterator<Item = (u64, &Word)> {
|
|||
self.leaves.iter().map(|(i, w)| (*i, w))
|
|||
}
|
|||
|
|||
/// Returns an iterator over the inner nodes of this [SimpleSmt].
|
|||
pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
|
|||
self.inner_nodes.values().map(|e| InnerNodeInfo {
|
|||
value: e.hash(),
|
|||
left: e.left,
|
|||
right: e.right,
|
|||
})
|
|||
}
|
|||
|
|||
// STATE MUTATORS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// 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
|
|||
/// [`EMPTY_WORD`].
|
|||
///
|
|||
/// This also recomputes all hashes between the leaf (associated with the key) and the root,
|
|||
/// updating the root itself.
|
|||
pub fn insert(&mut self, key: LeafIndex<DEPTH>, value: Word) -> Word {
|
|||
<Self as SparseMerkleTree<DEPTH>>::insert(self, key, value)
|
|||
}
|
|||
|
|||
/// Inserts a subtree at the specified index. The depth at which the subtree is inserted is
|
|||
/// computed as `DEPTH - SUBTREE_DEPTH`.
|
|||
///
|
|||
/// Returns the new root.
|
|||
pub fn set_subtree<const SUBTREE_DEPTH: u8>(
|
|||
&mut self,
|
|||
subtree_insertion_index: u64,
|
|||
subtree: SimpleSmt<SUBTREE_DEPTH>,
|
|||
) -> Result<RpoDigest, MerkleError> {
|
|||
if SUBTREE_DEPTH > DEPTH {
|
|||
return Err(MerkleError::InvalidSubtreeDepth {
|
|||
subtree_depth: SUBTREE_DEPTH,
|
|||
tree_depth: DEPTH,
|
|||
});
|
|||
}
|
|||
|
|||
// Verify that `subtree_insertion_index` is valid.
|
|||
let subtree_root_insertion_depth = DEPTH - SUBTREE_DEPTH;
|
|||
let subtree_root_index =
|
|||
NodeIndex::new(subtree_root_insertion_depth, subtree_insertion_index)?;
|
|||
|
|||
// add leaves
|
|||
// --------------
|
|||
|
|||
// The subtree's leaf indices live in their own context - i.e. a subtree of depth `d`. If we
|
|||
// insert the subtree at `subtree_insertion_index = 0`, then the subtree leaf indices are
|
|||
// valid as they are. However, consider what happens when we insert at
|
|||
// `subtree_insertion_index = 1`. The first leaf of our subtree now will have index `2^d`;
|
|||
// you can see it as there's a full subtree sitting on its left. In general, for
|
|||
// `subtree_insertion_index = i`, there are `i` subtrees sitting before the subtree we want
|
|||
// to insert, so we need to adjust all its leaves by `i * 2^d`.
|
|||
let leaf_index_shift: u64 = subtree_insertion_index * 2_u64.pow(SUBTREE_DEPTH.into());
|
|||
for (subtree_leaf_idx, leaf_value) in subtree.leaves() {
|
|||
let new_leaf_idx = leaf_index_shift + subtree_leaf_idx;
|
|||
debug_assert!(new_leaf_idx < 2_u64.pow(DEPTH.into()));
|
|||
|
|||
self.leaves.insert(new_leaf_idx, *leaf_value);
|
|||
}
|
|||
|
|||
// add subtree's branch nodes (which includes the root)
|
|||
// --------------
|
|||
for (branch_idx, branch_node) in subtree.inner_nodes {
|
|||
let new_branch_idx = {
|
|||
let new_depth = subtree_root_insertion_depth + branch_idx.depth();
|
|||
let new_value = subtree_insertion_index * 2_u64.pow(branch_idx.depth().into())
|
|||
+ branch_idx.value();
|
|||
|
|||
NodeIndex::new(new_depth, new_value).expect("index guaranteed to be valid")
|
|||
};
|
|||
|
|||
self.inner_nodes.insert(new_branch_idx, branch_node);
|
|||
}
|
|||
|
|||
// recompute nodes starting from subtree root
|
|||
// --------------
|
|||
self.recompute_nodes_from_index_to_root(subtree_root_index, subtree.root);
|
|||
|
|||
Ok(self.root)
|
|||
}
|
|||
}
|
|||
|
|||
impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
|
|||
type Key = LeafIndex<DEPTH>;
|
|||
type Value = Word;
|
|||
type Leaf = Word;
|
|||
type Opening = ValuePath;
|
|||
|
|||
const EMPTY_VALUE: Self::Value = EMPTY_WORD;
|
|||
|
|||
fn root(&self) -> RpoDigest {
|
|||
self.root
|
|||
}
|
|||
|
|||
fn set_root(&mut self, root: RpoDigest) {
|
|||
self.root = root;
|
|||
}
|
|||
|
|||
fn get_inner_node(&self, index: NodeIndex) -> InnerNode {
|
|||
self.inner_nodes.get(&index).cloned().unwrap_or_else(|| {
|
|||
let node = EmptySubtreeRoots::entry(DEPTH, index.depth() + 1);
|
|||
|
|||
InnerNode { left: *node, right: *node }
|
|||
})
|
|||
}
|
|||
|
|||
fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) {
|
|||
self.inner_nodes.insert(index, inner_node);
|
|||
}
|
|||
|
|||
fn insert_value(&mut self, key: LeafIndex<DEPTH>, value: Word) -> Option<Word> {
|
|||
self.leaves.insert(key.value(), value)
|
|||
}
|
|||
|
|||
fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
|
|||
// the lookup in empty_hashes could fail only if empty_hashes were not built correctly
|
|||
// by the constructor as we check the depth of the lookup above.
|
|||
let leaf_pos = key.value();
|
|||
|
|||
match self.leaves.get(&leaf_pos) {
|
|||
Some(word) => *word,
|
|||
None => Word::from(*EmptySubtreeRoots::entry(DEPTH, DEPTH)),
|
|||
}
|
|||
}
|
|||
|
|||
fn hash_leaf(leaf: &Word) -> RpoDigest {
|
|||
// `SimpleSmt` takes the leaf value itself as the hash
|
|||
leaf.into()
|
|||
}
|
|||
|
|||
fn key_to_leaf_index(key: &LeafIndex<DEPTH>) -> LeafIndex<DEPTH> {
|
|||
*key
|
|||
}
|
|||
}
|
|||
|
|||
// TRY APPLY DIFF
|
|||
// ================================================================================================
|
|||
impl<const DEPTH: u8> TryApplyDiff<RpoDigest, StoreNode> for SimpleSmt<DEPTH> {
|
|||
type Error = MerkleError;
|
|||
type DiffType = MerkleTreeDelta;
|
|||
|
|||
fn try_apply(&mut self, diff: MerkleTreeDelta) -> Result<(), MerkleError> {
|
|||
if diff.depth() != DEPTH {
|
|||
return Err(MerkleError::InvalidDepth { expected: DEPTH, provided: diff.depth() });
|
|||
}
|
|||
|
|||
for slot in diff.cleared_slots() {
|
|||
self.insert(LeafIndex::<DEPTH>::new(*slot)?, Self::EMPTY_VALUE);
|
|||
}
|
|||
|
|||
for (slot, value) in diff.updated_slots() {
|
|||
self.insert(LeafIndex::<DEPTH>::new(*slot)?, *value);
|
|||
}
|
|||
|
|||
Ok(())
|
|||
}
|
|||
}
|
|||