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)
|
## 0.12.0 (2024-10-30)
|
||||||
|
|
||||||
- [BREAKING] Updated Winterfell dependency to v0.10 (#338).
|
- [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)
|
## 0.11.0 (2024-10-17)
|
||||||
|
|
||||||
|
|
1
Cargo.lock
generated
1
Cargo.lock
generated
|
@ -541,6 +541,7 @@ dependencies = [
|
||||||
"rand",
|
"rand",
|
||||||
"rand_chacha",
|
"rand_chacha",
|
||||||
"rand_core",
|
"rand_core",
|
||||||
|
"rayon",
|
||||||
"seq-macro",
|
"seq-macro",
|
||||||
"serde",
|
"serde",
|
||||||
"sha3",
|
"sha3",
|
||||||
|
|
18
Cargo.toml
18
Cargo.toml
|
@ -27,13 +27,28 @@ harness = false
|
||||||
name = "smt"
|
name = "smt"
|
||||||
harness = false
|
harness = false
|
||||||
|
|
||||||
|
[[bench]]
|
||||||
|
name = "smt-subtree"
|
||||||
|
harness = false
|
||||||
|
required-features = ["internal"]
|
||||||
|
|
||||||
|
[[bench]]
|
||||||
|
name = "merkle"
|
||||||
|
harness = false
|
||||||
|
|
||||||
|
[[bench]]
|
||||||
|
name = "smt-with-entries"
|
||||||
|
harness = false
|
||||||
|
|
||||||
[[bench]]
|
[[bench]]
|
||||||
name = "store"
|
name = "store"
|
||||||
harness = false
|
harness = false
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = ["std"]
|
concurrent = ["dep:rayon"]
|
||||||
|
default = ["std", "concurrent"]
|
||||||
executable = ["dep:clap", "dep:rand-utils", "std"]
|
executable = ["dep:clap", "dep:rand-utils", "std"]
|
||||||
|
internal = []
|
||||||
serde = ["dep:serde", "serde?/alloc", "winter-math/serde"]
|
serde = ["dep:serde", "serde?/alloc", "winter-math/serde"]
|
||||||
std = [
|
std = [
|
||||||
"blake3/std",
|
"blake3/std",
|
||||||
|
@ -53,6 +68,7 @@ num-complex = { version = "0.4", default-features = false }
|
||||||
rand = { version = "0.8", default-features = false }
|
rand = { version = "0.8", default-features = false }
|
||||||
rand_core = { version = "0.6", default-features = false }
|
rand_core = { version = "0.6", default-features = false }
|
||||||
rand-utils = { version = "0.11", package = "winter-rand-utils", optional = true }
|
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"] }
|
serde = { version = "1.0", default-features = false, optional = true, features = ["derive"] }
|
||||||
sha3 = { version = "0.10", default-features = false }
|
sha3 = { version = "0.10", default-features = false }
|
||||||
thiserror = { version = "2.0", 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
|
.PHONY: bench-tx
|
||||||
bench-tx: ## Run crypto benchmarks
|
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:
|
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.
|
- `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.
|
- `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:
|
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
|
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
|
// PROVIDERS
|
||||||
// --------------------------------------------------------------------------------------------
|
// --------------------------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
|
@ -21,9 +21,11 @@ mod path;
|
||||||
pub use path::{MerklePath, RootPath, ValuePath};
|
pub use path::{MerklePath, RootPath, ValuePath};
|
||||||
|
|
||||||
mod smt;
|
mod smt;
|
||||||
|
#[cfg(feature = "internal")]
|
||||||
|
pub use smt::build_subtree_for_bench;
|
||||||
pub use smt::{
|
pub use smt::{
|
||||||
LeafIndex, MutationSet, SimpleSmt, Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError,
|
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;
|
mod mmr;
|
||||||
|
|
|
@ -71,12 +71,51 @@ impl Smt {
|
||||||
|
|
||||||
/// Returns a new [Smt] instantiated with leaves set as specified by the provided entries.
|
/// 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].
|
/// All leaves omitted from the entries list are set to [Self::EMPTY_VALUE].
|
||||||
///
|
///
|
||||||
/// # Errors
|
/// # Errors
|
||||||
/// Returns an error if the provided entries contain multiple values for the same key.
|
/// Returns an error if the provided entries contain multiple values for the same key.
|
||||||
pub fn with_entries(
|
pub fn with_entries(
|
||||||
entries: impl IntoIterator<Item = (RpoDigest, Word)>,
|
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> {
|
) -> Result<Self, MerkleError> {
|
||||||
// create an empty tree
|
// create an empty tree
|
||||||
let mut tree = Self::new();
|
let mut tree = Self::new();
|
||||||
|
@ -101,6 +140,23 @@ impl Smt {
|
||||||
Ok(tree)
|
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
|
// PUBLIC ACCESSORS
|
||||||
// --------------------------------------------------------------------------------------------
|
// --------------------------------------------------------------------------------------------
|
||||||
|
|
||||||
|
@ -260,6 +316,19 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
|
||||||
const EMPTY_VALUE: Self::Value = EMPTY_WORD;
|
const EMPTY_VALUE: Self::Value = EMPTY_WORD;
|
||||||
const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
|
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 {
|
fn root(&self) -> RpoDigest {
|
||||||
self.root
|
self.root
|
||||||
}
|
}
|
||||||
|
@ -344,6 +413,23 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
|
||||||
fn path_and_leaf_to_opening(path: MerklePath, leaf: SmtLeaf) -> SmtProof {
|
fn path_and_leaf_to_opening(path: MerklePath, leaf: SmtLeaf) -> SmtProof {
|
||||||
SmtProof::new_unchecked(path, leaf)
|
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 {
|
impl Default for Smt {
|
||||||
|
|
|
@ -1,4 +1,7 @@
|
||||||
use alloc::{collections::BTreeMap, vec::Vec};
|
use alloc::{collections::BTreeMap, vec::Vec};
|
||||||
|
use core::mem;
|
||||||
|
|
||||||
|
use num::Integer;
|
||||||
|
|
||||||
use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex};
|
use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex};
|
||||||
use crate::{
|
use crate::{
|
||||||
|
@ -62,6 +65,17 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
|
||||||
// PROVIDED METHODS
|
// 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
|
/// 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.
|
/// path to the leaf, as well as the leaf itself.
|
||||||
fn open(&self, key: &Self::Key) -> Self::Opening {
|
fn open(&self, key: &Self::Key) -> Self::Opening {
|
||||||
|
@ -292,6 +306,16 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
|
||||||
// REQUIRED METHODS
|
// 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
|
/// The root of the tree
|
||||||
fn root(&self) -> RpoDigest;
|
fn root(&self) -> RpoDigest;
|
||||||
|
|
||||||
|
@ -341,18 +365,137 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
|
||||||
/// Maps a key to a leaf index
|
/// Maps a key to a leaf index
|
||||||
fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
|
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.
|
/// Maps a (MerklePath, Self::Leaf) to an opening.
|
||||||
///
|
///
|
||||||
/// The length `path` is guaranteed to be equal to `DEPTH`
|
/// The length `path` is guaranteed to be equal to `DEPTH`
|
||||||
fn path_and_leaf_to_opening(path: MerklePath, leaf: Self::Leaf) -> Self::Opening;
|
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
|
// 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)]
|
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||||
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
|
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
|
||||||
pub(crate) struct InnerNode {
|
pub struct InnerNode {
|
||||||
pub left: RpoDigest,
|
pub left: RpoDigest,
|
||||||
pub right: RpoDigest,
|
pub right: RpoDigest,
|
||||||
}
|
}
|
||||||
|
@ -462,3 +605,197 @@ impl<const DEPTH: u8, K, V> MutationSet<DEPTH, K, V> {
|
||||||
self.new_root
|
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::{
|
use super::{
|
||||||
super::ValuePath, EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, MerkleError,
|
super::ValuePath, EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, MerkleError,
|
||||||
|
@ -97,6 +100,23 @@ impl<const DEPTH: u8> SimpleSmt<DEPTH> {
|
||||||
Ok(tree)
|
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
|
/// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
|
||||||
/// starting at index 0.
|
/// starting at index 0.
|
||||||
pub fn with_contiguous_leaves(
|
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_VALUE: Self::Value = EMPTY_WORD;
|
||||||
const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(DEPTH, 0);
|
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 {
|
fn root(&self) -> RpoDigest {
|
||||||
self.root
|
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 {
|
fn path_and_leaf_to_opening(path: MerklePath, leaf: Word) -> ValuePath {
|
||||||
(path, leaf).into()
|
(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