arnaucube/miden-crypto
1
0
Fork 0
mirror of https://github.com/arnaucube/miden-crypto.git synced 2026-01-28 23:06:45 +01:00
Code Issues Packages Projects Releases Wiki Activity

feat: refactor simple smt to use empty subtree constants

Browse Source 
Prior to this commit, there was an internal procedure with the merkle
trees to compute empty sub-tree for arbitrary depths.

However, this isn't ideal as this code can be reused in any merkle
implementation that uses RPO as backend.

This commit introduces a structure that will generate these empty
subtrees values.
This commit is contained in:
Victor Lopez
2023-02-19 18:30:57 +01:00
parent 7ffa0cd97d
commit 3a6a4fcce6
6 changed files with 560 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
src/merkle/simple_smt/mod.rs
View File

@@ -1,4 +1,6 @@
use super::{BTreeMap, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word};
use super::{
BTreeMap, EmptySubtreeRoots, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word,
};
#[cfg(test)]
mod tests;
@@ -29,38 +31,55 @@ impl SimpleSmt {
// CONSTRUCTORS
// --------------------------------------------------------------------------------------------
/// Creates a new simple SMT.
///
/// The provided entries will be tuples of the leaves and their corresponding keys.
/// Creates a new simple SMT with the provided depth.
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 (store, root) = Store::new(depth);
Ok(Self { root, depth, store })
}
/// Appends the provided entries as leaves of the tree.
///
/// # Errors
///
/// The function will fail if the provided entries count exceed the maximum tree capacity, that
/// is `2^{depth}`.
pub fn new<R, I>(entries: R, depth: u8) -> Result<Self, MerkleError>
pub fn with_leaves<R, I>(mut self, entries: R) -> Result<Self, MerkleError>
where
R: IntoIterator<IntoIter = I>,
I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
{
// check if the leaves count will fit the depth setup
let mut entries = entries.into_iter();
// validate the range of the depth.
let max = 1 << depth;
if depth < Self::MIN_DEPTH {
return Err(MerkleError::DepthTooSmall(depth));
} else if Self::MAX_DEPTH < depth {
return Err(MerkleError::DepthTooBig(depth as u64));
} else if entries.len() > max {
let max = 1 << self.depth;
if entries.len() > max {
return Err(MerkleError::InvalidEntriesCount(max, entries.len()));
}
let (store, root) = Store::new(depth);
let mut tree = Self { root, depth, store };
entries.try_for_each(|(key, leaf)| tree.insert_leaf(key, leaf))?;
Ok(tree)
// append leaves and return
entries.try_for_each(|(key, leaf)| self.insert_leaf(key, leaf))?;
Ok(self)
}
/// Replaces the internal empty digests used when a given depth doesn't contain a node.
pub fn with_empty_subtrees<I>(mut self, hashes: I) -> Self
where
I: IntoIterator<Item = RpoDigest>,
{
self.store
.replace_empty_subtrees(hashes.into_iter().collect());
self
}
// PUBLIC ACCESSORS
// --------------------------------------------------------------------------------------------
/// Returns the root of this Merkle tree.
pub const fn root(&self) -> Word {
self.root
@@ -71,6 +90,9 @@ impl SimpleSmt {
self.depth
}
// PROVIDERS
// --------------------------------------------------------------------------------------------
/// Returns the set count of the keys of the leaves.
pub fn leaves_count(&self) -> usize {
self.store.leaves_count()
@@ -81,16 +103,24 @@ impl SimpleSmt {
/// # Errors
/// Returns an error if:
/// * The specified depth is greater than the depth of the tree.
/// * The specified key does not exist
pub fn get_node(&self, index: &NodeIndex) -> Result<Word, 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() {
self.store.get_leaf_node(index.value())
self.store
.get_leaf_node(index.value())
.or_else(|| {
self.store
.empty_hashes
.get(index.depth() as usize)
.copied()
.map(Word::from)
})
.ok_or(MerkleError::InvalidIndex(*index))
} else {
let branch_node = self.store.get_branch_node(index)?;
let branch_node = self.store.get_branch_node(index);
Ok(Rpo256::merge(&[branch_node.left, branch_node.right]).into())
}
}
@@ -100,23 +130,19 @@ impl SimpleSmt {
///
/// # Errors
/// Returns an error if:
/// * The specified key does not exist as a branch or leaf node
/// * The specified depth is greater than the depth of the 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));
} else if index.depth() == self.depth() && !self.store.check_leaf_node_exists(index.value())
{
return Err(MerkleError::InvalidIndex(index.with_depth(self.depth())));
}
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.store.get_branch_node(&index)?;
let BranchNode { left, right } = self.store.get_branch_node(&index);
let value = if is_right { left } else { right };
path.push(*value);
}
@@ -133,6 +159,9 @@ impl SimpleSmt {
self.get_path(NodeIndex::new(self.depth(), key))
}
// STATE MUTATORS
// --------------------------------------------------------------------------------------------
/// Replaces the leaf located at the specified key, and recomputes hashes by walking up the tree
///
/// # Errors
@@ -156,10 +185,7 @@ impl SimpleSmt {
for _ in 0..index.depth() {
let is_right = index.is_value_odd();
index.move_up();
let BranchNode { left, right } = self
.store
.get_branch_node(&index)
.unwrap_or_else(|_| self.store.get_empty_node(index.depth() as usize + 1));
let BranchNode { left, right } = self.store.get_branch_node(&index);
let (left, right) = if is_right {
(left, value)
} else {
@@ -200,16 +226,7 @@ impl Store {
let leaves = BTreeMap::new();
// Construct empty node digests for each layer of the tree
let empty_hashes: Vec<RpoDigest> = (0..depth + 1)
.scan(Word::default().into(), |state, _| {
let value = *state;
*state = Rpo256::merge(&[value, value]);
Some(value)
})
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect();
let empty_hashes = EmptySubtreeRoots::empty_hashes(depth).to_vec();
let root = empty_hashes[0].into();
let store = Self {
@@ -222,34 +239,30 @@ impl Store {
(store, root)
}
fn get_empty_node(&self, depth: usize) -> BranchNode {
let digest = self.empty_hashes[depth];
BranchNode {
left: digest,
right: digest,
}
fn replace_empty_subtrees(&mut self, hashes: Vec<RpoDigest>) {
self.empty_hashes = hashes;
}
fn check_leaf_node_exists(&self, key: u64) -> bool {
self.leaves.contains_key(&key)
}
fn get_leaf_node(&self, key: u64) -> Result<Word, MerkleError> {
self.leaves
.get(&key)
.cloned()
.ok_or(MerkleError::InvalidIndex(NodeIndex::new(self.depth, key)))
fn get_leaf_node(&self, key: u64) -> Option<Word> {
self.leaves.get(&key).copied()
}
fn insert_leaf_node(&mut self, key: u64, node: Word) {
self.leaves.insert(key, node);
}
fn get_branch_node(&self, index: &NodeIndex) -> Result<BranchNode, MerkleError> {
self.branches
.get(index)
.cloned()
.ok_or(MerkleError::InvalidIndex(*index))
fn get_branch_node(&self, index: &NodeIndex) -> BranchNode {
self.branches.get(index).cloned().unwrap_or_else(|| {
let node = self.empty_hashes[index.depth() as usize + 1];
BranchNode {
left: node,
right: node,
}
})
}
fn insert_branch_node(&mut self, index: NodeIndex, left: RpoDigest, right: RpoDigest) {