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 ff9c8242ee
commit 95aeea9e07
1 changed files with 34 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -383,7 +383,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
let hash = Self::hash_leaf(&leaf); let hash = Self::hash_leaf(&leaf);
accumulator.nodes.insert(col, 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()); 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)] #[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct PairComputations<K, L> { pub(crate) struct PairComputations<K, L> {
/// Literal leaves to be added to the sparse Merkle tree's internal mapping. /// Literal leaves to be added to the sparse Merkle tree's internal mapping.
@ -639,12 +640,22 @@ pub(crate) struct PairComputations<K, L> {
pub leaves: Vec<Vec<SubtreeLeaf>>, pub leaves: Vec<Vec<SubtreeLeaf>>,
} }
impl<K, L> PrecomputedSubtrees<K, L> { // Derive requires `L` to impl Default, even though we don't actually need that.
pub fn add_leaf(&mut self, leaf: SubtreeLeaf) { impl<K, L> Default for PairComputations<K, L> {
let last_subtree = match self.leaves.last_mut() { 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. // Base case.
None => { None => {
self.leaves.push(vec![leaf]); subtrees.push(vec![leaf]);
return; return;
}, },
Some(last_subtree) => last_subtree, Some(last_subtree) => last_subtree,
@ -660,22 +671,12 @@ impl<K, L> PrecomputedSubtrees<K, L> {
} else { } else {
last_subtree_col.next_multiple_of(COLS_PER_SUBTREE) last_subtree_col.next_multiple_of(COLS_PER_SUBTREE)
}; };
if leaf.col < next_subtree_col { if leaf.col < next_subtree_col {
last_subtree.push(leaf); last_subtree.push(leaf);
} else { } else {
let next_subtree = vec![leaf]; let next_subtree = vec![leaf];
self.leaves.push(next_subtree); subtrees.push(next_subtree);
}
}
}
// 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(),
}
} }
} }
@ -687,9 +688,7 @@ mod test {
use alloc::{collections::BTreeMap, vec::Vec}; use alloc::{collections::BTreeMap, vec::Vec};
use num::Integer; use super::{SparseMerkleTree, SubtreeLeaf};
use super::{SparseMerkleTree, SubtreeLeaf, COLS_PER_SUBTREE};
use crate::{ use crate::{
hash::rpo::RpoDigest, hash::rpo::RpoDigest,
merkle::{smt::InnerNode, NodeIndex, Smt, SmtLeaf, SMT_DEPTH}, merkle::{smt::InnerNode, NodeIndex, Smt, SmtLeaf, SMT_DEPTH},
@ -904,25 +903,7 @@ mod test {
accumulated_nodes.extend(nodes); accumulated_nodes.extend(nodes);
for subtree_leaf in next_leaves { for subtree_leaf in next_leaves {
if leaf_subtrees.is_empty() { super::add_subtree_leaf(&mut leaf_subtrees, subtree_leaf);
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]);
}
} }
} }