feat: implement a persistent backend for the Large SMT Forest (#876)

This commit implements a persistent backend for the SMT forest that both
allows it to start up rapidly from an on-disk state, and allows the
offloading of many portions of the forest from resident memory.

It stores the full tree data for each lineage in the forest in a RocksDB
instance. As it stands, the backend has not undergone any particular
optimization work, instead predominantly relying on the optimizations to
access patterns developed for `LargeSmt`.

Comparative performance analysis has been performed against Large SMT,
and we have found that in like-for-like scenarios, the forest ranges
from 1.5x to 2x slower than `LargeSmt`. This is perfectly in line with
estimations, as in the worst case the forest has to perform 2x the
amount of I/O due to its lack of an in-memory prefix.

The commit also includes basic benchmarks for the large SMT forest.
While they do not cover every single piece of functionality, they
currently cover the following, specifically for the persistent backend:

- `forest.open(...)`: The time it takes to get a single opening from
  some arbitrary tree in the forest, both for the current tree and from
  the history.
- `forest.add_lineage(...)`: The time it takes to add a new lineage to
  the forest.
- `forest.update_tree(...)`: The time it takes to update an existing
  lineage in the forest.
- `forest.update_forest(...)`: The time it takes to update multiple
  lineages in the forest in a single batch.

There may be further opportunities for optimization, based around
tailoring the database parameters better for the forest, but the current
performance is well-within the expected bounds.

Author: iamrecursion

miden-crypto/Cargo.toml

@@ -46,9 +46,14 @@ required-features = ["std"]
 
 [[bench]]
 harness           = false
-name              = "large-smt"
+name              = "large_smt"
 required-features = ["rocksdb"]
 
+[[bench]]
+harness           = false
+name              = "large_smt_forest"
+required-features = ["persistent-forest"]
+
 [[bench]]
 harness           = false
 name              = "store"
@@ -75,15 +80,16 @@ name              = "transpose"
 required-features = ["std"]
 
 [features]
-concurrent = ["dep:rayon", "p3-maybe-rayon/parallel", "p3-miden-prover/parallel", "p3-util/parallel", "std"]
-default    = ["concurrent", "std"]
-executable = ["concurrent", "dep:clap", "dep:rand-utils"]
-fuzzing    = []
-internal   = ["concurrent"]
-rocksdb    = ["concurrent", "dep:rocksdb"]
-serde      = ["dep:serde", "serde?/alloc"]
-std        = ["blake3/std", "dep:cc", "miden-serde-utils/std", "rand/std", "rand/thread_rng", "serde?/std"]
-testing    = ["dep:proptest", "miden-field/testing"]
+concurrent        = ["dep:rayon", "p3-maybe-rayon/parallel", "p3-miden-prover/parallel", "p3-util/parallel", "std"]
+default           = ["concurrent", "std"]
+executable        = ["concurrent", "dep:clap", "dep:rand-utils"]
+fuzzing           = []
+internal          = ["concurrent"]
+persistent-forest = ["rocksdb", "serde"]
+rocksdb           = ["concurrent", "dep:rocksdb"]
+serde             = ["dep:serde", "serde?/alloc"]
+std               = ["blake3/std", "dep:cc", "miden-serde-utils/std", "rand/std", "rand/thread_rng", "serde?/std"]
+testing           = ["dep:proptest", "miden-field/testing"]
 
 [dependencies]
 # Miden dependencies

miden-crypto/benches/large-smt.rs miden-crypto/benches/large_smt.rsmiden-crypto/benches/large-smt.rs renamed to miden-crypto/benches/large_smt.rs

@@ -1,6 +1,6 @@
-use std::hint;
+use std::{hint, iter::empty};
 
-use criterion::{Criterion, criterion_group, criterion_main};
+use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
 use miden_crypto::{
     Felt, Word,
     merkle::{
@@ -13,11 +13,10 @@ use miden_crypto::{
 };
 
 mod common;
-use common::*;
 
-use crate::{
-    common::data::{generate_smt_entries_sequential, generate_test_keys_sequential},
+use crate::common::{
     config::{DEFAULT_MEASUREMENT_TIME, DEFAULT_SAMPLE_SIZE},
+    data::{generate_smt_entries_sequential, generate_test_keys_sequential},
 };
 
 // SUBTREE SERIALIZATION BENCHMARKS
@@ -138,6 +137,28 @@ benchmark_with_setup_data! {
     },
 }
 
+benchmark_with_setup_data! {
+    large_smt_open_in_large_tree,
+    DEFAULT_MEASUREMENT_TIME,
+    DEFAULT_SAMPLE_SIZE,
+    "rocksdb_smt_open_in_large_tree",
+    || {
+        let entries = generate_smt_entries_sequential(10_000);
+        let keys = generate_test_keys_sequential(10);
+        let temp_dir = tempfile::TempDir::new().unwrap();
+        let storage = RocksDbStorage::open(RocksDbConfig::new(temp_dir.path())).unwrap();
+        let smt = LargeSmt::with_entries(storage, entries).unwrap();
+        (smt, keys, temp_dir)
+    },
+    |b: &mut criterion::Bencher, (smt, keys, _temp_dir): &(LargeSmt<RocksDbStorage>, Vec<Word>, tempfile::TempDir)| {
+        b.iter(|| {
+            for key in keys {
+                hint::black_box(smt.open(key));
+            }
+        })
+    },
+}
+
 benchmark_with_setup_data! {
     large_smt_compute_mutations,
     DEFAULT_MEASUREMENT_TIME,
@@ -241,6 +262,52 @@ benchmark_batch! {
     |size| Some(criterion::Throughput::Elements(size as u64))
 }
 
+benchmark_batch! {
+    large_smt_insert_batch_to_empty_tree,
+    &[100, 1_000, 10_000],
+    |b: &mut criterion::Bencher, insert_count: usize| {
+        b.iter_batched(
+            || {
+                let temp_dir = tempfile::TempDir::new().unwrap();
+                let storage = RocksDbStorage::open(RocksDbConfig::new(temp_dir.path())).unwrap();
+                let smt = LargeSmt::with_entries(storage, empty()).unwrap();
+                let batch = generate_smt_entries_sequential(insert_count);
+
+                (temp_dir, smt, batch)
+            },
+            |(_temp_dir, mut smt, batch)| {
+                smt.insert_batch(batch).unwrap();
+            },
+            BatchSize::LargeInput
+        )
+    },
+    |size| Some(criterion::Throughput::Elements(size as u64))
+}
+
+benchmark_batch! {
+    large_smt_insert_batch_to_populated_tree,
+    &[100, 1_000, 10_000],
+    |b: &mut criterion::Bencher, insert_count: usize| {
+        let initial_entries = generate_smt_entries_sequential(10_000);
+
+        b.iter_batched(
+            || {
+                let temp_dir = tempfile::TempDir::new().unwrap();
+                let storage = RocksDbStorage::open(RocksDbConfig::new(temp_dir.path())).unwrap();
+                let smt = LargeSmt::with_entries(storage, initial_entries.clone()).unwrap();
+                let batch = generate_smt_entries_sequential(insert_count);
+
+                (temp_dir, smt, batch)
+            },
+            |(_temp_dir, mut smt, batch)| {
+                smt.insert_batch(batch).unwrap();
+            },
+            BatchSize::LargeInput
+        )
+    },
+    |size| Some(criterion::Throughput::Elements(size as u64))
+}
+
 // MEMORY STORAGE BENCHMARKS
 // ================================================================================================
 
@@ -333,10 +400,13 @@ benchmark_batch! {
 criterion_group!(
     large_smt_benchmark_group,
     large_smt_open,
+    large_smt_open_in_large_tree,
     large_smt_compute_mutations,
     large_smt_apply_mutations,
     large_smt_apply_mutations_with_reversion,
     large_smt_insert_batch,
+    large_smt_insert_batch_to_empty_tree,
+    large_smt_insert_batch_to_populated_tree,
 );
 
 criterion_group!(

miden-crypto/benches/large_smt_forest.rs

@@ -0,0 +1,228 @@
+//! This module contains the benchmarks for the large SMT forest, focusing on the performance of key
+//! operations.
+
+mod common;
+
+use std::hint;
+
+use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
+use miden_crypto::{
+    merkle::smt::{
+        Backend, ForestPersistentBackend, LargeSmtForest, LineageId, PersistentBackendConfig,
+        SmtForestUpdateBatch, SmtUpdateBatch, TreeId,
+    },
+    rand::test_utils::rand_value,
+};
+use miden_field::Word;
+
+use crate::common::{
+    config::{DEFAULT_MEASUREMENT_TIME, DEFAULT_SAMPLE_SIZE},
+    data::{generate_smt_entries_sequential, generate_test_keys_sequential},
+};
+
+// CONSTANTS
+// ================================================================================================
+
+/// The number of entries to modify in an arbitrary batch of updates.
+const BATCH_SIZE: usize = 10_000;
+
+/// The number of trees we update in a single whole-forest batch.
+const TREES_PER_BATCH: usize = 50;
+
+// SETUP FUNCTIONALITY
+// ================================================================================================
+
+/// The setup for a benchmark over the smt forest.
+#[derive(Debug)]
+struct ForestSetup<B: Backend> {
+    pub forest: LargeSmtForest<B>,
+    _file: Option<tempfile::TempDir>,
+}
+impl ForestSetup<ForestPersistentBackend> {
+    /// Sets up a new persistent forest as a benchmark setup.
+    fn new_persistent() -> Self {
+        let dir = tempfile::tempdir().unwrap();
+        let backend =
+            ForestPersistentBackend::load(PersistentBackendConfig::new(dir.path()).unwrap())
+                .unwrap();
+        let forest = LargeSmtForest::new(backend).unwrap();
+        let _file = Some(dir);
+
+        Self { forest, _file }
+    }
+}
+
+/// Generates a tree update batch containing `count` entries which may be additions or removals.
+fn generate_tree_update_batch(count: usize) -> SmtUpdateBatch {
+    let entries = generate_smt_entries_sequential(count);
+    SmtUpdateBatch::from(entries.into_iter())
+}
+
+/// Generates a forest update batch containing `count` entries which may be additions or removals
+/// and which are allocated equally over the `lineages` in the forest.
+fn generate_forest_update_batch(lineages: &[LineageId], count: usize) -> SmtForestUpdateBatch {
+    let mut updates = SmtForestUpdateBatch::empty();
+    for lineage in lineages {
+        *updates.operations(*lineage) = generate_tree_update_batch(count / lineages.len());
+    }
+    updates
+}
+
+/// Generates `count` lineage identifiers.
+fn generate_lineages(count: usize) -> Vec<LineageId> {
+    let mut lineages = Vec::new();
+    for _ in 0..count {
+        lineages.push(LineageId::new(rand_value()));
+    }
+    lineages
+}
+
+// FOREST WITH PERSISTENT BACKEND
+// ================================================================================================
+
+// Roughly equivalent to large_smt::large_smt_open in functionality.
+benchmark_with_setup_data! {
+    large_smt_forest_persistent_open_full_tree,
+    DEFAULT_MEASUREMENT_TIME,
+    DEFAULT_SAMPLE_SIZE,
+    "large_smt_forest_persistent_open_full_tree",
+    || {
+        let mut setup = ForestSetup::new_persistent();
+        let batch = generate_tree_update_batch(BATCH_SIZE);
+        let lineage = LineageId::new([0x42; 32]);
+        let version = 0;
+        setup.forest.add_lineage(lineage, version, batch).unwrap();
+        let keys = generate_test_keys_sequential(10);
+        let tree = TreeId::new(lineage, version);
+        (setup, keys, tree)
+    },
+    |b: &mut criterion::Bencher, (setup, keys, tree): &(ForestSetup<_>, Vec<Word>, TreeId)| {
+        b.iter(|| {
+            for key in keys {
+                hint::black_box(setup.forest.open(*tree, *key).unwrap());
+            }
+        })
+    }
+}
+
+// Doesn't have a direct analogue in large SMT, but should be roughly equivalent in performance to
+// large_smt_forest_persistent_open_full_tree above, as the historical portion should not dominate.
+benchmark_with_setup_data! {
+    large_smt_forest_persistent_open_historical_tree,
+    DEFAULT_MEASUREMENT_TIME,
+    DEFAULT_SAMPLE_SIZE,
+    "large_smt_forest_persistent_open_historical_tree",
+    || {
+        let mut setup = ForestSetup::new_persistent();
+        let initial_batch = generate_tree_update_batch(BATCH_SIZE);
+        let lineage = LineageId::new([0x42; 32]);
+        let version = 0;
+        setup.forest.add_lineage(lineage, version, initial_batch).unwrap();
+        let update_batch = generate_tree_update_batch(BATCH_SIZE);
+        setup.forest.update_tree(lineage, 1, update_batch).unwrap();
+
+        let keys = generate_test_keys_sequential(10);
+        let tree = TreeId::new(lineage, version);
+        (setup, keys, tree)
+    },
+    |b: &mut criterion::Bencher, (setup, keys, tree): &(ForestSetup<_>, Vec<Word>, TreeId)| {
+        b.iter(|| {
+            for key in keys {
+                hint::black_box(setup.forest.open(*tree, *key).unwrap());
+            }
+        })
+    },
+}
+
+// Roughly equivalent to large_smt::large_smt_insert_batch_to_empty_tree in functionality.
+benchmark_batch! {
+    large_smt_forest_persistent_add_lineage,
+    &[100, 1_000, 10_000],
+    |b: &mut criterion::Bencher, entry_count: usize| {
+        let lineage = LineageId::new([0x42; 32]);
+        let version = 0;
+
+        b.iter_batched(
+            || {
+                let batch = generate_tree_update_batch(entry_count);
+                let setup = ForestSetup::new_persistent();
+                (setup, batch)
+            },
+            |(mut setup, batch)| {
+                setup.forest.add_lineage(lineage, version, batch).unwrap()
+            },
+            BatchSize::LargeInput
+        )
+    },
+    |size| Some(criterion::Throughput::Elements(size as u64))
+}
+
+// Roughly equivalent to large_smt::large_smt_insert_batch_to_populated_tree in functionality.
+benchmark_batch! {
+    large_smt_forest_persistent_update_tree,
+    &[100, 1_000, 10_000],
+    |b: &mut criterion::Bencher, entry_count: usize| {
+        let initial_batch = generate_tree_update_batch(BATCH_SIZE);
+        let lineage = LineageId::new([0x42; 32]);
+        let version = 0;
+
+        b.iter_batched(
+            || {
+                let mut setup = ForestSetup::new_persistent();
+                setup.forest.add_lineage(lineage, version, initial_batch.clone()).unwrap();
+                let batch = generate_tree_update_batch(entry_count);
+                (setup, batch)
+            },
+            |(mut setup, batch)| {
+                setup.forest.update_tree(lineage, version + 1, batch).unwrap();
+            },
+            BatchSize::LargeInput
+        )
+    },
+    |size| Some(criterion::Throughput::Elements(size as u64))
+}
+
+// Has no direct equivalent in the large smt, but should be broadly equivalent workwise to the
+// large_smt_forest_persistent_update_tree above in time as we try and do as much in parallel as
+// possible.
+benchmark_batch! {
+    large_smt_forest_persistent_update_forest,
+    &[100, 1_000, 10_000],
+    |b: &mut criterion::Bencher, entry_count: usize| {
+        let initial_batch = generate_tree_update_batch(100);
+        let lineages = generate_lineages(TREES_PER_BATCH);
+
+        b.iter_batched(
+            || {
+                let mut setup = ForestSetup::new_persistent();
+                let version = 0;
+                for lineage in &lineages {
+                    setup.forest.add_lineage(*lineage, version, initial_batch.clone()).unwrap();
+                }
+
+                let batch = generate_forest_update_batch(&lineages, entry_count);
+
+                (setup, batch)
+            },
+            |(mut setup, batch)| {
+                hint::black_box(setup.forest.update_forest(1, batch).unwrap())
+            },
+            BatchSize::LargeInput
+        )
+    },
+    |size| Some(criterion::Throughput::Elements(size as u64))
+}
+
+// BENCHMARK RUNS
+// ================================================================================================
+
+criterion_group!(
+    large_smt_forest_persistent_group,
+    large_smt_forest_persistent_open_full_tree,
+    large_smt_forest_persistent_open_historical_tree,
+    large_smt_forest_persistent_add_lineage,
+    large_smt_forest_persistent_update_tree,
+    large_smt_forest_persistent_update_forest,
+);
+
+criterion_main!(large_smt_forest_persistent_group);