feat: implement concurrent Smt construction (#341)
* merkle: add parent() helper function on NodeIndex * smt: add pairs_to_leaf() to trait * smt: add sorted_pairs_to_leaves() and test for it * smt: implement single subtree-8 hashing, w/ benchmarks & tests This will be composed into depth-8-subtree-based computation of entire sparse Merkle trees. * merkle: add a benchmark for constructing 256-balanced trees This is intended for comparison with the benchmarks from the previous commit. This benchmark represents the theoretical perfect-efficiency performance we could possibly (but impractically) get for computing depth-8 sparse Merkle subtrees. * smt: test that SparseMerkleTree::build_subtree() is composable * smt: test that subtree logic can correctly construct an entire tree This commit ensures that `SparseMerkleTree::build_subtree()` can correctly compose into building an entire sparse Merkle tree, without yet getting into potential complications concurrency introduces. * smt: implement test for basic parallelized subtree computation w/ rayon Building on the previous commit, this commit implements a test proving that `SparseMerkleTree::build_subtree()` can be composed into itself not just concurrently, but in parallel, without issue. * smt: add from_raw_parts() to trait interface This commit adds a new required method to the SparseMerkleTree trait, to allow generic construction from pre-computed parts. This will be used to add a generic version of `with_entries()` in a later commit. * smt: add parallel constructors to Smt and SimpleSmt What the previous few commits have been leading up to: SparseMerkleTree now has a function to construct the tree from existing data in parallel. This is significantly faster than the singlethreaded equivalent. Benchmarks incoming! --------- Co-authored-by: krushimir <krushimir@reilabs.co> Co-authored-by: krushimir <kresimir.grofelnik@reilabs.io>
This commit is contained in:
parent
1867f842d3
commit
b151773b0d
14 changed files with 1194 additions and 6 deletions
|
@ -7,6 +7,7 @@
|
|||
## 0.12.0 (2024-10-30)
|
||||
|
||||
- [BREAKING] Updated Winterfell dependency to v0.10 (#338).
|
||||
- Added parallel implementation of `Smt::with_entries()` with significantly better performance when the `concurrent` feature is enabled (#341).
|
||||
|
||||
## 0.11.0 (2024-10-17)
|
||||
|
||||
|
|
1
Cargo.lock
generated
1
Cargo.lock
generated
|
@ -541,6 +541,7 @@ dependencies = [
|
|||
"rand",
|
||||
"rand_chacha",
|
||||
"rand_core",
|
||||
"rayon",
|
||||
"seq-macro",
|
||||
"serde",
|
||||
"sha3",
|
||||
|
|
18
Cargo.toml
18
Cargo.toml
|
@ -27,13 +27,28 @@ harness = false
|
|||
name = "smt"
|
||||
harness = false
|
||||
|
||||
[[bench]]
|
||||
name = "smt-subtree"
|
||||
harness = false
|
||||
required-features = ["internal"]
|
||||
|
||||
[[bench]]
|
||||
name = "merkle"
|
||||
harness = false
|
||||
|
||||
[[bench]]
|
||||
name = "smt-with-entries"
|
||||
harness = false
|
||||
|
||||
[[bench]]
|
||||
name = "store"
|
||||
harness = false
|
||||
|
||||
[features]
|
||||
default = ["std"]
|
||||
concurrent = ["dep:rayon"]
|
||||
default = ["std", "concurrent"]
|
||||
executable = ["dep:clap", "dep:rand-utils", "std"]
|
||||
internal = []
|
||||
serde = ["dep:serde", "serde?/alloc", "winter-math/serde"]
|
||||
std = [
|
||||
"blake3/std",
|
||||
|
@ -53,6 +68,7 @@ num-complex = { version = "0.4", default-features = false }
|
|||
rand = { version = "0.8", default-features = false }
|
||||
rand_core = { version = "0.6", default-features = false }
|
||||
rand-utils = { version = "0.11", package = "winter-rand-utils", optional = true }
|
||||
rayon = { version = "1.10", optional = true }
|
||||
serde = { version = "1.0", default-features = false, optional = true, features = ["derive"] }
|
||||
sha3 = { version = "0.10", default-features = false }
|
||||
thiserror = { version = "2.0", default-features = false }
|
||||
|
|
2
Makefile
2
Makefile
|
@ -83,4 +83,4 @@ build-sve: ## Build with sve support
|
|||
|
||||
.PHONY: bench-tx
|
||||
bench-tx: ## Run crypto benchmarks
|
||||
cargo bench
|
||||
cargo bench --features="concurrent"
|
||||
|
|
|
@ -60,10 +60,11 @@ make
|
|||
|
||||
This crate can be compiled with the following features:
|
||||
|
||||
- `concurrent`- enabled by default; enables multi-threaded implementation of `Smt::with_entries()` which significantly improves performance on multi-core CPUs.
|
||||
- `std` - enabled by default and relies on the Rust standard library.
|
||||
- `no_std` does not rely on the Rust standard library and enables compilation to WebAssembly.
|
||||
|
||||
Both of these features imply the use of [alloc](https://doc.rust-lang.org/alloc/) to support heap-allocated collections.
|
||||
All of these features imply the use of [alloc](https://doc.rust-lang.org/alloc/) to support heap-allocated collections.
|
||||
|
||||
To compile with `no_std`, disable default features via `--no-default-features` flag or using the following command:
|
||||
|
||||
|
|
66
benches/merkle.rs
Normal file
66
benches/merkle.rs
Normal file
|
@ -0,0 +1,66 @@
|
|||
//! Benchmark for building a [`miden_crypto::merkle::MerkleTree`]. This is intended to be compared
|
||||
//! with the results from `benches/smt-subtree.rs`, as building a fully balanced Merkle tree with
|
||||
//! 256 leaves should indicate the *absolute best* performance we could *possibly* get for building
|
||||
//! a depth-8 sparse Merkle subtree, though practically speaking building a fully balanced Merkle
|
||||
//! tree will perform better than the sparse version. At the time of this writing (2024/11/24), this
|
||||
//! benchmark is about four times more efficient than the equivalent benchmark in
|
||||
//! `benches/smt-subtree.rs`.
|
||||
use std::{hint, mem, time::Duration};
|
||||
|
||||
use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
|
||||
use miden_crypto::{merkle::MerkleTree, Felt, Word, ONE};
|
||||
use rand_utils::prng_array;
|
||||
|
||||
fn balanced_merkle_even(c: &mut Criterion) {
|
||||
c.bench_function("balanced-merkle-even", |b| {
|
||||
b.iter_batched(
|
||||
|| {
|
||||
let entries: Vec<Word> =
|
||||
(0..256).map(|i| [Felt::new(i), ONE, ONE, Felt::new(i)]).collect();
|
||||
assert_eq!(entries.len(), 256);
|
||||
entries
|
||||
},
|
||||
|leaves| {
|
||||
let tree = MerkleTree::new(hint::black_box(leaves)).unwrap();
|
||||
assert_eq!(tree.depth(), 8);
|
||||
},
|
||||
BatchSize::SmallInput,
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
fn balanced_merkle_rand(c: &mut Criterion) {
|
||||
let mut seed = [0u8; 32];
|
||||
c.bench_function("balanced-merkle-rand", |b| {
|
||||
b.iter_batched(
|
||||
|| {
|
||||
let entries: Vec<Word> = (0..256).map(|_| generate_word(&mut seed)).collect();
|
||||
assert_eq!(entries.len(), 256);
|
||||
entries
|
||||
},
|
||||
|leaves| {
|
||||
let tree = MerkleTree::new(hint::black_box(leaves)).unwrap();
|
||||
assert_eq!(tree.depth(), 8);
|
||||
},
|
||||
BatchSize::SmallInput,
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
criterion_group! {
|
||||
name = smt_subtree_group;
|
||||
config = Criterion::default()
|
||||
.measurement_time(Duration::from_secs(20))
|
||||
.configure_from_args();
|
||||
targets = balanced_merkle_even, balanced_merkle_rand
|
||||
}
|
||||
criterion_main!(smt_subtree_group);
|
||||
|
||||
// HELPER FUNCTIONS
|
||||
// --------------------------------------------------------------------------------------------
|
||||
|
||||
fn generate_word(seed: &mut [u8; 32]) -> Word {
|
||||
mem::swap(seed, &mut prng_array(*seed));
|
||||
let nums: [u64; 4] = prng_array(*seed);
|
||||
[Felt::new(nums[0]), Felt::new(nums[1]), Felt::new(nums[2]), Felt::new(nums[3])]
|
||||
}
|
142
benches/smt-subtree.rs
Normal file
142
benches/smt-subtree.rs
Normal file
|
@ -0,0 +1,142 @@
|
|||
use std::{fmt::Debug, hint, mem, time::Duration};
|
||||
|
||||
use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
|
||||
use miden_crypto::{
|
||||
hash::rpo::RpoDigest,
|
||||
merkle::{build_subtree_for_bench, NodeIndex, SmtLeaf, SubtreeLeaf, SMT_DEPTH},
|
||||
Felt, Word, ONE,
|
||||
};
|
||||
use rand_utils::prng_array;
|
||||
use winter_utils::Randomizable;
|
||||
|
||||
const PAIR_COUNTS: [u64; 5] = [1, 64, 128, 192, 256];
|
||||
|
||||
fn smt_subtree_even(c: &mut Criterion) {
|
||||
let mut seed = [0u8; 32];
|
||||
|
||||
let mut group = c.benchmark_group("subtree8-even");
|
||||
|
||||
for pair_count in PAIR_COUNTS {
|
||||
let bench_id = BenchmarkId::from_parameter(pair_count);
|
||||
group.bench_with_input(bench_id, &pair_count, |b, &pair_count| {
|
||||
b.iter_batched(
|
||||
|| {
|
||||
// Setup.
|
||||
let entries: Vec<(RpoDigest, Word)> = (0..pair_count)
|
||||
.map(|n| {
|
||||
// A single depth-8 subtree can have a maximum of 255 leaves.
|
||||
let leaf_index = ((n as f64 / pair_count as f64) * 255.0) as u64;
|
||||
let key = RpoDigest::new([
|
||||
generate_value(&mut seed),
|
||||
ONE,
|
||||
Felt::new(n),
|
||||
Felt::new(leaf_index),
|
||||
]);
|
||||
let value = generate_word(&mut seed);
|
||||
(key, value)
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut leaves: Vec<_> = entries
|
||||
.iter()
|
||||
.map(|(key, value)| {
|
||||
let leaf = SmtLeaf::new_single(*key, *value);
|
||||
let col = NodeIndex::from(leaf.index()).value();
|
||||
let hash = leaf.hash();
|
||||
SubtreeLeaf { col, hash }
|
||||
})
|
||||
.collect();
|
||||
leaves.sort();
|
||||
leaves.dedup_by_key(|leaf| leaf.col);
|
||||
leaves
|
||||
},
|
||||
|leaves| {
|
||||
// Benchmarked function.
|
||||
let (subtree, _) = build_subtree_for_bench(
|
||||
hint::black_box(leaves),
|
||||
hint::black_box(SMT_DEPTH),
|
||||
hint::black_box(SMT_DEPTH),
|
||||
);
|
||||
assert!(!subtree.is_empty());
|
||||
},
|
||||
BatchSize::SmallInput,
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn smt_subtree_random(c: &mut Criterion) {
|
||||
let mut seed = [0u8; 32];
|
||||
|
||||
let mut group = c.benchmark_group("subtree8-rand");
|
||||
|
||||
for pair_count in PAIR_COUNTS {
|
||||
let bench_id = BenchmarkId::from_parameter(pair_count);
|
||||
group.bench_with_input(bench_id, &pair_count, |b, &pair_count| {
|
||||
b.iter_batched(
|
||||
|| {
|
||||
// Setup.
|
||||
let entries: Vec<(RpoDigest, Word)> = (0..pair_count)
|
||||
.map(|i| {
|
||||
let leaf_index: u8 = generate_value(&mut seed);
|
||||
let key = RpoDigest::new([
|
||||
ONE,
|
||||
ONE,
|
||||
Felt::new(i),
|
||||
Felt::new(leaf_index as u64),
|
||||
]);
|
||||
let value = generate_word(&mut seed);
|
||||
(key, value)
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut leaves: Vec<_> = entries
|
||||
.iter()
|
||||
.map(|(key, value)| {
|
||||
let leaf = SmtLeaf::new_single(*key, *value);
|
||||
let col = NodeIndex::from(leaf.index()).value();
|
||||
let hash = leaf.hash();
|
||||
SubtreeLeaf { col, hash }
|
||||
})
|
||||
.collect();
|
||||
leaves.sort();
|
||||
leaves
|
||||
},
|
||||
|leaves| {
|
||||
let (subtree, _) = build_subtree_for_bench(
|
||||
hint::black_box(leaves),
|
||||
hint::black_box(SMT_DEPTH),
|
||||
hint::black_box(SMT_DEPTH),
|
||||
);
|
||||
assert!(!subtree.is_empty());
|
||||
},
|
||||
BatchSize::SmallInput,
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
criterion_group! {
|
||||
name = smt_subtree_group;
|
||||
config = Criterion::default()
|
||||
.measurement_time(Duration::from_secs(40))
|
||||
.sample_size(60)
|
||||
.configure_from_args();
|
||||
targets = smt_subtree_even, smt_subtree_random
|
||||
}
|
||||
criterion_main!(smt_subtree_group);
|
||||
|
||||
// HELPER FUNCTIONS
|
||||
// --------------------------------------------------------------------------------------------
|
||||
|
||||
fn generate_value<T: Copy + Debug + Randomizable>(seed: &mut [u8; 32]) -> T {
|
||||
mem::swap(seed, &mut prng_array(*seed));
|
||||
let value: [T; 1] = rand_utils::prng_array(*seed);
|
||||
value[0]
|
||||
}
|
||||
|
||||
fn generate_word(seed: &mut [u8; 32]) -> Word {
|
||||
mem::swap(seed, &mut prng_array(*seed));
|
||||
let nums: [u64; 4] = prng_array(*seed);
|
||||
[Felt::new(nums[0]), Felt::new(nums[1]), Felt::new(nums[2]), Felt::new(nums[3])]
|
||||
}
|
71
benches/smt-with-entries.rs
Normal file
71
benches/smt-with-entries.rs
Normal file
|
@ -0,0 +1,71 @@
|
|||
use std::{fmt::Debug, hint, mem, time::Duration};
|
||||
|
||||
use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
|
||||
use miden_crypto::{hash::rpo::RpoDigest, merkle::Smt, Felt, Word, ONE};
|
||||
use rand_utils::prng_array;
|
||||
use winter_utils::Randomizable;
|
||||
|
||||
// 2^0, 2^4, 2^8, 2^12, 2^16
|
||||
const PAIR_COUNTS: [u64; 6] = [1, 16, 256, 4096, 65536, 1_048_576];
|
||||
|
||||
fn smt_with_entries(c: &mut Criterion) {
|
||||
let mut seed = [0u8; 32];
|
||||
|
||||
let mut group = c.benchmark_group("smt-with-entries");
|
||||
|
||||
for pair_count in PAIR_COUNTS {
|
||||
let bench_id = BenchmarkId::from_parameter(pair_count);
|
||||
group.bench_with_input(bench_id, &pair_count, |b, &pair_count| {
|
||||
b.iter_batched(
|
||||
|| {
|
||||
// Setup.
|
||||
prepare_entries(pair_count, &mut seed)
|
||||
},
|
||||
|entries| {
|
||||
// Benchmarked function.
|
||||
Smt::with_entries(hint::black_box(entries)).unwrap();
|
||||
},
|
||||
BatchSize::SmallInput,
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
criterion_group! {
|
||||
name = smt_with_entries_group;
|
||||
config = Criterion::default()
|
||||
//.measurement_time(Duration::from_secs(960))
|
||||
.measurement_time(Duration::from_secs(60))
|
||||
.sample_size(10)
|
||||
.configure_from_args();
|
||||
targets = smt_with_entries
|
||||
}
|
||||
criterion_main!(smt_with_entries_group);
|
||||
|
||||
// HELPER FUNCTIONS
|
||||
// --------------------------------------------------------------------------------------------
|
||||
|
||||
fn prepare_entries(pair_count: u64, seed: &mut [u8; 32]) -> Vec<(RpoDigest, [Felt; 4])> {
|
||||
let entries: Vec<(RpoDigest, Word)> = (0..pair_count)
|
||||
.map(|i| {
|
||||
let count = pair_count as f64;
|
||||
let idx = ((i as f64 / count) * (count)) as u64;
|
||||
let key = RpoDigest::new([generate_value(seed), ONE, Felt::new(i), Felt::new(idx)]);
|
||||
let value = generate_word(seed);
|
||||
(key, value)
|
||||
})
|
||||
.collect();
|
||||
entries
|
||||
}
|
||||
|
||||
fn generate_value<T: Copy + Debug + Randomizable>(seed: &mut [u8; 32]) -> T {
|
||||
mem::swap(seed, &mut prng_array(*seed));
|
||||
let value: [T; 1] = rand_utils::prng_array(*seed);
|
||||
value[0]
|
||||
}
|
||||
|
||||
fn generate_word(seed: &mut [u8; 32]) -> Word {
|
||||
mem::swap(seed, &mut prng_array(*seed));
|
||||
let nums: [u64; 4] = prng_array(*seed);
|
||||
[Felt::new(nums[0]), Felt::new(nums[1]), Felt::new(nums[2]), Felt::new(nums[3])]
|
||||
}
|
|
@ -97,6 +97,14 @@ impl NodeIndex {
|
|||
self
|
||||
}
|
||||
|
||||
/// Returns the parent of the current node. This is the same as [`Self::move_up()`], but returns
|
||||
/// a new value instead of mutating `self`.
|
||||
pub const fn parent(mut self) -> Self {
|
||||
self.depth = self.depth.saturating_sub(1);
|
||||
self.value >>= 1;
|
||||
self
|
||||
}
|
||||
|
||||
// PROVIDERS
|
||||
// --------------------------------------------------------------------------------------------
|
||||
|
||||
|
|
|
@ -21,9 +21,11 @@ mod path;
|
|||
pub use path::{MerklePath, RootPath, ValuePath};
|
||||
|
||||
mod smt;
|
||||
#[cfg(feature = "internal")]
|
||||
pub use smt::build_subtree_for_bench;
|
||||
pub use smt::{
|
||||
LeafIndex, MutationSet, SimpleSmt, Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError,
|
||||
SMT_DEPTH, SMT_MAX_DEPTH, SMT_MIN_DEPTH,
|
||||
SubtreeLeaf, SMT_DEPTH, SMT_MAX_DEPTH, SMT_MIN_DEPTH,
|
||||
};
|
||||
|
||||
mod mmr;
|
||||
|
|
|
@ -71,12 +71,51 @@ impl Smt {
|
|||
|
||||
/// Returns a new [Smt] instantiated with leaves set as specified by the provided entries.
|
||||
///
|
||||
/// If the `concurrent` feature is enabled, this function uses a parallel implementation to
|
||||
/// process the entries efficiently, otherwise it defaults to the sequential implementation.
|
||||
///
|
||||
/// All leaves omitted from the entries list are set to [Self::EMPTY_VALUE].
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if the provided entries contain multiple values for the same key.
|
||||
pub fn with_entries(
|
||||
entries: impl IntoIterator<Item = (RpoDigest, Word)>,
|
||||
) -> Result<Self, MerkleError> {
|
||||
#[cfg(feature = "concurrent")]
|
||||
{
|
||||
let mut seen_keys = BTreeSet::new();
|
||||
let entries: Vec<_> = entries
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
if seen_keys.insert(key) {
|
||||
Ok((key, value))
|
||||
} else {
|
||||
Err(MerkleError::DuplicateValuesForIndex(
|
||||
LeafIndex::<SMT_DEPTH>::from(key).value(),
|
||||
))
|
||||
}
|
||||
})
|
||||
.collect::<Result<_, _>>()?;
|
||||
if entries.is_empty() {
|
||||
return Ok(Self::default());
|
||||
}
|
||||
<Self as SparseMerkleTree<SMT_DEPTH>>::with_entries_par(entries)
|
||||
}
|
||||
#[cfg(not(feature = "concurrent"))]
|
||||
{
|
||||
Self::with_entries_sequential(entries)
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a new [Smt] instantiated with leaves set as specified by the provided entries.
|
||||
///
|
||||
/// This sequential implementation processes entries one at a time to build the tree.
|
||||
/// All leaves omitted from the entries list are set to [Self::EMPTY_VALUE].
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if the provided entries contain multiple values for the same key.
|
||||
pub fn with_entries_sequential(
|
||||
entries: impl IntoIterator<Item = (RpoDigest, Word)>,
|
||||
) -> Result<Self, MerkleError> {
|
||||
// create an empty tree
|
||||
let mut tree = Self::new();
|
||||
|
@ -101,6 +140,23 @@ impl Smt {
|
|||
Ok(tree)
|
||||
}
|
||||
|
||||
/// Returns a new [`Smt`] instantiated from already computed leaves and nodes.
|
||||
///
|
||||
/// This function performs minimal consistency checking. It is the caller's responsibility to
|
||||
/// ensure the passed arguments are correct and consistent with each other.
|
||||
///
|
||||
/// # Panics
|
||||
/// With debug assertions on, this function panics if `root` does not match the root node in
|
||||
/// `inner_nodes`.
|
||||
pub fn from_raw_parts(
|
||||
inner_nodes: BTreeMap<NodeIndex, InnerNode>,
|
||||
leaves: BTreeMap<u64, SmtLeaf>,
|
||||
root: RpoDigest,
|
||||
) -> Self {
|
||||
// Our particular implementation of `from_raw_parts()` never returns `Err`.
|
||||
<Self as SparseMerkleTree<SMT_DEPTH>>::from_raw_parts(inner_nodes, leaves, root).unwrap()
|
||||
}
|
||||
|
||||
// PUBLIC ACCESSORS
|
||||
// --------------------------------------------------------------------------------------------
|
||||
|
||||
|
@ -260,6 +316,19 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
|
|||
const EMPTY_VALUE: Self::Value = EMPTY_WORD;
|
||||
const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
|
||||
|
||||
fn from_raw_parts(
|
||||
inner_nodes: BTreeMap<NodeIndex, InnerNode>,
|
||||
leaves: BTreeMap<u64, SmtLeaf>,
|
||||
root: RpoDigest,
|
||||
) -> Result<Self, MerkleError> {
|
||||
if cfg!(debug_assertions) {
|
||||
let root_node = inner_nodes.get(&NodeIndex::root()).unwrap();
|
||||
assert_eq!(root_node.hash(), root);
|
||||
}
|
||||
|
||||
Ok(Self { root, inner_nodes, leaves })
|
||||
}
|
||||
|
||||
fn root(&self) -> RpoDigest {
|
||||
self.root
|
||||
}
|
||||
|
@ -344,6 +413,23 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
|
|||
fn path_and_leaf_to_opening(path: MerklePath, leaf: SmtLeaf) -> SmtProof {
|
||||
SmtProof::new_unchecked(path, leaf)
|
||||
}
|
||||
|
||||
fn pairs_to_leaf(mut pairs: Vec<(RpoDigest, Word)>) -> SmtLeaf {
|
||||
assert!(!pairs.is_empty());
|
||||
|
||||
if pairs.len() > 1 {
|
||||
SmtLeaf::new_multiple(pairs).unwrap()
|
||||
} else {
|
||||
let (key, value) = pairs.pop().unwrap();
|
||||
// TODO: should we ever be constructing empty leaves from pairs?
|
||||
if value == Self::EMPTY_VALUE {
|
||||
let index = Self::key_to_leaf_index(&key);
|
||||
SmtLeaf::new_empty(index)
|
||||
} else {
|
||||
SmtLeaf::new_single(key, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for Smt {
|
||||
|
|
|
@ -1,4 +1,7 @@
|
|||
use alloc::{collections::BTreeMap, vec::Vec};
|
||||
use core::mem;
|
||||
|
||||
use num::Integer;
|
||||
|
||||
use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex};
|
||||
use crate::{
|
||||
|
@ -62,6 +65,17 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
|
|||
// PROVIDED METHODS
|
||||
// ---------------------------------------------------------------------------------------------
|
||||
|
||||
/// Creates a new sparse Merkle tree from an existing set of key-value pairs, in parallel.
|
||||
#[cfg(feature = "concurrent")]
|
||||
fn with_entries_par(entries: Vec<(Self::Key, Self::Value)>) -> Result<Self, MerkleError>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
let (inner_nodes, leaves) = Self::build_subtrees(entries);
|
||||
let root = inner_nodes.get(&NodeIndex::root()).unwrap().hash();
|
||||
Self::from_raw_parts(inner_nodes, leaves, root)
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
|
@ -292,6 +306,16 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
|
|||
// REQUIRED METHODS
|
||||
// ---------------------------------------------------------------------------------------------
|
||||
|
||||
/// Construct this type from already computed leaves and nodes. The caller ensures passed
|
||||
/// arguments are correct and consistent with each other.
|
||||
fn from_raw_parts(
|
||||
inner_nodes: BTreeMap<NodeIndex, InnerNode>,
|
||||
leaves: BTreeMap<u64, Self::Leaf>,
|
||||
root: RpoDigest,
|
||||
) -> Result<Self, MerkleError>
|
||||
where
|
||||
Self: Sized;
|
||||
|
||||
/// The root of the tree
|
||||
fn root(&self) -> RpoDigest;
|
||||
|
||||
|
@ -341,18 +365,137 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
|
|||
/// 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;
|
||||
|
||||
/// Performs the initial transforms for constructing a [`SparseMerkleTree`] by composing
|
||||
/// subtrees. In other words, this function takes the key-value inputs to the tree, and produces
|
||||
/// the inputs to feed into [`build_subtree()`].
|
||||
///
|
||||
/// `pairs` *must* already be sorted **by leaf index column**, not simply sorted by key. If
|
||||
/// `pairs` is not correctly sorted, the returned computations will be incorrect.
|
||||
///
|
||||
/// # Panics
|
||||
/// With debug assertions on, this function panics if it detects that `pairs` is not correctly
|
||||
/// sorted. Without debug assertions, the returned computations will be incorrect.
|
||||
fn sorted_pairs_to_leaves(
|
||||
pairs: Vec<(Self::Key, Self::Value)>,
|
||||
) -> PairComputations<u64, Self::Leaf> {
|
||||
debug_assert!(pairs.is_sorted_by_key(|(key, _)| Self::key_to_leaf_index(key).value()));
|
||||
|
||||
let mut accumulator: PairComputations<u64, Self::Leaf> = Default::default();
|
||||
let mut accumulated_leaves: Vec<SubtreeLeaf> = Vec::with_capacity(pairs.len() / 2);
|
||||
|
||||
// As we iterate, we'll keep track of the kv-pairs we've seen so far that correspond to a
|
||||
// single leaf. When we see a pair that's in a different leaf, we'll swap these pairs
|
||||
// out and store them in our accumulated leaves.
|
||||
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);
|
||||
accumulated_leaves.push(SubtreeLeaf { col, hash });
|
||||
|
||||
debug_assert!(current_leaf_buffer.is_empty());
|
||||
}
|
||||
|
||||
// TODO: determine is there is any notable performance difference between computing
|
||||
// subtree boundaries after the fact as an iterator adapter (like this), versus computing
|
||||
// subtree boundaries as we go. Either way this function is only used at the beginning of a
|
||||
// parallel construction, so it should not be a critical path.
|
||||
accumulator.leaves = SubtreeLeavesIter::from_leaves(&mut accumulated_leaves).collect();
|
||||
accumulator
|
||||
}
|
||||
|
||||
/// Computes the raw parts for a new sparse Merkle tree from a set of key-value pairs.
|
||||
///
|
||||
/// `entries` need not be sorted. This function will sort them.
|
||||
#[cfg(feature = "concurrent")]
|
||||
fn build_subtrees(
|
||||
mut entries: Vec<(Self::Key, Self::Value)>,
|
||||
) -> (BTreeMap<NodeIndex, InnerNode>, BTreeMap<u64, Self::Leaf>) {
|
||||
entries.sort_by_key(|item| {
|
||||
let index = Self::key_to_leaf_index(&item.0);
|
||||
index.value()
|
||||
});
|
||||
Self::build_subtrees_from_sorted_entries(entries)
|
||||
}
|
||||
|
||||
/// Computes the raw parts for a new sparse Merkle tree from a set of key-value pairs.
|
||||
///
|
||||
/// This function is mostly an implementation detail of
|
||||
/// [`SparseMerkleTree::with_entries_par()`].
|
||||
#[cfg(feature = "concurrent")]
|
||||
fn build_subtrees_from_sorted_entries(
|
||||
entries: Vec<(Self::Key, Self::Value)>,
|
||||
) -> (BTreeMap<NodeIndex, InnerNode>, BTreeMap<u64, Self::Leaf>) {
|
||||
use rayon::prelude::*;
|
||||
|
||||
let mut accumulated_nodes: BTreeMap<NodeIndex, InnerNode> = Default::default();
|
||||
|
||||
let PairComputations {
|
||||
leaves: mut leaf_subtrees,
|
||||
nodes: initial_leaves,
|
||||
} = Self::sorted_pairs_to_leaves(entries);
|
||||
|
||||
for current_depth in (SUBTREE_DEPTH..=DEPTH).step_by(SUBTREE_DEPTH as usize).rev() {
|
||||
let (nodes, mut subtree_roots): (Vec<BTreeMap<_, _>>, Vec<SubtreeLeaf>) = leaf_subtrees
|
||||
.into_par_iter()
|
||||
.map(|subtree| {
|
||||
debug_assert!(subtree.is_sorted());
|
||||
debug_assert!(!subtree.is_empty());
|
||||
|
||||
let (nodes, subtree_root) = build_subtree(subtree, DEPTH, current_depth);
|
||||
(nodes, subtree_root)
|
||||
})
|
||||
.unzip();
|
||||
|
||||
leaf_subtrees = SubtreeLeavesIter::from_leaves(&mut subtree_roots).collect();
|
||||
accumulated_nodes.extend(nodes.into_iter().flatten());
|
||||
|
||||
debug_assert!(!leaf_subtrees.is_empty());
|
||||
}
|
||||
(accumulated_nodes, initial_leaves)
|
||||
}
|
||||
}
|
||||
|
||||
// INNER NODE
|
||||
// ================================================================================================
|
||||
|
||||
/// This struct is public so functions returning it can be used in `benches/`, but is otherwise not
|
||||
/// part of the public API.
|
||||
#[doc(hidden)]
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
|
||||
pub(crate) struct InnerNode {
|
||||
pub struct InnerNode {
|
||||
pub left: RpoDigest,
|
||||
pub right: RpoDigest,
|
||||
}
|
||||
|
@ -462,3 +605,197 @@ impl<const DEPTH: u8, K, V> MutationSet<DEPTH, K, V> {
|
|||
self.new_root
|
||||
}
|
||||
}
|
||||
|
||||
// SUBTREES
|
||||
// ================================================================================================
|
||||
/// A subtree is of depth 8.
|
||||
const SUBTREE_DEPTH: u8 = 8;
|
||||
|
||||
/// A depth-8 subtree contains 256 "columns" that can possibly be occupied.
|
||||
const COLS_PER_SUBTREE: u64 = u64::pow(2, SUBTREE_DEPTH as u32);
|
||||
|
||||
/// Helper struct for organizing the data we care about when computing Merkle subtrees.
|
||||
///
|
||||
/// Note that these represet "conceptual" leaves of some subtree, not necessarily
|
||||
/// the leaf type for the sparse Merkle tree.
|
||||
#[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,
|
||||
}
|
||||
|
||||
/// 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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct SubtreeLeavesIter<'s> {
|
||||
leaves: core::iter::Peekable<alloc::vec::Drain<'s, SubtreeLeaf>>,
|
||||
}
|
||||
impl<'s> SubtreeLeavesIter<'s> {
|
||||
fn from_leaves(leaves: &'s mut Vec<SubtreeLeaf>) -> Self {
|
||||
// TODO: determine if there is any notable performance difference between taking a Vec,
|
||||
// which many need flattening first, vs storing a `Box<dyn Iterator<Item = SubtreeLeaf>>`.
|
||||
// The latter may have self-referential properties that are impossible to express in purely
|
||||
// safe Rust Rust.
|
||||
Self { leaves: leaves.drain(..).peekable() }
|
||||
}
|
||||
}
|
||||
impl core::iter::Iterator for SubtreeLeavesIter<'_> {
|
||||
type Item = Vec<SubtreeLeaf>;
|
||||
|
||||
/// Each `next()` collects an entire subtree.
|
||||
fn next(&mut self) -> Option<Vec<SubtreeLeaf>> {
|
||||
let mut subtree: Vec<SubtreeLeaf> = Default::default();
|
||||
|
||||
let mut last_subtree_col = 0;
|
||||
|
||||
while let Some(leaf) = self.leaves.peek() {
|
||||
last_subtree_col = u64::max(1, last_subtree_col);
|
||||
let is_exact_multiple = Integer::is_multiple_of(&last_subtree_col, &COLS_PER_SUBTREE);
|
||||
let next_subtree_col = if is_exact_multiple {
|
||||
u64::next_multiple_of(last_subtree_col + 1, COLS_PER_SUBTREE)
|
||||
} else {
|
||||
last_subtree_col.next_multiple_of(COLS_PER_SUBTREE)
|
||||
};
|
||||
|
||||
last_subtree_col = leaf.col;
|
||||
if leaf.col < next_subtree_col {
|
||||
subtree.push(self.leaves.next().unwrap());
|
||||
} else if subtree.is_empty() {
|
||||
continue;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if subtree.is_empty() {
|
||||
debug_assert!(self.leaves.peek().is_none());
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(subtree)
|
||||
}
|
||||
}
|
||||
|
||||
// HELPER FUNCTIONS
|
||||
// ================================================================================================
|
||||
|
||||
/// Builds Merkle nodes from a bottom layer of "leaves" -- represented by a horizontal index and
|
||||
/// the hash of the leaf at that index. `leaves` *must* be sorted by horizontal index, and
|
||||
/// `leaves` must not contain more than one depth-8 subtree's worth of leaves.
|
||||
///
|
||||
/// This function will then calculate the inner nodes above each leaf for 8 layers, as well as
|
||||
/// the "leaves" for the next 8-deep subtree, so this function can effectively be chained into
|
||||
/// itself.
|
||||
///
|
||||
/// # Panics
|
||||
/// With debug assertions on, this function panics under invalid inputs: if `leaves` contains
|
||||
/// more entries than can fit in a depth-8 subtree, if `leaves` contains leaves belonging to
|
||||
/// different depth-8 subtrees, if `bottom_depth` is lower in the tree than the specified
|
||||
/// maximum depth (`DEPTH`), or if `leaves` is not sorted.
|
||||
fn build_subtree(
|
||||
mut leaves: Vec<SubtreeLeaf>,
|
||||
tree_depth: u8,
|
||||
bottom_depth: u8,
|
||||
) -> (BTreeMap<NodeIndex, InnerNode>, SubtreeLeaf) {
|
||||
debug_assert!(bottom_depth <= tree_depth);
|
||||
debug_assert!(Integer::is_multiple_of(&bottom_depth, &SUBTREE_DEPTH));
|
||||
debug_assert!(leaves.len() <= usize::pow(2, SUBTREE_DEPTH as u32));
|
||||
let subtree_root = bottom_depth - SUBTREE_DEPTH;
|
||||
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(tree_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(tree_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(tree_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);
|
||||
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);
|
||||
}
|
||||
debug_assert_eq!(leaves.len(), 1);
|
||||
let root = leaves.pop().unwrap();
|
||||
(inner_nodes, root)
|
||||
}
|
||||
|
||||
#[cfg(feature = "internal")]
|
||||
pub fn build_subtree_for_bench(
|
||||
leaves: Vec<SubtreeLeaf>,
|
||||
tree_depth: u8,
|
||||
bottom_depth: u8,
|
||||
) -> (BTreeMap<NodeIndex, InnerNode>, SubtreeLeaf) {
|
||||
build_subtree(leaves, tree_depth, bottom_depth)
|
||||
}
|
||||
|
||||
// TESTS
|
||||
// ================================================================================================
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
|
|
@ -1,4 +1,7 @@
|
|||
use alloc::collections::{BTreeMap, BTreeSet};
|
||||
use alloc::{
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
vec::Vec,
|
||||
};
|
||||
|
||||
use super::{
|
||||
super::ValuePath, EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, MerkleError,
|
||||
|
@ -97,6 +100,23 @@ impl<const DEPTH: u8> SimpleSmt<DEPTH> {
|
|||
Ok(tree)
|
||||
}
|
||||
|
||||
/// Returns a new [`SimpleSmt`] instantiated from already computed leaves and nodes.
|
||||
///
|
||||
/// This function performs minimal consistency checking. It is the caller's responsibility to
|
||||
/// ensure the passed arguments are correct and consistent with each other.
|
||||
///
|
||||
/// # Panics
|
||||
/// With debug assertions on, this function panics if `root` does not match the root node in
|
||||
/// `inner_nodes`.
|
||||
pub fn from_raw_parts(
|
||||
inner_nodes: BTreeMap<NodeIndex, InnerNode>,
|
||||
leaves: BTreeMap<u64, Word>,
|
||||
root: RpoDigest,
|
||||
) -> Self {
|
||||
// Our particular implementation of `from_raw_parts()` never returns `Err`.
|
||||
<Self as SparseMerkleTree<DEPTH>>::from_raw_parts(inner_nodes, leaves, root).unwrap()
|
||||
}
|
||||
|
||||
/// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
|
||||
/// starting at index 0.
|
||||
pub fn with_contiguous_leaves(
|
||||
|
@ -306,6 +326,19 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
|
|||
const EMPTY_VALUE: Self::Value = EMPTY_WORD;
|
||||
const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(DEPTH, 0);
|
||||
|
||||
fn from_raw_parts(
|
||||
inner_nodes: BTreeMap<NodeIndex, InnerNode>,
|
||||
leaves: BTreeMap<u64, Word>,
|
||||
root: RpoDigest,
|
||||
) -> Result<Self, MerkleError> {
|
||||
if cfg!(debug_assertions) {
|
||||
let root_node = inner_nodes.get(&NodeIndex::root()).unwrap();
|
||||
assert_eq!(root_node.hash(), root);
|
||||
}
|
||||
|
||||
Ok(Self { root, inner_nodes, leaves })
|
||||
}
|
||||
|
||||
fn root(&self) -> RpoDigest {
|
||||
self.root
|
||||
}
|
||||
|
@ -370,4 +403,11 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
|
|||
fn path_and_leaf_to_opening(path: MerklePath, leaf: Word) -> ValuePath {
|
||||
(path, leaf).into()
|
||||
}
|
||||
|
||||
fn pairs_to_leaf(mut pairs: Vec<(LeafIndex<DEPTH>, Word)>) -> Word {
|
||||
// SimpleSmt can't have more than one value per key.
|
||||
assert_eq!(pairs.len(), 1);
|
||||
let (_key, value) = pairs.pop().unwrap();
|
||||
value
|
||||
}
|
||||
}
|
||||
|
|
417
src/merkle/smt/tests.rs
Normal file
417
src/merkle/smt/tests.rs
Normal file
|
@ -0,0 +1,417 @@
|
|||
use alloc::{collections::BTreeMap, vec::Vec};
|
||||
|
||||
use super::{
|
||||
build_subtree, InnerNode, LeafIndex, NodeIndex, PairComputations, SmtLeaf, SparseMerkleTree,
|
||||
SubtreeLeaf, SubtreeLeavesIter, COLS_PER_SUBTREE, SUBTREE_DEPTH,
|
||||
};
|
||||
use crate::{
|
||||
hash::rpo::RpoDigest,
|
||||
merkle::{Smt, SMT_DEPTH},
|
||||
Felt, Word, ONE,
|
||||
};
|
||||
|
||||
fn smtleaf_to_subtree_leaf(leaf: &SmtLeaf) -> SubtreeLeaf {
|
||||
SubtreeLeaf {
|
||||
col: leaf.index().index.value(),
|
||||
hash: leaf.hash(),
|
||||
}
|
||||
}
|
||||
|
||||
#[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_sequential(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(smtleaf_to_subtree_leaf).collect())
|
||||
.to_vec();
|
||||
assert_eq!(control_leaves_iter.next(), None);
|
||||
control_subtree_leaves
|
||||
};
|
||||
|
||||
let subtrees: PairComputations<u64, SmtLeaf> = 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);
|
||||
|
||||
// Flattening and re-separating out the leaves into subtrees should have the same result.
|
||||
let mut all_leaves: Vec<SubtreeLeaf> = subtrees.leaves.clone().into_iter().flatten().collect();
|
||||
let re_grouped: Vec<Vec<_>> = SubtreeLeavesIter::from_leaves(&mut all_leaves).collect();
|
||||
assert_eq!(subtrees.leaves, re_grouped);
|
||||
|
||||
// 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 = COLS_PER_SUBTREE;
|
||||
|
||||
let entries = generate_entries(PAIR_COUNT);
|
||||
|
||||
let control = Smt::with_entries_sequential(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, subtree_root) = build_subtree(leaves, SMT_DEPTH, SMT_DEPTH);
|
||||
assert!(!first_subtree.is_empty());
|
||||
|
||||
// The inner nodes computed from that subtree should match the nodes in our control tree.
|
||||
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",
|
||||
);
|
||||
}
|
||||
|
||||
// The root returned should also match the equivalent node in the control tree.
|
||||
let control_root_index =
|
||||
NodeIndex::new(SMT_DEPTH - SUBTREE_DEPTH, subtree_root.col).expect("Valid root index");
|
||||
let control_root_node = control.get_inner_node(control_root_index);
|
||||
let control_hash = control_root_node.hash();
|
||||
assert_eq!(
|
||||
control_hash, subtree_root.hash,
|
||||
"Subtree-computed root at index {control_root_index:?} does not match control"
|
||||
);
|
||||
}
|
||||
|
||||
// Test that not just can we compute a subtree correctly, but we can feed the results of one
|
||||
// subtree into computing another. In other words, test that `build_subtree()` is correctly
|
||||
// composable.
|
||||
#[test]
|
||||
fn test_two_subtrees() {
|
||||
// Two subtrees' worth of leaves.
|
||||
const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 2;
|
||||
|
||||
let entries = generate_entries(PAIR_COUNT);
|
||||
|
||||
let control = Smt::with_entries_sequential(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]: [Vec<_>; 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, first_root) = build_subtree(first, SMT_DEPTH, current_depth);
|
||||
next_leaves.push(first_root);
|
||||
|
||||
let (second_nodes, second_root) = build_subtree(second, SMT_DEPTH, current_depth);
|
||||
next_leaves.push(second_root);
|
||||
|
||||
// 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 -= SUBTREE_DEPTH;
|
||||
|
||||
let (nodes, root_leaf) = build_subtree(next_leaves, SMT_DEPTH, current_depth);
|
||||
assert_eq!(nodes.len(), SUBTREE_DEPTH as usize);
|
||||
assert_eq!(root_leaf.col, 0);
|
||||
|
||||
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",
|
||||
);
|
||||
}
|
||||
|
||||
let index = NodeIndex::new(current_depth - SUBTREE_DEPTH, root_leaf.col).unwrap();
|
||||
let control_root = control.get_inner_node(index).hash();
|
||||
assert_eq!(control_root, root_leaf.hash, "Root mismatch");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_singlethreaded_subtrees() {
|
||||
const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 64;
|
||||
|
||||
let entries = generate_entries(PAIR_COUNT);
|
||||
|
||||
let control = Smt::with_entries_sequential(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 (SUBTREE_DEPTH..=SMT_DEPTH).step_by(SUBTREE_DEPTH as usize).rev() {
|
||||
// There's no flat_map_unzip(), so this is the best we can do.
|
||||
let (nodes, mut subtree_roots): (Vec<BTreeMap<_, _>>, Vec<SubtreeLeaf>) = leaf_subtrees
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, subtree)| {
|
||||
// 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, subtree_root) = build_subtree(subtree, SMT_DEPTH, current_depth);
|
||||
|
||||
// Post-assertions.
|
||||
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,
|
||||
);
|
||||
}
|
||||
|
||||
(nodes, subtree_root)
|
||||
})
|
||||
.unzip();
|
||||
|
||||
// Update state between each depth iteration.
|
||||
|
||||
leaf_subtrees = SubtreeLeavesIter::from_leaves(&mut subtree_roots).collect();
|
||||
accumulated_nodes.extend(nodes.into_iter().flatten());
|
||||
|
||||
assert!(!leaf_subtrees.is_empty(), "on depth {current_depth}");
|
||||
}
|
||||
|
||||
// Make sure the true leaves match, first checking length and then checking 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, "test leaf at column {col} does not match control");
|
||||
}
|
||||
|
||||
// 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]: [Vec<_>; 1] = leaf_subtrees.try_into().unwrap();
|
||||
// which itself contains only one 'leaf'...
|
||||
let [root_leaf]: [SubtreeLeaf; 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);
|
||||
}
|
||||
|
||||
/// The parallel version of `test_singlethreaded_subtree()`.
|
||||
#[test]
|
||||
#[cfg(feature = "concurrent")]
|
||||
fn test_multithreaded_subtrees() {
|
||||
use rayon::prelude::*;
|
||||
|
||||
const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 64;
|
||||
|
||||
let entries = generate_entries(PAIR_COUNT);
|
||||
|
||||
let control = Smt::with_entries_sequential(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 (SUBTREE_DEPTH..=SMT_DEPTH).step_by(SUBTREE_DEPTH as usize).rev() {
|
||||
let (nodes, mut subtree_roots): (Vec<BTreeMap<_, _>>, Vec<SubtreeLeaf>) = leaf_subtrees
|
||||
.into_par_iter()
|
||||
.enumerate()
|
||||
.map(|(i, subtree)| {
|
||||
// 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!",
|
||||
);
|
||||
|
||||
let (nodes, subtree_root) = build_subtree(subtree, SMT_DEPTH, current_depth);
|
||||
|
||||
// Post-assertions.
|
||||
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,
|
||||
);
|
||||
}
|
||||
|
||||
(nodes, subtree_root)
|
||||
})
|
||||
.unzip();
|
||||
|
||||
leaf_subtrees = SubtreeLeavesIter::from_leaves(&mut subtree_roots).collect();
|
||||
accumulated_nodes.extend(nodes.into_iter().flatten());
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "concurrent")]
|
||||
fn test_with_entries_parallel() {
|
||||
const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 64;
|
||||
|
||||
let entries = generate_entries(PAIR_COUNT);
|
||||
|
||||
let control = Smt::with_entries_sequential(entries.clone()).unwrap();
|
||||
|
||||
let smt = Smt::with_entries(entries.clone()).unwrap();
|
||||
assert_eq!(smt.root(), control.root());
|
||||
assert_eq!(smt, control);
|
||||
}
|
Loading…
Add table
Reference in a new issue