[WIP] implement subtree-based SMT computations

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## 0.11.0 (2024-10-30)
 
 - [BREAKING] Updated Winterfell dependency to v0.10 (#338).
+- Added `Smt::with_entries_par()`, a parallel version of `with_entries()` with significantly better performance (#341).
 
 ## 0.11.0 (2024-10-17)
 

Cargo.toml

@@ -27,14 +27,28 @@ harness = false
 name = "smt"
 harness = false
 
+[[bench]]
+name = "smt-subtree"
+harness = false
+
+[[bench]]
+name = "merkle"
+harness = false
+
+[[bench]]
+name = "parallel-subtree"
+harness = false
+required-features = ["concurrent"]
+
 [[bench]]
 name = "store"
 harness = false
 
 [features]
-default = ["std"]
+default = ["std", "concurrent"]
 executable = ["dep:clap", "dep:rand-utils", "std"]
 serde = ["dep:serde", "serde?/alloc", "winter-math/serde"]
+concurrent = ["dep:rayon"]
 std = [
     "blake3/std",
     "dep:cc",
@@ -58,6 +72,7 @@ sha3 = { version = "0.10", default-features = false }
 winter-crypto = { version = "0.10", default-features = false }
 winter-math = { version = "0.10", default-features = false }
 winter-utils = { version = "0.10", default-features = false }
+rayon = { version = "1.10.0", optional = true }
 
 [dev-dependencies]
 criterion = { version = "0.5", features = ["html_reports"] }

benches/merkle.rs

@@ -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])]
+}

benches/parallel-subtree.rs

@@ -0,0 +1,75 @@
+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_parallel_subtree(c: &mut Criterion) {
+    let mut seed = [0u8; 32];
+
+    let mut group = c.benchmark_group("parallel-subtrees");
+
+    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 count = pair_count as f64;
+                            let idx = ((i as f64 / count) * (count)) as u64;
+                            let key = RpoDigest::new([
+                                generate_value(&mut seed),
+                                ONE,
+                                Felt::new(i),
+                                Felt::new(idx),
+                            ]);
+                            let value = generate_word(&mut seed);
+                            (key, value)
+                        })
+                        .collect();
+
+                    let control = Smt::with_entries(entries.clone()).unwrap();
+                    (entries, control)
+                },
+                |(entries, control)| {
+                    // Benchmarked function.
+                    let tree = Smt::with_entries_par(hint::black_box(entries)).unwrap();
+                    assert_eq!(tree.root(), control.root());
+                },
+                BatchSize::SmallInput,
+            );
+        });
+    }
+}
+
+criterion_group! {
+    name = smt_subtree_group;
+    config = Criterion::default()
+        //.measurement_time(Duration::from_secs(960))
+        .measurement_time(Duration::from_secs(60))
+        .sample_size(10)
+        .configure_from_args();
+    targets = smt_parallel_subtree
+}
+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])]
+}

benches/smt-subtree.rs

@@ -0,0 +1,136 @@
+use std::{fmt::Debug, hint, mem, time::Duration};
+
+use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
+use miden_crypto::{
+    hash::rpo::RpoDigest,
+    merkle::{NodeIndex, Smt, 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, _) =
+                        Smt::build_subtree(hint::black_box(leaves), 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, _) =
+                        Smt::build_subtree(hint::black_box(leaves), 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])]
+}

src/main.rs

@@ -33,7 +33,12 @@ pub fn benchmark_smt() {
         entries.push((key, value));
     }
 
-    let mut tree = construction(entries, tree_size).unwrap();
+    let mut tree = construction(entries.clone(), tree_size).unwrap();
+    #[cfg(feature = "concurrent")]
+    {
+        let parallel = parallel_construction(entries, tree_size).unwrap();
+        assert_eq!(tree, parallel);
+    }
     insertion(&mut tree, tree_size).unwrap();
     batched_insertion(&mut tree, tree_size).unwrap();
     proof_generation(&mut tree, tree_size).unwrap();
@@ -56,6 +61,27 @@ pub fn construction(entries: Vec<(RpoDigest, Word)>, size: u64) -> Result<Smt, M
     Ok(tree)
 }
 
+#[cfg(feature = "concurrent")]
+pub fn parallel_construction(
+    entries: Vec<(RpoDigest, Word)>,
+    size: u64,
+) -> Result<Smt, MerkleError> {
+    println!("Running a parallel construction benchmark:");
+    let now = Instant::now();
+
+    let tree = Smt::with_entries_par(entries).unwrap();
+
+    let elapsed = now.elapsed();
+    println!(
+        "Parallel-constructed an SMT with {} key-value pairs in {:.3} seconds",
+        size,
+        elapsed.as_secs_f32(),
+    );
+    println!("Number of leaf nodes: {}\n", tree.leaves().count());
+
+    Ok(tree)
+}
+
 /// Runs the insertion benchmark for the [`Smt`].
 pub fn insertion(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
     println!("Running an insertion benchmark:");

src/merkle/index.rs

@@ -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
     // --------------------------------------------------------------------------------------------
 

src/merkle/mod.rs

@@ -23,7 +23,7 @@ pub use path::{MerklePath, RootPath, ValuePath};
 mod smt;
 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;