qyriad/miden-crypto
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

factor out subtree-append logic

Browse source
This commit is contained in:
Qyriad 2024-11-04 12:53:27 -07:00
parent 2b04a93a15
commit 60f4dd2161
1 changed files with 203 additions and 192 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

79
src/merkle/smt/mod.rs
View file

@ -89,7 +89,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
};
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
@ -114,7 +114,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
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.
@ -122,7 +122,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
&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();
@ -144,7 +144,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
}
}
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.
@ -157,7 +157,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
fn compute_mutations(
&self,
kv_pairs: impl IntoIterator<Item = (Self::Key, Self::Value)>,
) -> MutationSet<DEPTH, Self::Key, Self::Value> {
) -> MutationSet<DEPTH, Self::Key, Self::Value> {
use NodeMutation::*;
let mut new_root = self.root();
@ -242,7 +242,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
node_mutations,
new_pairs,
}
}
}
/// Apply the prospective mutations computed with [`SparseMerkleTree::compute_mutations()`] to
/// this tree.
@ -255,10 +255,10 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
fn apply_mutations(
&mut self,
mutations: MutationSet<DEPTH, Self::Key, Self::Value>,
) -> Result<(), MerkleError>
where
) -> Result<(), MerkleError>
where
Self: Sized,
{
{
use NodeMutation::*;
let MutationSet {
old_root,
@ -287,7 +287,7 @@ where
self.set_root(new_root);
Ok(())
}
}
// REQUIRED METHODS
// ---------------------------------------------------------------------------------------------
@ -336,7 +336,7 @@ where
existing_leaf: Self::Leaf,
key: &Self::Key,
value: &Self::Value,
) -> Self::Leaf;
) -> Self::Leaf;
/// Maps a key to a leaf index
fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
@ -383,7 +383,7 @@ where
let hash = Self::hash_leaf(&leaf);
accumulator.nodes.insert(col, leaf);
accumulator.add_leaf(SubtreeLeaf { col, hash });
add_subtree_leaf(&mut accumulator.leaves, SubtreeLeaf { col, hash });
debug_assert!(current_leaf_buffer.is_empty());
}
@ -631,6 +631,7 @@ impl SubtreeLeaf {
}
}
/// 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.
@ -679,6 +680,36 @@ impl<K, L> Default for PairComputations<K, L> {
}
}
/// 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)]
@ -687,9 +718,7 @@ mod test {
use alloc::{collections::BTreeMap, vec::Vec};
use num::Integer;
use super::{InnerNode, PairComputations, SparseMerkleTree, SubtreeLeaf, COLS_PER_SUBTREE};
use super::{InnerNode, PairComputations, SparseMerkleTree, SubtreeLeaf};
use crate::{
hash::rpo::RpoDigest,
merkle::{NodeIndex, Smt, SmtLeaf, SMT_DEPTH},
@ -904,25 +933,7 @@ mod test {
accumulated_nodes.extend(nodes);
for subtree_leaf in next_leaves {
if leaf_subtrees.is_empty() {
leaf_subtrees.push(vec![subtree_leaf]);
continue;
}
let buffer_max_col =
u64::max(1, leaf_subtrees.last().unwrap().last().unwrap().col);
let next_subtree_col =
if Integer::is_multiple_of(&buffer_max_col, &COLS_PER_SUBTREE) {
u64::next_multiple_of(buffer_max_col + 1, COLS_PER_SUBTREE)
} else {
buffer_max_col.next_multiple_of(COLS_PER_SUBTREE)
};
if subtree_leaf.col < next_subtree_col {
leaf_subtrees.last_mut().unwrap().push(subtree_leaf);
} else {
leaf_subtrees.push(vec![subtree_leaf]);
}
super::add_subtree_leaf(&mut leaf_subtrees, subtree_leaf);
}
}