diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..c9b65f3 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,2 @@ +[env] +RUST_MIN_STACK = "8388608" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d60d595..3147209 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -4,6 +4,7 @@ on: push: branches: [main] pull_request: + branches: [main] env: CARGO_TERM_COLOR: always diff --git a/.gitignore b/.gitignore index ad67955..cf2f4b0 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ # will have compiled files and executables debug target +Cargo.lock # These are backup files generated by rustfmt **/*.rs.bk @@ -13,6 +14,14 @@ target # Contains mutation testing data **/mutants.out*/ +# Benchmark artifacts +profile.json.gz +.db +*.csv +bench_results/ + +**/.DS_Store + # RustRover # JetBrains specific template is maintained in a separate JetBrains.gitignore that can # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore diff --git a/Cargo.toml b/Cargo.toml index a807033..ed62f5a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,3 +8,11 @@ edition = "2024" rust-version = "1.95.0" license = "MIT OR Apache-2.0" repository = "https://github.com/ethsystems/mono" + +[profile.release] +lto = "fat" +codegen-units = 1 + +[profile.profiling] +inherits = "release" +debug = true diff --git a/README.md b/README.md index 63f8d51..d55ec09 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,10 @@ Rust 1.95.0 (pinned via [`rust-toolchain.toml`](rust-toolchain.toml)). +## Crates + +- [`rotortree`](crates/rotortree): n-ary leanIMT implementation with persistence, for high-throughput append-only merkle trees. + ## Security See [SECURITY.md](SECURITY.md) for our vulnerability disclosure policy. diff --git a/crates/.gitkeep b/crates/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/crates/rotortree/Cargo.toml b/crates/rotortree/Cargo.toml new file mode 100644 index 0000000..d4a804f --- /dev/null +++ b/crates/rotortree/Cargo.toml @@ -0,0 +1,84 @@ +[package] +name = "rotortree" +version = "0.17.0" +edition.workspace = true +license.workspace = true +rust-version.workspace = true +repository.workspace = true +readme = "README.md" +description = "rotortree: a fast n-ary leanIMT implementation with persistence" + +[features] +default = ["blake3"] +blake3 = ["dep:blake3"] +wincode = ["dep:wincode"] +std = ["blake3?/std", "wincode?/std", "crc-fast?/std", "serde?/std"] +concurrent = ["std", "dep:parking_lot"] +parallel = ["std", "dep:rayon"] +storage = ["std", "dep:parking_lot", "dep:arc-swap", "wincode", "dep:crc-fast", "dep:fs4", "dep:memmap2"] +test-helpers = [] +docs = ["dep:include-utils"] +serde = ["dep:serde", "dep:serde_with"] + +[dependencies] +# Pinned exactly: the SIMD batch parent-hashing path (src/adapters/blake3.rs) +# drives the #[doc(hidden)] blake3::platform::Platform::hash_many API with +# hardcoded IV + flag constants, and relies on its single-chunk output being +# byte-identical to blake3::hash. That contract is not part of blake3's stable +# surface, so a minor bump could silently change it. The byte-identity unit +# test guards behavior; this pin guards the API. +blake3 = { version = "=1.8.5", default-features = false, optional = true } +wincode = { version = "0.4", default-features = false, features = ["alloc", "derive"], optional = true } +parking_lot = { version = "0.12", optional = true } +arc-swap = { version = "1", optional = true } +rayon = { version = "1", optional = true } +crc-fast = { version = "1.10", default-features = false, optional = true } +fs4 = { version = "0.13", optional = true, features = ["sync"] } +memmap2 = { version = "0.9", optional = true } +include-utils = { version = "0.2.4", optional = true } +serde = { version = "1", default-features = false, optional = true, features = ["derive"] } +serde_with = { version = "3.17", default-features = false, optional = true } + +[dev-dependencies] +proptest = "1" +criterion = { version = "0.8", features = ["html_reports"] } +crabtime = "1.1.4" +tempfile = "3" +tikv-jemallocator = "0.6" +rotortree = { path = ".", features = ["test-helpers"] } + +[[bench]] +name = "tree_bench" +harness = false + +[[bench]] +name = "tree_bench_concurrent" +harness = false +required-features = ["concurrent"] + +[[bench]] +name = "tree_bench_parallel" +harness = false +required-features = ["parallel"] + +[[bench]] +name = "tree_bench_all" +harness = false +required-features = ["concurrent", "parallel"] + +[[bench]] +name = "tree_bench_storage" +harness = false +required-features = ["storage"] + +[[example]] +name = "bulk_load" +required-features = ["storage", "parallel", "blake3"] + +[[example]] +name = "light_client" +required-features = ["storage", "blake3"] + +[package.metadata.docs.rs] +all-features = true +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/rotortree/README.md b/crates/rotortree/README.md new file mode 100644 index 0000000..56c29fb --- /dev/null +++ b/crates/rotortree/README.md @@ -0,0 +1,221 @@ +# rotortree + +A fast n-ary leanIMT implementation with persistence, for high-throughput append-only merkle trees. + + +Most privacy protocols reuse general-purpose databases such as rocksdb or postgres, which may not suit high-throughput use cases; this is also why layerzero built [qmdb](https://arxiv.org/pdf/2501.05262). rotortree's design differs in that it focuses exclusively on append-only merkle trees and does not support updating leaves in place. + + + +This approach makes significant tradeoffs and is not intended for production use. + +The tree design is inspired by [lean-imt](https://zkkit.org/leanimt-paper.pdf), based on work by cedoor and vivian at [PSE](https://pse.dev). This design allows for functional equivalents in zk DSLs and Solidity; the main deviation here is an n-ary leanIMT, which reduces tree depth while maintaining the same leaf count and allows on-disk storage blocks to group leaves together efficiently. + + + +## Design decisions + +- you should have a k-v database in tandem with this to ensure you don't insert the same nullifier twice. +- you should constrain node values to the finite field you're using before insertion +- generic hasher, blake3 default +- batteries included for playing with different branching factors and max depths +- wal for persistence and recovery, with checkpointing to prevent unbounded wal growth +- [wincode](https://github.com/anza-xyz/wincode) & regular serde for compatibility with ecosystem (e.g the jolt prover needs serde) +- no_std by default, persistence requires std +- benchmarks driven and configured by divan + crabtime +- by default your tree lives in memory, but with the `storage` feature you can tier cold levels to mmap'd data files via `TieringConfig::pin_above_level` + - with N=4, MAX_DEPTH=16, you can fit ~4.3B nullifiers in 41 GiB + - with N=8, MAX_DEPTH=10, you can fit ~1B nullifiers in 37 GiB + - which are quite feasible, but expensive. just use a new tree per generation and encode your nullifiers with the generation pls + - in most cases, you just need the tree in memory without crash persistence (as long as there is a bootstrapping sync mechanism), just use the single threaded variant, its _MUCH_ better if you have a low number of insertions + - writes go to the wal + memory, reads are always from memory or mmap. one cannot do an apples to apples comparison with other merkle tree dbs that read from disk on every query +- assumes that leaves are prehashed, partial support for rfc 6962 only for internal nodes +- few dependencies ~ 65 (active + transitive, excluding dev deps), zero deps for just the algorithm + + + + + +## Usage + +### In-memory (no persistence) `default-features = false` + +use `LeanIMT` + +```rust +use rotortree::{LeanIMT, Blake3Hasher}; + +// N=4 branching factor, MAX_DEPTH=20 +let mut tree = LeanIMT::::new(Blake3Hasher); + +// single insert +let leaf = [1u8; 32]; +let root = tree.insert(leaf).unwrap(); + +// batch insert +let leaves: Vec<[u8; 32]> = (0..1000u32) + .map(|i| { + let mut h = [0u8; 32]; + h[..4].copy_from_slice(&i.to_le_bytes()); + h + }) + .collect(); +let root = tree.insert_many(&leaves).unwrap(); + +// proof generation & verification +let snap = tree.snapshot(); +let proof = snap.generate_proof(0).unwrap(); +assert!(proof.verify(&Blake3Hasher).unwrap()); +``` + +optional feature flags for the in-memory mode: +- `concurrent`: switches to `&self` methods with internal `RwLock` (via `parking_lot`) +- `parallel`: enables rayon-parallelized `insert_many` for large batches (this works really well) +- `wincode`: adds wincode serde derives to proof types + +### With WAL persistence (`storage` feature) + +```rust +use rotortree::{ + Blake3Hasher, RotorTree, RotorTreeConfig, + FlushPolicy, CheckpointPolicy, TieringConfig, +}; +use std::path::PathBuf; + +let config = RotorTreeConfig { + path: PathBuf::from("/tmp/my-tree"), + flush_policy: FlushPolicy::default(), // fsync every 10ms + checkpoint_policy: CheckpointPolicy::default(), // manual + tiering: TieringConfig::default(), // all in memory + verify_checkpoint: true, // recompute root on recovery +}; + +// opens existing WAL or creates a new one +let tree = RotorTree::::open(Blake3Hasher, config).unwrap(); + +// insert: returns root + a durability token +let (root, token) = tree.insert([42u8; 32]).unwrap(); +// token.wait() blocks until the entry is fsynced + +// or insert + wait for fsync in one call +let root = tree.insert_durable([43u8; 32]).unwrap(); + +// batch insert +let leaves = vec![[1u8; 32]; 500]; +let (root, token) = tree.insert_many(&leaves).unwrap(); + +// lock-free snapshot for proof generation (same as in-memory) +let snap = tree.snapshot(); +let proof = snap.generate_proof(0).unwrap(); +assert!(proof.verify(&Blake3Hasher).unwrap()); + +// explicit flush & close +tree.flush().unwrap(); +tree.close().unwrap(); +// (also flushes + releases file lock on drop) +``` + +`FlushPolicy` options: +- `Interval(Duration)`: background thread fsyncs periodically +- `Manual`: caller controls flushing via `tree.flush()` (works well if you're following a blockchain as the canonical source of state transitions) + +`CheckpointPolicy` options (materializes wal state to data files, allowing wal truncation): +- `Manual`: caller calls `checkpoint()` explicitly +- `EveryNEntries(n)`: auto-checkpoint after every `n` wal entries +- `MemoryThreshold(bytes)`: auto-checkpoint when uncheckpointed memory exceeds threshold +- `OnClose`: checkpoint only on graceful close + +`TieringConfig` controls which levels stay in memory vs get mmap'd after checkpoint: +- `pin_above_level`: levels below this value have their committed chunks mmap'd from data files after checkpoint. set to `usize::MAX` to mmap everything (default), `0` to keep everything in memory + +### Provability + +#### Jolt + +```math +\text{execution\_cycles} \approx 2948.4 + 1832.9 \cdot N \cdot D + (899.5 + 2993.8N - 103.4N^2) \cdot d +``` + +```math +\text{where } N \in \{2, 4, 8, 16\}, \quad D = \text{max\_depth}, \quad d = \lceil \log_N(\text{num\_leaves}) \rceil +``` + +#### Noir (UltraHonk via bb) + +```math +\text{expression\_width} = 68 + (217N - 100) \cdot d +``` +```math +\text{circuit\_size} = 2962 + 30N + (551N^2 + 2629N - 2448) \cdot d +``` + +```math +\text{where } N \in \{2, 4, 8, 16\}, \quad d = \lceil \log_N(\text{num\_leaves}) \rceil +``` + +### Tuning + +of course, you'll have to benchmark your unique workload to see if this database suits your use case and requirements. here are some constants and env vars you can play around with to alter behaviour: + +compile-time constants (change in source and recompile): +- `CHUNK_SIZE` (default `128`): hashes per chunk for structural sharing. 128 × 32 bytes = 4 KB per chunk. affects snapshot copy cost and arc granularity +- `CHUNKS_PER_SEGMENT` (default `256`): chunks per immutable segment. controls how many chunks are frozen into a single arc slab +- `PAR_CHUNK_SIZE` (default `64`, `parallel` feature): parents per rayon work unit. smaller values = more parallelism but more scheduling overhead +- `MAX_FRAME_PAYLOAD` (default `128 MB`): maximum wal/checkpoint frame payload size. change this if you insert_many more than 128 MB worth of leaves + +runtime env vars: +- `ROTORTREE_PARALLEL_THRESHOLD` (default `1024`, `parallel` feature): minimum parent count before rayon parallelism kicks in + + + +## Development + +### Prerequisites + +- [cargo-hack](https://github.com/taiki-e/cargo-hack?tab=readme-ov-file#installation): to test all combinations of feature flags +- [cargo-nextest](https://nexte.st/): rust test runner + +### Check + +``` +cargo hack check --feature-powerset +``` + +### Clippy + +``` +cargo hack clippy --feature-powerset +``` + +### Format + +``` +cargo +nightly fmt +``` + +## Testing + +``` +cargo hack nextest run --feature-powerset +``` + +## Benchmarks + +``` +cargo bench -- --list +``` + +there are some feature flagged benchmarks, refer to the [Cargo.toml entry](Cargo.toml) for more details + + + +the pure in-memory benchmark (tree_bench_parallel) exhibits lesser variance in the benchmarks, and achieves peak throughput upto ~190M leaves/sec; single threaded by far has the best performance characteristic in terms of variance though, useful to keep in mind if that is a constraint; trading off performance for predictability under load. + + + +There are more realistic benchmarks that simulate performance under load, i.e concurrent reads / proof generation + insertions + +## Future work + +1. optimize `ceil_log_n` by precomputing the table +2. run benchmarks in an isolated environment for better estimations diff --git a/crates/rotortree/benches/common/mod.rs b/crates/rotortree/benches/common/mod.rs new file mode 100644 index 0000000..bcf1a0c --- /dev/null +++ b/crates/rotortree/benches/common/mod.rs @@ -0,0 +1,11 @@ +use rotortree::Hash; + +pub fn generate_leaves(count: usize) -> Vec { + (0..count) + .map(|i| { + let mut leaf = [0u8; 32]; + leaf[..8].copy_from_slice(&(i as u64).to_le_bytes()); + leaf + }) + .collect() +} diff --git a/crates/rotortree/benches/tree_bench.rs b/crates/rotortree/benches/tree_bench.rs new file mode 100644 index 0000000..3b85617 --- /dev/null +++ b/crates/rotortree/benches/tree_bench.rs @@ -0,0 +1,251 @@ +use criterion::{ + BatchSize, + BenchmarkId, + Throughput, + criterion_group, + criterion_main, +}; +use rotortree::{ + Blake3Hasher, + LeanIMT, +}; + +mod common; +use common::generate_leaves; + +#[crabtime::function] +fn bench_insert_single(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_single_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_single/n", {{n}})); + for count in [1_000usize, 10_000, 100_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || LeanIMT::::new(Blake3Hasher), + |tree| { + for &leaf in leaves { + std::hint::black_box(tree.insert(leaf).unwrap()); + } + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_insert_many(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_many_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_many/n", {{n}})); + for count in [1_000usize, 10_000, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || LeanIMT::::new(Blake3Hasher), + |tree| { + std::hint::black_box(tree.insert_many(leaves).unwrap()); + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_insert_many_chunked_100(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_many_chunked_100_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_many_chunked_100/n", {{n}})); + for count in [10_000usize, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || LeanIMT::::new(Blake3Hasher), + |tree| { + for chunk in leaves.chunks(100) { + std::hint::black_box(tree.insert_many(chunk).unwrap()); + } + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_insert_many_chunked_1000(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_many_chunked_1000_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_many_chunked_1000/n", {{n}})); + for count in [10_000usize, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || LeanIMT::::new(Blake3Hasher), + |tree| { + for chunk in leaves.chunks(1000) { + std::hint::black_box(tree.insert_many(chunk).unwrap()); + } + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_insert_incremental(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_incremental_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_incremental/n", {{n}})); + for count in [10_000usize, 100_000, 1_000_000] { + let all_leaves = generate_leaves(count); + let half = count / 2; + let (first_half, second_half) = all_leaves.split_at(half); + let first_half = first_half.to_vec(); + let second_half = second_half.to_vec(); + group.throughput(Throughput::Elements(second_half.len() as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &(first_half.clone(), second_half.clone()), + |b, (first_half, second_half)| { + b.iter_batched_ref( + || { + let mut tree = LeanIMT::::new(Blake3Hasher); + tree.insert_many(first_half).unwrap(); + tree + }, + |tree| { + std::hint::black_box(tree.insert_many(second_half).unwrap()); + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_generate_proof(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn generate_proof_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("generate_proof/n", {{n}})); + for count in [1_000usize, 10_000, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + let mut tree = LeanIMT::::new(Blake3Hasher); + tree.insert_many(&leaves).unwrap(); + let snap = tree.snapshot(); + let mid = (count / 2) as u64; + group.bench_function(BenchmarkId::from_parameter(count), |b| { + b.iter(|| { + std::hint::black_box(snap.generate_proof(mid).unwrap()); + }); + }); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_verify_proof(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn verify_proof_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("verify_proof/n", {{n}})); + for count in [1_000usize, 10_000, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + let mut tree = LeanIMT::::new(Blake3Hasher); + tree.insert_many(&leaves).unwrap(); + let snap = tree.snapshot(); + let proof = snap.generate_proof(0).unwrap(); + let th = Blake3Hasher; + group.bench_function(BenchmarkId::from_parameter(count), |b| { + b.iter(|| { + std::hint::black_box(proof.verify(&th).unwrap()); + }); + }); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn define_harness(n_values: Vec) { + for n in n_values { + crabtime::output! { + bench_insert_single!([{{n}}]); + bench_insert_many!([{{n}}]); + bench_insert_many_chunked_100!([{{n}}]); + bench_insert_many_chunked_1000!([{{n}}]); + bench_insert_incremental!([{{n}}]); + bench_generate_proof!([{{n}}]); + bench_verify_proof!([{{n}}]); + + criterion_group!( + benches_n{{n}}, + insert_single_n{{n}}, + insert_many_n{{n}}, + insert_many_chunked_100_n{{n}}, + insert_many_chunked_1000_n{{n}}, + insert_incremental_n{{n}}, + generate_proof_n{{n}}, + verify_proof_n{{n}} + ); + } + } +} + +define_harness!([2, 4, 8, 16]); +criterion_main!(benches_n2, benches_n4, benches_n8, benches_n16); diff --git a/crates/rotortree/benches/tree_bench_all.rs b/crates/rotortree/benches/tree_bench_all.rs new file mode 100644 index 0000000..51c54ed --- /dev/null +++ b/crates/rotortree/benches/tree_bench_all.rs @@ -0,0 +1,150 @@ +use criterion::{ + BatchSize, + BenchmarkId, + Throughput, + criterion_group, + criterion_main, +}; +use rotortree::{ + Blake3Hasher, + LeanIMT, +}; +use std::sync::Arc; + +mod common; +use common::generate_leaves; + +#[crabtime::function] +fn bench_insert_many(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_many_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_many/n", {{n}})); + for count in [1_000usize, 10_000, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || LeanIMT::::new(Blake3Hasher), + |tree| { + std::hint::black_box(tree.insert_many(leaves).unwrap()); + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_insert_incremental(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_incremental_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_incremental/n", {{n}})); + for count in [10_000usize, 100_000, 1_000_000] { + let all_leaves = generate_leaves(count); + let half = count / 2; + let (first_half, second_half) = all_leaves.split_at(half); + let first_half = first_half.to_vec(); + let second_half = second_half.to_vec(); + group.throughput(Throughput::Elements(second_half.len() as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &(first_half.clone(), second_half.clone()), + |b, (first_half, second_half)| { + b.iter_batched_ref( + || { + let tree = LeanIMT::::new(Blake3Hasher); + tree.insert_many(first_half).unwrap(); + tree + }, + |tree| { + std::hint::black_box(tree.insert_many(second_half).unwrap()); + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_concurrent_contention(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn concurrent_contention_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("concurrent_contention/n", {{n}})); + for count in [10_000usize, 100_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched( + || Arc::new(LeanIMT::::new(Blake3Hasher)), + |tree| { + std::thread::scope(|s| { + for _ in 0..4 { + let tree = Arc::clone(&tree); + s.spawn(move || { + loop { + let snap = tree.snapshot(); + let size = snap.size(); + if size > 0 { + let _ = std::hint::black_box( + snap.generate_proof(size / 2) + ); + } + if size >= count as u64 { + break; + } + } + }); + } + tree.insert_many(leaves).unwrap(); + }); + }, + BatchSize::PerIteration, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn define_harness(n_values: Vec) { + for n in n_values { + crabtime::output! { + bench_insert_many!([{{n}}]); + bench_insert_incremental!([{{n}}]); + bench_concurrent_contention!([{{n}}]); + + criterion_group!( + benches_n{{n}}, + insert_many_n{{n}}, + insert_incremental_n{{n}}, + concurrent_contention_n{{n}} + ); + } + } +} + +define_harness!([2, 4, 8, 16]); +criterion_main!(benches_n2, benches_n4, benches_n8, benches_n16); diff --git a/crates/rotortree/benches/tree_bench_concurrent.rs b/crates/rotortree/benches/tree_bench_concurrent.rs new file mode 100644 index 0000000..013504a --- /dev/null +++ b/crates/rotortree/benches/tree_bench_concurrent.rs @@ -0,0 +1,260 @@ +use criterion::{ + BatchSize, + BenchmarkId, + Throughput, + criterion_group, + criterion_main, +}; +use rotortree::{ + Blake3Hasher, + LeanIMT, +}; +use std::sync::Arc; + +mod common; +use common::generate_leaves; + +#[crabtime::function] +fn bench_insert_single(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_single_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_single/n", {{n}})); + for count in [1_000usize, 10_000, 100_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || LeanIMT::::new(Blake3Hasher), + |tree| { + for &leaf in leaves { + std::hint::black_box(tree.insert(leaf).unwrap()); + } + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_insert_many(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_many_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_many/n", {{n}})); + for count in [1_000usize, 10_000, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || LeanIMT::::new(Blake3Hasher), + |tree| { + std::hint::black_box(tree.insert_many(leaves).unwrap()); + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_insert_incremental(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_incremental_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_incremental/n", {{n}})); + for count in [10_000usize, 100_000, 1_000_000] { + let all_leaves = generate_leaves(count); + let half = count / 2; + let (first_half, second_half) = all_leaves.split_at(half); + let first_half = first_half.to_vec(); + let second_half = second_half.to_vec(); + group.throughput(Throughput::Elements(second_half.len() as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &(first_half.clone(), second_half.clone()), + |b, (first_half, second_half)| { + b.iter_batched_ref( + || { + let tree = LeanIMT::::new(Blake3Hasher); + tree.insert_many(first_half).unwrap(); + tree + }, + |tree| { + std::hint::black_box(tree.insert_many(second_half).unwrap()); + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_generate_proof(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn generate_proof_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("generate_proof/n", {{n}})); + for count in [1_000usize, 10_000, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + let tree = LeanIMT::::new(Blake3Hasher); + tree.insert_many(&leaves).unwrap(); + let snap = tree.snapshot(); + let mid = (count / 2) as u64; + group.bench_function(BenchmarkId::from_parameter(count), |b| { + b.iter(|| { + std::hint::black_box(snap.generate_proof(mid).unwrap()); + }); + }); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_verify_proof(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn verify_proof_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("verify_proof/n", {{n}})); + for count in [1_000usize, 10_000, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + let tree = LeanIMT::::new(Blake3Hasher); + tree.insert_many(&leaves).unwrap(); + let snap = tree.snapshot(); + let proof = snap.generate_proof(0).unwrap(); + let th = Blake3Hasher; + group.bench_function(BenchmarkId::from_parameter(count), |b| { + b.iter(|| { + std::hint::black_box(proof.verify(&th).unwrap()); + }); + }); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_snapshot(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn snapshot_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("snapshot/n", {{n}})); + for count in [1_000usize, 10_000, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + let tree = LeanIMT::::new(Blake3Hasher); + tree.insert_many(&leaves).unwrap(); + group.bench_function(BenchmarkId::from_parameter(count), |b| { + b.iter(|| { + std::hint::black_box(tree.snapshot()); + }); + }); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_concurrent_contention(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn concurrent_contention_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("concurrent_contention/n", {{n}})); + for count in [10_000usize, 100_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched( + || Arc::new(LeanIMT::::new(Blake3Hasher)), + |tree| { + std::thread::scope(|s| { + for _ in 0..4 { + let tree = Arc::clone(&tree); + s.spawn(move || { + loop { + let snap = tree.snapshot(); + let size = snap.size(); + if size > 0 { + let _ = std::hint::black_box( + snap.generate_proof(size / 2) + ); + } + if size >= count as u64 { + break; + } + } + }); + } + tree.insert_many(leaves).unwrap(); + }); + }, + BatchSize::PerIteration, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn define_harness(n_values: Vec) { + for n in n_values { + crabtime::output! { + bench_insert_single!([{{n}}]); + bench_insert_many!([{{n}}]); + bench_insert_incremental!([{{n}}]); + bench_generate_proof!([{{n}}]); + bench_verify_proof!([{{n}}]); + bench_snapshot!([{{n}}]); + bench_concurrent_contention!([{{n}}]); + + criterion_group!( + benches_n{{n}}, + insert_single_n{{n}}, + insert_many_n{{n}}, + insert_incremental_n{{n}}, + generate_proof_n{{n}}, + verify_proof_n{{n}}, + snapshot_n{{n}}, + concurrent_contention_n{{n}} + ); + } + } +} + +define_harness!([2, 4, 8, 16]); +criterion_main!(benches_n2, benches_n4, benches_n8, benches_n16); diff --git a/crates/rotortree/benches/tree_bench_parallel.rs b/crates/rotortree/benches/tree_bench_parallel.rs new file mode 100644 index 0000000..fbdf766 --- /dev/null +++ b/crates/rotortree/benches/tree_bench_parallel.rs @@ -0,0 +1,178 @@ +use criterion::{ + BatchSize, + BenchmarkId, + Throughput, + criterion_group, + criterion_main, +}; +use rotortree::{ + Blake3Hasher, + LeanIMT, +}; + +mod common; +use common::generate_leaves; + +#[crabtime::function] +fn bench_insert_many(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_many_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_many/n", {{n}})); + for count in [1_000usize, 10_000, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || { + let tree = LeanIMT::::new(Blake3Hasher); + rayon::broadcast(|_| {}); + tree + }, + |tree| { + std::hint::black_box(tree.insert_many(leaves).unwrap()); + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_insert_many_chunked_100(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_many_chunked_100_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_many_chunked_100/n", {{n}})); + for count in [10_000usize, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || { + let tree = LeanIMT::::new(Blake3Hasher); + rayon::broadcast(|_| {}); + tree + }, + |tree| { + for chunk in leaves.chunks(100) { + std::hint::black_box(tree.insert_many(chunk).unwrap()); + } + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_insert_many_chunked_1000(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_many_chunked_1000_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_many_chunked_1000/n", {{n}})); + for count in [10_000usize, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || { + let tree = LeanIMT::::new(Blake3Hasher); + rayon::broadcast(|_| {}); + tree + }, + |tree| { + for chunk in leaves.chunks(1000) { + std::hint::black_box(tree.insert_many(chunk).unwrap()); + } + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_insert_incremental(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_incremental_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_incremental/n", {{n}})); + for count in [10_000usize, 100_000, 1_000_000] { + let all_leaves = generate_leaves(count); + let half = count / 2; + let (first_half, second_half) = all_leaves.split_at(half); + let first_half = first_half.to_vec(); + let second_half = second_half.to_vec(); + group.throughput(Throughput::Elements(second_half.len() as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &(first_half.clone(), second_half.clone()), + |b, (first_half, second_half)| { + b.iter_batched_ref( + || { + let mut tree = LeanIMT::::new(Blake3Hasher); + tree.insert_many(first_half).unwrap(); + rayon::broadcast(|_| {}); + tree + }, + |tree| { + std::hint::black_box(tree.insert_many(second_half).unwrap()); + }, + BatchSize::LargeInput, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn define_harness(n_values: Vec) { + for n in n_values { + crabtime::output! { + bench_insert_many!([{{n}}]); + bench_insert_many_chunked_100!([{{n}}]); + bench_insert_many_chunked_1000!([{{n}}]); + bench_insert_incremental!([{{n}}]); + + criterion_group!( + benches_n{{n}}, + insert_many_n{{n}}, + insert_many_chunked_100_n{{n}}, + insert_many_chunked_1000_n{{n}}, + insert_incremental_n{{n}} + ); + } + } +} + +define_harness!([2, 4, 8, 16]); +criterion_main!(benches_n2, benches_n4, benches_n8, benches_n16); diff --git a/crates/rotortree/benches/tree_bench_storage.rs b/crates/rotortree/benches/tree_bench_storage.rs new file mode 100644 index 0000000..9264214 --- /dev/null +++ b/crates/rotortree/benches/tree_bench_storage.rs @@ -0,0 +1,328 @@ +use criterion::{ + BatchSize, + BenchmarkId, + Throughput, + criterion_group, + criterion_main, +}; +use rotortree::{ + Blake3Hasher, + CheckpointPolicy, + FlushPolicy, + RotorTree, + RotorTreeConfig, +}; + +mod common; +use common::generate_leaves; + +#[crabtime::function] +fn bench_insert_single(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_single_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_single/n", {{n}})); + for count in [1_000usize, 10_000, 100_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || { + let dir = tempfile::tempdir().unwrap(); + let config = RotorTreeConfig { + path: dir.path().to_path_buf(), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: Default::default(), + tiering: Default::default(), + verify_checkpoint: true, + }; + let tree = RotorTree::::open(Blake3Hasher, config).unwrap(); + (tree, dir) + }, + |(tree, _dir)| { + for &leaf in leaves { + std::hint::black_box(tree.insert(leaf).unwrap()); + } + }, + BatchSize::PerIteration, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_insert_many(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn insert_many_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("insert_many/n", {{n}})); + for count in [1_000usize, 10_000, 100_000, 1_000_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || { + let dir = tempfile::tempdir().unwrap(); + let config = RotorTreeConfig { + path: dir.path().to_path_buf(), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: Default::default(), + tiering: Default::default(), + verify_checkpoint: true, + }; + let tree = RotorTree::::open(Blake3Hasher, config).unwrap(); + #[cfg(feature = "parallel")] + rayon::broadcast(|_| {}); + (tree, dir) + }, + |(tree, _dir)| { + std::hint::black_box(tree.insert_many(leaves).unwrap()); + }, + BatchSize::PerIteration, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_flush(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn flush_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("flush/n", {{n}})); + for count in [1_000usize, 10_000, 100_000] { + let leaves = generate_leaves(count); + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(count), + &leaves, + |b, leaves| { + b.iter_batched( + || { + let dir = tempfile::tempdir().unwrap(); + let config = RotorTreeConfig { + path: dir.path().to_path_buf(), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: Default::default(), + tiering: Default::default(), + verify_checkpoint: true, + }; + let tree = RotorTree::::open(Blake3Hasher, config).unwrap(); + tree.insert_many(leaves).unwrap(); + (tree, dir) + }, + |(tree, _dir)| { + std::hint::black_box(tree.flush().unwrap()); + tree.close().unwrap(); + }, + BatchSize::PerIteration, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_open_recover(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn open_recover_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("open_recover/n", {{n}})); + for count in [1_000usize, 10_000, 100_000] { + let leaves = generate_leaves(count); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().to_path_buf(); + { + let config = RotorTreeConfig { + path: path.clone(), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: Default::default(), + tiering: Default::default(), + verify_checkpoint: true, + }; + let tree = RotorTree::::open(Blake3Hasher, config).unwrap(); + tree.insert_many(&leaves).unwrap(); + tree.flush().unwrap(); + tree.close().unwrap(); + } + group.throughput(Throughput::Elements(count as u64)); + group.bench_function(BenchmarkId::from_parameter(count), |b| { + b.iter_batched( + || path.clone(), + |path| { + let config = RotorTreeConfig { + path, + flush_policy: FlushPolicy::Manual, + checkpoint_policy: Default::default(), + tiering: Default::default(), + verify_checkpoint: true, + }; + let tree = RotorTree::::open(Blake3Hasher, config).unwrap(); + std::hint::black_box(tree.root()); + tree.close().unwrap(); + }, + BatchSize::PerIteration, + ); + }); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_mixed_workload(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn mixed_workload_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("mixed_workload/n", {{n}})); + for tick in [100usize, 1_000, 10_000, 100_000] { + let prepop_leaves = generate_leaves(10_000); + let tick_leaves = generate_leaves(tick); + let th = Blake3Hasher; + group.throughput(Throughput::Elements(tick as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(tick), + &(prepop_leaves.clone(), tick_leaves.clone()), + |b, (prepop_leaves, tick_leaves)| { + b.iter_batched( + || { + let dir = tempfile::tempdir().unwrap(); + let config = RotorTreeConfig { + path: dir.path().to_path_buf(), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: Default::default(), + tiering: Default::default(), + verify_checkpoint: true, + }; + let tree = RotorTree::::open(Blake3Hasher, config).unwrap(); + tree.insert_many(prepop_leaves).unwrap(); + tree.flush().unwrap(); + (tree, dir) + }, + |(tree, _dir)| { + let (root, _token) = tree.insert_many(tick_leaves).unwrap(); + std::hint::black_box(root); + + std::hint::black_box(tree.root()); + + let snap = tree.snapshot(); + let proof_index = snap.size() / 2; + let proof = snap.generate_proof(proof_index).unwrap(); + std::hint::black_box(&proof); + + std::hint::black_box(proof.verify(&th).unwrap()); + tree.close().unwrap(); + }, + BatchSize::PerIteration, + ); + }, + ); + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn bench_sustained_checkpoint(n_values: Vec) { + for n in n_values { + crabtime::output! { + fn sustained_checkpoint_n{{n}}(c: &mut criterion::Criterion) { + let mut group = c.benchmark_group(concat!("sustained_checkpoint/n", {{n}})); + for count in [100_000usize, 1_000_000] { + for freq in [5usize, 25, 100, 500] { + let leaves = generate_leaves(count); + let th = Blake3Hasher; + group.throughput(Throughput::Elements(count as u64)); + group.bench_with_input( + BenchmarkId::new(format!("every{freq}"), count), + &leaves, + |b, leaves| { + b.iter_batched_ref( + || { + let dir = tempfile::tempdir().unwrap(); + let config = RotorTreeConfig { + path: dir.path().to_path_buf(), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: CheckpointPolicy::EveryNEntries(freq as u64), + tiering: Default::default(), + verify_checkpoint: true, + }; + let tree = RotorTree::::open(Blake3Hasher, config).unwrap(); + (tree, dir) + }, + |(tree, _dir)| { + for chunk in leaves.chunks(10_000) { + std::hint::black_box(tree.insert_many(chunk).unwrap()); + + let snap = tree.snapshot(); + let size = snap.size(); + let proof = snap.generate_proof(0).unwrap(); + std::hint::black_box(proof.verify(&th).unwrap()); + let proof = snap.generate_proof(size / 2).unwrap(); + std::hint::black_box(proof.verify(&th).unwrap()); + let proof = snap.generate_proof(size - 1).unwrap(); + std::hint::black_box(proof.verify(&th).unwrap()); + } + }, + BatchSize::PerIteration, + ); + }, + ); + } + } + group.finish(); + } + } + } +} + +#[crabtime::function] +fn define_harness(n_values: Vec) { + for n in n_values { + crabtime::output! { + bench_insert_single!([{{n}}]); + bench_insert_many!([{{n}}]); + bench_flush!([{{n}}]); + bench_open_recover!([{{n}}]); + bench_mixed_workload!([{{n}}]); + bench_sustained_checkpoint!([{{n}}]); + + criterion_group!( + benches_n{{n}}, + insert_single_n{{n}}, + insert_many_n{{n}}, + flush_n{{n}}, + open_recover_n{{n}}, + mixed_workload_n{{n}}, + sustained_checkpoint_n{{n}} + ); + } + } +} + +define_harness!([2, 4, 8, 16]); +criterion_main!(benches_n2, benches_n4, benches_n8, benches_n16); diff --git a/crates/rotortree/examples/bulk_load.rs b/crates/rotortree/examples/bulk_load.rs new file mode 100644 index 0000000..06d66fc --- /dev/null +++ b/crates/rotortree/examples/bulk_load.rs @@ -0,0 +1,113 @@ +use rotortree::{ + Blake3Hasher, + CheckpointPolicy, + FlushPolicy, + RotorTree, + RotorTreeConfig, + TieringConfig, +}; +use std::{ + env, + fs::File, + io::Write, + path::PathBuf, + time::Instant, +}; + +const N: usize = 4; +const MAX_DEPTH: usize = 14; +const DB_PATH: &str = ".db"; + +use tikv_jemallocator::Jemalloc; + +#[global_allocator] +static GLOBAL: Jemalloc = Jemalloc; + +fn env_or(key: &str, default: T) -> T { + env::var(key) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) +} + +fn generate_leaves(start: u64, count: u64) -> Vec<[u8; 32]> { + (start..start + count) + .map(|i| { + let mut h = [0u8; 32]; + h[..8].copy_from_slice(&i.to_le_bytes()); + h + }) + .collect() +} + +fn main() { + let total_leaves: u64 = env_or("TOTAL_LEAVES", 100_000_000); + let block_size: u64 = env_or("BLOCK_SIZE", 1_000_000); + let proof_iters: usize = env_or("PROOF_ITERS", 100); + let csv_path: String = env_or("CSV_PATH", "bulk_load.csv".to_string()); + + let _ = std::fs::remove_dir_all(DB_PATH); + + let config = RotorTreeConfig { + path: PathBuf::from(DB_PATH), + flush_policy: FlushPolicy::Interval(std::time::Duration::from_millis(10)), + checkpoint_policy: CheckpointPolicy::MemoryThreshold(256 * 1024 * 1024), + tiering: TieringConfig::default(), + verify_checkpoint: false, + }; + let tree = + RotorTree::::open(Blake3Hasher, config).unwrap(); + + let mut csv = File::create(&csv_path).unwrap(); + writeln!( + csv, + "leaves,block_ins_per_sec,durable_ins_per_sec,proof_gen_ns,proof_verify_ns,depth" + ) + .unwrap(); + + let mut inserted: u64 = 0; + + while inserted < total_leaves { + let this_block = block_size.min(total_leaves - inserted); + let leaves = generate_leaves(inserted, this_block); + + let block_start = Instant::now(); + let (_root, token) = tree.insert_many(&leaves).unwrap(); + let insert_elapsed = block_start.elapsed(); + token.wait(); + let durable_elapsed = block_start.elapsed(); + + inserted += this_block; + + let probe_index = inserted - this_block / 2; + let snap = tree.snapshot(); + + let mut gen_times = Vec::with_capacity(proof_iters); + let mut verify_times = Vec::with_capacity(proof_iters); + for _ in 0..proof_iters { + let t = Instant::now(); + let proof = snap.generate_proof(probe_index).unwrap(); + gen_times.push(t.elapsed()); + let t = Instant::now(); + assert!(proof.verify(&Blake3Hasher).unwrap()); + verify_times.push(t.elapsed()); + } + gen_times.sort(); + verify_times.sort(); + let proof_gen_ns = gen_times[gen_times.len() / 2].as_nanos(); + let proof_verify_ns = verify_times[verify_times.len() / 2].as_nanos(); + + let block_ins = (this_block as f64 / insert_elapsed.as_secs_f64()) as u64; + let durable_ins = (this_block as f64 / durable_elapsed.as_secs_f64()) as u64; + + writeln!( + csv, + "{inserted},{block_ins},{durable_ins},{proof_gen_ns},{proof_verify_ns},{}", + tree.depth() + ) + .unwrap(); + } + + tree.close().unwrap(); + std::fs::remove_dir_all(DB_PATH).unwrap(); +} diff --git a/crates/rotortree/examples/light_client.rs b/crates/rotortree/examples/light_client.rs new file mode 100644 index 0000000..f322ffc --- /dev/null +++ b/crates/rotortree/examples/light_client.rs @@ -0,0 +1,212 @@ +#![cfg_attr(feature = "concurrent", allow(unused_mut))] + +//! Full-node / light-client consistency proof demo. +//! +//! Two threads communicate over `mpsc` channels: +//! +//! - **Full node**: persists batches to a `RotorTree` (WAL-backed) and sends +//! consistency proofs after each batch. +//! - **Light client**: bootstraps once with the initial leaves, then drops the +//! tree. Subsequent batches are verified using *only* the consistency proof: +//! no leaves are transferred, and the tracked inclusion proof is updated +//! purely from consistency data. +//! +//! ```sh +//! cargo run --example light_client --features storage,blake3 --release +//! ``` + +use rotortree::{ + Blake3Hasher, + CheckpointPolicy, + ConsistencyProof, + FlushPolicy, + Hash, + LeanIMT, + NaryProof, + RotorTree, + RotorTreeConfig, + TieringConfig, +}; +use std::{ + mem, + path::PathBuf, + sync::mpsc, + thread, + time::Instant, +}; + +const N: usize = 4; +const MAX_DEPTH: usize = 14; +const BATCH_SIZE: u64 = 10_000; +const NUM_BATCHES: usize = 5; +const TRACKED_LEAF: u64 = 812; +const DB_PATH: &str = ".db"; + +enum NodeMessage { + Bootstrap { + leaves: Vec, + }, + Update { + consistency_proof: ConsistencyProof, + }, + Shutdown, +} + +enum ClientMessage { + InclusionProof(NaryProof), +} + +fn generate_leaves(start: u64, count: u64) -> Vec { + (start..start + count) + .map(|i| *blake3::hash(&i.to_le_bytes()).as_bytes()) + .collect() +} + +fn full_node(tx: mpsc::Sender, rx: mpsc::Receiver) { + let hasher = Blake3Hasher; + let th = Blake3Hasher; + let _ = std::fs::remove_dir_all(DB_PATH); + let config = RotorTreeConfig { + path: PathBuf::from(DB_PATH), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: CheckpointPolicy::OnClose, + tiering: TieringConfig::default(), + verify_checkpoint: false, + }; + let tree = RotorTree::::open(hasher, config) + .expect("failed to open tree"); + + let mut inserted: u64 = 0; + + for batch in 0..NUM_BATCHES { + let leaves = generate_leaves(inserted, BATCH_SIZE); + + if batch == 0 { + let (_root, _token) = tree.insert_many(&leaves).expect("insert_many failed"); + inserted += BATCH_SIZE; + + let leaf_bytes = leaves.len() * mem::size_of::(); + println!(" [full node] batch 0: sending {leaf_bytes} bytes (raw leaves)"); + tx.send(NodeMessage::Bootstrap { leaves }).unwrap(); + + let ClientMessage::InclusionProof(proof) = rx.recv().unwrap(); + let snap = tree.snapshot(); + assert!(proof.verify(&th).unwrap()); + assert_eq!(proof.root, snap.root().unwrap()); + println!(" [full node] batch 0: client proof verified"); + } else { + let old_snap = tree.snapshot(); + let old_size = old_snap.size(); + let old_root = old_snap.root().unwrap(); + + let (_root, _token) = tree.insert_many(&leaves).expect("insert_many failed"); + inserted += BATCH_SIZE; + + let new_snap = tree.snapshot(); + let consistency_proof = new_snap + .generate_consistency_proof(old_size, old_root) + .expect("generate_consistency_proof failed"); + + let proof_bytes = mem::size_of_val(&consistency_proof); + let leaf_bytes = BATCH_SIZE as usize * mem::size_of::(); + println!( + " [full node] batch {batch}: sending {proof_bytes} bytes \ + (consistency proof, vs {leaf_bytes} bytes for raw leaves)" + ); + tx.send(NodeMessage::Update { consistency_proof }).unwrap(); + + let ClientMessage::InclusionProof(proof) = rx.recv().unwrap(); + assert!(proof.verify(&th).unwrap()); + assert_eq!(proof.root, new_snap.root().unwrap()); + println!(" [full node] batch {batch}: client proof verified"); + } + } + + tx.send(NodeMessage::Shutdown).unwrap(); + tree.close().expect("close failed"); + let _ = std::fs::remove_dir_all(DB_PATH); +} + +fn light_client(tx: mpsc::Sender, rx: mpsc::Receiver) { + let hasher = Blake3Hasher; + let th = Blake3Hasher; + let mut tracked_proof: Option> = None; + let mut current_root: Option = None; + + loop { + match rx.recv().unwrap() { + NodeMessage::Bootstrap { leaves } => { + let mut tree = LeanIMT::::new(hasher); + tree.insert_many(&leaves).expect("insert_many failed"); + let snap = tree.snapshot(); + let proof = snap + .generate_proof(TRACKED_LEAF) + .expect("generate_proof failed"); + assert!(proof.verify(&th).unwrap()); + + current_root = snap.root(); + println!( + " [light client] bootstrap: {} leaves, proof OK", + leaves.len() + ); + let send_proof = snap + .generate_proof(TRACKED_LEAF) + .expect("generate_proof failed"); + tracked_proof = Some(proof); + tx.send(ClientMessage::InclusionProof(send_proof)).unwrap(); + } + + NodeMessage::Update { consistency_proof } => { + assert!( + consistency_proof + .verify_transition(&th, current_root.unwrap()) + .unwrap() + ); + + let old_proof = tracked_proof.as_ref().expect("tracked proof must exist"); + let updated_proof = consistency_proof + .update_inclusion_proof(old_proof, &th) + .expect("update_inclusion_proof failed"); + + assert!(updated_proof.verify(&th).unwrap()); + assert_eq!(updated_proof.root, consistency_proof.new_root); + + current_root = Some(consistency_proof.new_root); + println!( + " [light client] update: consistency OK, proof updated \ + (old_size={}, new_size={})", + consistency_proof.old_size, consistency_proof.new_size, + ); + let send_proof = consistency_proof + .update_inclusion_proof(old_proof, &th) + .expect("update_inclusion_proof failed"); + tracked_proof = Some(updated_proof); + tx.send(ClientMessage::InclusionProof(send_proof)).unwrap(); + } + + NodeMessage::Shutdown => { + println!(" [light client] shutdown"); + break; + } + } + } +} + +fn main() { + println!( + "light_client: N={N}, MAX_DEPTH={MAX_DEPTH}, \ + {NUM_BATCHES} batches x {BATCH_SIZE} leaves, tracked leaf #{TRACKED_LEAF}" + ); + let start = Instant::now(); + + let (node_tx, client_rx) = mpsc::channel(); + let (client_tx, node_rx) = mpsc::channel(); + + let node_handle = thread::spawn(move || full_node(node_tx, node_rx)); + let client_handle = thread::spawn(move || light_client(client_tx, client_rx)); + + node_handle.join().expect("full node panicked"); + client_handle.join().expect("light client panicked"); + + println!("done in {:.2?}", start.elapsed()); +} diff --git a/crates/rotortree/proptest-regressions/simd_batch.txt b/crates/rotortree/proptest-regressions/simd_batch.txt new file mode 100644 index 0000000..60d4972 --- /dev/null +++ b/crates/rotortree/proptest-regressions/simd_batch.txt @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 70c2f5ce8ef8c20d76748ebdc0ffc9e37fcfda611c54bf93020740e89041444d # shrinks to n = 127, seed = 0 diff --git a/crates/rotortree/src/adapters/blake3.rs b/crates/rotortree/src/adapters/blake3.rs new file mode 100644 index 0000000..6aa1cb4 --- /dev/null +++ b/crates/rotortree/src/adapters/blake3.rs @@ -0,0 +1,243 @@ +use blake3::{ + IncrementCounter, + platform::{ + MAX_SIMD_DEGREE, + Platform, + }, +}; + +use crate::{ + Hash, + HashState, + Hasher, +}; + +// BLAKE3 spec constants, hardcoded because blake3's internal `IV` and flag +// consts are private. These pin `hash_many` to the exact configuration that +// makes a single-chunk multi-input hash byte-identical to `blake3::hash`: +// key = IV, counter = 0, increment_counter = No, flags = 0, +// flags_start = CHUNK_START, flags_end = CHUNK_END | ROOT. +// Identity holds ONLY when each input is a whole number of 64-byte blocks and +// is <= 1024 bytes (a single chunk). It diverges past 1024 bytes (multi-chunk), +// so the batch path below is restricted to arities whose byte length satisfies +// `len % 64 == 0 && len <= 1024`. See `Cargo.toml` (blake3 = "=1.8.5"). +const B3_IV: &[u32; 8] = &[ + 0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A, 0x510E527F, 0x9B05688C, 0x1F83D9AB, + 0x5BE0CD19, +]; +const CHUNK_START: u8 = 1; // 1 << 0 +const CHUNK_END: u8 = 2; // 1 << 1 +const ROOT: u8 = 8; // 1 << 3 +const FLAGS_END: u8 = CHUNK_END | ROOT; // 10 + +#[derive(Debug, Clone, Copy, Default)] +pub struct Blake3Hasher; + +impl Blake3Hasher { + pub fn new() -> Self { + Self + } +} + +/// The detected BLAKE3 platform (SIMD level). `detect()` is cheap (a couple of +/// cfg/cpuid branches) but std lets us cache it for free. +#[cfg(feature = "std")] +#[inline] +fn platform() -> Platform { + static PLATFORM: std::sync::OnceLock = std::sync::OnceLock::new(); + *PLATFORM.get_or_init(Platform::detect) +} + +#[cfg(not(feature = "std"))] +#[inline] +fn platform() -> Platform { + Platform::detect() +} + +/// Hash up to `MAX_SIMD_DEGREE` independent single-chunk inputs of exactly +/// `LEN` bytes each through the multi-input SIMD kernel. +/// +/// `LEN` MUST satisfy `LEN % 64 == 0 && LEN <= 1024`; otherwise the result is +/// NOT byte-identical to `blake3::hash` (silent corruption on the NEON path). +/// `inputs.len() <= MAX_SIMD_DEGREE` and `out.len() >= inputs.len()`. +#[inline] +fn batch_root( + plat: &Platform, + inputs: &[&[u8; LEN]], + out: &mut [Hash], +) { + debug_assert!( + LEN.is_multiple_of(64) && LEN <= 1024, + "ineligible input length" + ); + debug_assert!(inputs.len() <= MAX_SIMD_DEGREE); + debug_assert!(out.len() >= inputs.len()); + let flat: &mut [u8] = out[..inputs.len()].as_flattened_mut(); + plat.hash_many::( + inputs, + B3_IV, + 0, + IncrementCounter::No, + 0, + CHUNK_START, + FLAGS_END, + flat, + ); +} + +/// Reslice an eligible same-arity batch and dispatch it through `batch_root`. +/// Each group's bytes are contiguous (full groups at these arities are +/// chunk-aligned since `128 % N == 0`), so the checked `try_from` always +/// succeeds. Returns `false` (untouched `out`) for any ineligible arity so the +/// caller can take the scalar path. +fn batch_eligible( + plat: &Platform, + arity: usize, + groups: &[&[Hash]], + out: &mut [Hash], +) -> bool { + macro_rules! dispatch { + ($len:literal) => {{ + // Stack-collect &[u8; LEN] refs (at most MAX_SIMD_DEGREE). + let mut refs: [&[u8; $len]; MAX_SIMD_DEGREE] = + [&[0u8; $len]; MAX_SIMD_DEGREE]; + for (slot, g) in refs.iter_mut().zip(groups.iter()) { + match <&[u8; $len]>::try_from(g.as_flattened()) { + Ok(r) => *slot = r, + Err(_) => return false, // non-full / discontiguous group + } + } + batch_root::<$len>(plat, &refs[..groups.len()], out); + true + }}; + } + match arity { + 2 => dispatch!(64), + 4 => dispatch!(128), + 8 => dispatch!(256), + 16 => dispatch!(512), + _ => false, + } +} + +impl HashState for blake3::Hasher { + #[inline] + fn update(&mut self, data: &[u8]) { + blake3::Hasher::update(self, data); + } + + #[inline] + fn finalize(self) -> Hash { + *blake3::Hasher::finalize(&self).as_bytes() + } +} + +impl Hasher for Blake3Hasher { + type State = blake3::Hasher; + + #[inline] + fn new_state(&self) -> Self::State { + blake3::Hasher::new() + } + + /// One-shot scalar parent hash: `blake3::hash(child_0 || child_1 || ...)`. + /// + /// Byte-identical to the streaming default but skips the incremental + /// `Hasher` state. This is the fallback for every group the batched + /// `hash_many_into` path declines. + #[inline] + fn hash_children(&self, children: &[Hash]) -> Hash { + *blake3::hash(children.as_flattened()).as_bytes() + } + + /// Batched parent hashing through BLAKE3's multi-input SIMD kernel. + /// + /// Eligible only when every group has the same arity N in {2,4,8,16} + /// (byte length in {64,128,256,512}, all `% 64 == 0` and `<= 1024`), where + /// the multi-input hash is byte-identical to `hash_children`. Such groups + /// are processed `MAX_SIMD_DEGREE` at a time. Any other shape (mixed + /// arity, odd/lift/partial group, arity > 16, or a discontiguous slice) + /// falls back to the scalar `hash_children` per group. + fn hash_many_into(&self, groups: &[&[Hash]], out: &mut [Hash]) { + debug_assert!(out.len() >= groups.len()); + + // Eligible iff all groups share one arity in {2,4,8,16}. + let arity = groups.first().map_or(0, |g| g.len()); + let uniform_eligible = + matches!(arity, 2 | 4 | 8 | 16) && groups.iter().all(|g| g.len() == arity); + + if !uniform_eligible { + for (g, o) in groups.iter().zip(out.iter_mut()) { + *o = self.hash_children(g); + } + return; + } + + let plat = platform(); + let degree = plat.simd_degree().clamp(1, MAX_SIMD_DEGREE); + let mut g = groups; + let mut o = &mut out[..groups.len()]; + while !g.is_empty() { + let take = g.len().min(degree); + let (g_now, g_rest) = g.split_at(take); + let (o_now, o_rest) = o.split_at_mut(take); + // try_from inside batch_eligible re-checks contiguity; a false + // return means a group was not a full contiguous slice, so hash it + // scalar. Full groups at these arities are always contiguous. + if !batch_eligible(&plat, arity, g_now, o_now) { + for (gg, oo) in g_now.iter().zip(o_now.iter_mut()) { + *oo = self.hash_children(gg); + } + } + g = g_rest; + o = o_rest; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Each batched group hash must be byte-identical to the scalar + /// `hash_children` for the eligible arities {2,4,8,16}. + #[test] + fn batch_root_byte_identical_to_hash_children() { + let h = Blake3Hasher; + + fn make_groups(arity: usize) -> [[Hash; 16]; 4] { + let mut groups = [[[0u8; 32]; 16]; 4]; + for (k, group) in groups.iter_mut().enumerate() { + for (j, child) in group.iter_mut().enumerate().take(arity) { + for (b, byte) in child.iter_mut().enumerate() { + *byte = (k as u8) + .wrapping_mul(31) + .wrapping_add(j as u8) + .wrapping_add(b as u8) + ^ 0xA5; + } + } + } + groups + } + + for arity in [2usize, 4, 8, 16] { + let groups = make_groups(arity); + let refs: [&[Hash]; 4] = [ + &groups[0][..arity], + &groups[1][..arity], + &groups[2][..arity], + &groups[3][..arity], + ]; + let mut out = [[0u8; 32]; 4]; + h.hash_many_into(&refs, &mut out); + for (i, g) in refs.iter().enumerate() { + assert_eq!( + out[i], + h.hash_children(g), + "arity {arity} input {i} batched != scalar", + ); + } + } + } +} diff --git a/crates/rotortree/src/adapters/mod.rs b/crates/rotortree/src/adapters/mod.rs new file mode 100644 index 0000000..39f26c6 --- /dev/null +++ b/crates/rotortree/src/adapters/mod.rs @@ -0,0 +1,3 @@ +#[cfg(feature = "blake3")] +#[cfg_attr(docsrs, doc(cfg(feature = "blake3")))] +pub mod blake3; diff --git a/crates/rotortree/src/chunked_level.rs b/crates/rotortree/src/chunked_level.rs new file mode 100644 index 0000000..b7e131e --- /dev/null +++ b/crates/rotortree/src/chunked_level.rs @@ -0,0 +1,514 @@ +#[cfg(not(feature = "std"))] +use alloc::{ + boxed::Box, + sync::Arc, + vec::Vec, +}; +#[cfg(feature = "std")] +use std::{ + sync::Arc, + vec::Vec, +}; + +use crate::{ + Hash, + TreeError, +}; + +/// Number of hashes per chunk for structural sharing. +pub(crate) const CHUNK_SIZE: usize = 128; + +/// Number of chunks per immutable segment +pub(crate) const CHUNKS_PER_SEGMENT: usize = 256; + +#[cfg(not(feature = "storage"))] +#[derive(Clone)] +pub(crate) struct Chunk(Arc<[Hash; CHUNK_SIZE]>); + +#[cfg(not(feature = "storage"))] +impl Chunk { + #[inline(always)] + pub(crate) fn as_slice(&self) -> &[Hash; CHUNK_SIZE] { + &self.0 + } + + #[inline(always)] + pub(crate) fn make_mut(&mut self) -> &mut [Hash; CHUNK_SIZE] { + Arc::make_mut(&mut self.0) + } + + #[inline] + pub(crate) fn new_memory(data: [Hash; CHUNK_SIZE]) -> Self { + Self(Arc::new(data)) + } + + #[cfg(test)] + pub(crate) fn ptr_eq(a: &Self, b: &Self) -> bool { + Arc::ptr_eq(&a.0, &b.0) + } +} + +#[cfg(feature = "storage")] +#[derive(Clone)] +pub(crate) struct Chunk(ChunkInner); + +#[cfg(feature = "storage")] +#[derive(Clone)] +enum ChunkInner { + Memory(Arc<[Hash; CHUNK_SIZE]>), + Mapped { + region: Arc, + offset: usize, + }, +} + +#[cfg(feature = "storage")] +impl Chunk { + #[inline(always)] + pub(crate) fn as_slice(&self) -> &[Hash; CHUNK_SIZE] { + match &self.0 { + ChunkInner::Memory(arc) => arc, + ChunkInner::Mapped { region, offset } => { + // SAFETY: offset validated at construction + unsafe { &*(region.as_ptr().add(*offset).cast::<[Hash; CHUNK_SIZE]>()) } + } + } + } + + #[inline(always)] + pub(crate) fn make_mut(&mut self) -> &mut [Hash; CHUNK_SIZE] { + if matches!(&self.0, ChunkInner::Mapped { .. }) { + let data = *self.as_slice(); + self.0 = ChunkInner::Memory(Arc::new(data)); + } + match &mut self.0 { + ChunkInner::Memory(arc) => Arc::make_mut(arc), + ChunkInner::Mapped { .. } => unreachable!(), + } + } + + #[inline] + pub(crate) fn new_memory(data: [Hash; CHUNK_SIZE]) -> Self { + Self(ChunkInner::Memory(Arc::new(data))) + } + + pub(crate) fn new_mapped( + region: Arc, + offset: usize, + ) -> Self { + const CHUNK_BYTE_SIZE: usize = CHUNK_SIZE * 32; + assert!( + offset + CHUNK_BYTE_SIZE <= region.valid_len(), + "Chunk::new_mapped: offset {offset} + {CHUNK_BYTE_SIZE} exceeds valid_len {}", + region.valid_len() + ); + Self(ChunkInner::Mapped { region, offset }) + } + + #[cfg(test)] + pub(crate) fn ptr_eq(a: &Self, b: &Self) -> bool { + match (&a.0, &b.0) { + (ChunkInner::Memory(a), ChunkInner::Memory(b)) => Arc::ptr_eq(a, b), + _ => false, + } + } +} + +/// A single level of the tree stored as segmented chunks plus a +/// fixed-size tail buffer. +#[derive(Clone)] +pub(crate) struct ChunkedLevel { + /// Immutable segments of committed chunks, shared with snapshots. + segments: Vec>, + /// Mutable buffer of committed chunks not yet frozen into a segment. + /// At most `CHUNKS_PER_SEGMENT - 1` items. + pending: Vec, + /// Fixed-size tail buffer (partially filled). + tail: [Hash; CHUNK_SIZE], + /// Number of valid entries in `tail`. + tail_len: usize, + /// Total number of hashes in this level. + len: usize, +} + +impl ChunkedLevel { + pub(crate) fn new() -> Self { + Self { + segments: Vec::new(), + pending: Vec::new(), + tail: [[0u8; 32]; CHUNK_SIZE], + tail_len: 0, + len: 0, + } + } + + /// Construct a level from checkpoint data, partitioning chunks into + /// segments and pending. + #[cfg(feature = "storage")] + pub(crate) fn from_parts( + chunks: Vec, + tail: [Hash; CHUNK_SIZE], + tail_len: usize, + len: usize, + ) -> Self { + let full_segments = chunks.len() / CHUNKS_PER_SEGMENT; + let mut segments = Vec::with_capacity(full_segments); + let mut drain = chunks.into_iter(); + for _ in 0..full_segments { + let seg: Vec = drain.by_ref().take(CHUNKS_PER_SEGMENT).collect(); + let boxed: Box<[Chunk; CHUNKS_PER_SEGMENT]> = seg + .into_boxed_slice() + .try_into() + .unwrap_or_else(|_| unreachable!()); + segments.push(Arc::from(boxed)); + } + let pending: Vec = drain.collect(); + + Self { + segments, + pending, + tail, + tail_len, + len, + } + } + + /// Total number of hashes in this level. + #[inline] + pub(crate) fn len(&self) -> usize { + self.len + } + + /// Total number of committed chunks (segments + pending). + #[inline] + pub(crate) fn chunk_count(&self) -> usize { + self.segments.len() * CHUNKS_PER_SEGMENT + self.pending.len() + } + + /// Resolve a chunk index to a slice reference. + #[inline(always)] + fn chunk_slice(&self, chunk_idx: usize) -> &[Hash; CHUNK_SIZE] { + let committed = self.segments.len() * CHUNKS_PER_SEGMENT; + if chunk_idx < committed { + let seg_idx = chunk_idx / CHUNKS_PER_SEGMENT; + let seg_off = chunk_idx % CHUNKS_PER_SEGMENT; + self.segments[seg_idx][seg_off].as_slice() + } else { + self.pending[chunk_idx - committed].as_slice() + } + } + + /// Read a hash at the given index. + #[inline] + pub(crate) fn get(&self, index: usize) -> Result { + if index >= self.len { + return Err(TreeError::IndexOutOfRange { + index: index as u64, + size: self.len as u64, + }); + } + let chunk_idx = index / CHUNK_SIZE; + let offset = index % CHUNK_SIZE; + if chunk_idx < self.chunk_count() { + Ok(self.chunk_slice(chunk_idx)[offset]) + } else { + Ok(self.tail[offset]) + } + } + + /// Copy a contiguous group of hashes into `out`. + /// Fast path when the group falls within a single chunk or tail. + #[inline(always)] + pub(crate) fn get_group(&self, start: usize, count: usize, out: &mut [Hash]) { + let chunk_idx = start / CHUNK_SIZE; + let offset = start % CHUNK_SIZE; + if offset + count <= CHUNK_SIZE { + let src = if chunk_idx < self.chunk_count() { + &self.chunk_slice(chunk_idx)[offset..offset + count] + } else { + &self.tail[offset..offset + count] + }; + out[..count].copy_from_slice(src); + } else { + for (i, item) in out.iter_mut().enumerate().take(count) { + *item = self.get(start + i).expect("checked prev; qed"); + } + } + } + + /// Borrow a contiguous group of `count` hashes starting at `start`, + /// when it lies wholly within a single chunk or the tail. + /// + /// Returns `None` if the group straddles a chunk/tail boundary, in which + /// case the caller must fall back to a copying read. For the batched + /// parent path, full groups of arity N in {2,4,8,16} are chunk-aligned + /// (`CHUNK_SIZE % N == 0`), so this always returns `Some`. + #[inline] + pub(crate) fn group_slice(&self, start: usize, count: usize) -> Option<&[Hash]> { + if start + count > self.len { + return None; + } + let chunk_idx = start / CHUNK_SIZE; + let offset = start % CHUNK_SIZE; + if offset + count > CHUNK_SIZE { + return None; + } + if chunk_idx < self.chunk_count() { + Some(&self.chunk_slice(chunk_idx)[offset..offset + count]) + } else { + Some(&self.tail[offset..offset + count]) + } + } + + /// Write a hash at the given index + #[inline] + pub(crate) fn set(&mut self, index: usize, value: Hash) -> Result<(), TreeError> { + if self.len <= index { + self.ensure_len(index + 1)?; + } + self.set_preallocated(index, value); + Ok(()) + } + + /// Caller must ensure `index < self.len` + #[inline(always)] + pub(crate) fn set_preallocated(&mut self, index: usize, value: Hash) { + debug_assert!( + index < self.len, + "set_preallocated: index {index} >= len {}", + self.len + ); + let chunk_idx = index / CHUNK_SIZE; + let offset = index % CHUNK_SIZE; + let committed = self.segments.len() * CHUNKS_PER_SEGMENT; + if chunk_idx < committed { + let seg_idx = chunk_idx / CHUNKS_PER_SEGMENT; + let seg_off = chunk_idx % CHUNKS_PER_SEGMENT; + Arc::make_mut(&mut self.segments[seg_idx])[seg_off].make_mut()[offset] = + value; + } else if chunk_idx - committed < self.pending.len() { + self.pending[chunk_idx - committed].make_mut()[offset] = value; + } else { + self.tail[offset] = value; + } + } + + /// Append a hash. Promotes the tail when it reaches + /// `CHUNK_SIZE`. + #[cfg(test)] + #[inline] + pub(crate) fn push(&mut self, value: Hash) -> Result<(), TreeError> { + self.tail[self.tail_len] = value; + self.tail_len = self.tail_len.checked_add(1).ok_or(TreeError::MathError)?; + self.len = self.len.checked_add(1).ok_or(TreeError::MathError)?; + if self.tail_len == CHUNK_SIZE { + self.promote_tail(); + } + Ok(()) + } + + pub(crate) fn extend(&mut self, values: &[Hash]) -> Result<(), TreeError> { + if values.is_empty() { + return Ok(()); + } + let new_len = self + .len + .checked_add(values.len()) + .ok_or(TreeError::MathError)?; + + let mut remaining = values; + + // fill current tail + if self.tail_len > 0 { + let space = CHUNK_SIZE - self.tail_len; + let to_copy = space.min(remaining.len()); + self.tail[self.tail_len..self.tail_len + to_copy] + .copy_from_slice(&remaining[..to_copy]); + self.tail_len += to_copy; + remaining = &remaining[to_copy..]; + if self.tail_len == CHUNK_SIZE { + self.promote_tail(); + } + } + + // full chunks — bypass tail + let full_chunks = remaining.len() / CHUNK_SIZE; + if full_chunks > 0 { + self.pending.reserve(full_chunks.min(CHUNKS_PER_SEGMENT)); + for i in 0..full_chunks { + let start = i * CHUNK_SIZE; + let chunk: [Hash; CHUNK_SIZE] = remaining[start..start + CHUNK_SIZE] + .try_into() + .expect("slice len == CHUNK_SIZE; qed"); + self.push_chunk(Chunk::new_memory(chunk)); + } + remaining = &remaining[full_chunks * CHUNK_SIZE..]; + } + + // tail remainder + if !remaining.is_empty() { + self.tail[..remaining.len()].copy_from_slice(remaining); + self.tail_len = remaining.len(); + } + + self.len = new_len; + Ok(()) + } + + pub(crate) fn ensure_len(&mut self, target: usize) -> Result<(), TreeError> { + if self.len >= target { + return Ok(()); + } + let needed = target - self.len; + + let tail_space = CHUNK_SIZE - self.tail_len; + let fill_tail = tail_space.min(needed); + debug_assert!( + self.tail[self.tail_len..self.tail_len + fill_tail] + .iter() + .all(|h| *h == [0u8; 32]), + "ensure_len: tail slots must be zeroed" + ); + self.tail_len += fill_tail; + let mut filled = fill_tail; + if self.tail_len == CHUNK_SIZE { + self.promote_tail(); + } + + let remaining = needed - filled; + let full_chunks = remaining / CHUNK_SIZE; + if full_chunks > 0 { + for _ in 0..full_chunks { + self.push_chunk(Chunk::new_memory([[0u8; 32]; CHUNK_SIZE])); + } + filled += full_chunks * CHUNK_SIZE; + } + + let leftover = needed - filled; + self.tail_len += leftover; + + self.len = target; + Ok(()) + } + + /// Promote the full tail into a chunk, freezing pending if full + fn promote_tail(&mut self) { + debug_assert_eq!(self.tail_len, CHUNK_SIZE); + self.push_chunk(Chunk::new_memory(self.tail)); + self.tail = [[0u8; 32]; CHUNK_SIZE]; + self.tail_len = 0; + } + + /// Push a chunk to pending, freezing into a segment when full + fn push_chunk(&mut self, chunk: Chunk) { + self.pending.push(chunk); + if self.pending.len() == CHUNKS_PER_SEGMENT { + self.freeze_pending(); + } + } + + /// Freeze the full pending buffer into an immutable segment + fn freeze_pending(&mut self) { + debug_assert_eq!(self.pending.len(), CHUNKS_PER_SEGMENT); + let pending = core::mem::take(&mut self.pending); + let boxed_arr: Box<[Chunk; CHUNKS_PER_SEGMENT]> = pending + .into_boxed_slice() + .try_into() + .unwrap_or_else(|_| unreachable!()); // qed + self.segments.push(Arc::from(boxed_arr)); + } + + /// Collect chunks from index `already` onward + #[cfg(feature = "storage")] + pub(crate) fn chunks_since(&self, already: usize) -> Vec { + let total = self.chunk_count(); + if already >= total { + return Vec::new(); + } + let committed = self.segments.len() * CHUNKS_PER_SEGMENT; + let mut result = Vec::with_capacity(total - already); + + // Collect from segments + if already < committed { + let start_seg = already / CHUNKS_PER_SEGMENT; + let start_off = already % CHUNKS_PER_SEGMENT; + for (seg_i, segment) in self.segments.iter().enumerate().skip(start_seg) { + let from = if seg_i == start_seg { start_off } else { 0 }; + for chunk in &segment[from..] { + result.push(chunk.clone()); + } + } + } + + // Collect from pending + let pending_start = already.saturating_sub(committed); + if pending_start < self.pending.len() { + for chunk in &self.pending[pending_start..] { + result.push(chunk.clone()); + } + } + + result + } + + /// Remap the first `count` chunks to mmap-backed chunks (one region per shard) + #[cfg(feature = "storage")] + pub(crate) fn remap_chunks( + &mut self, + count: usize, + regions: &[Arc], + ) { + use crate::storage::checkpoint::shard_address; + + let total = self.chunk_count(); + let remap_count = count.min(total); + if remap_count == 0 || regions.is_empty() { + return; + } + + let committed = self.segments.len() * CHUNKS_PER_SEGMENT; + let mut unmapped: Vec = + Vec::with_capacity(total.saturating_sub(remap_count)); + for chunk_idx in remap_count..total { + if chunk_idx < committed { + let seg_idx = chunk_idx / CHUNKS_PER_SEGMENT; + let seg_off = chunk_idx % CHUNKS_PER_SEGMENT; + unmapped.push(self.segments[seg_idx][seg_off].clone()); + } else { + unmapped.push(self.pending[chunk_idx - committed].clone()); + } + } + + self.segments.clear(); + self.pending.clear(); + + (0..remap_count) + .map(|chunk_idx| { + let (shard_idx, offset_in_shard) = shard_address(chunk_idx); + Chunk::new_mapped(Arc::clone(®ions[shard_idx]), offset_in_shard) + }) + .chain(unmapped) + .for_each(|chunk| self.push_chunk(chunk)); + } + + /// Access the tail buffer + #[cfg(feature = "storage")] + pub(crate) fn tail_data(&self) -> &[Hash; CHUNK_SIZE] { + &self.tail + } + + #[cfg(test)] + pub(crate) fn tail_len(&self) -> usize { + self.tail_len + } + + #[cfg(test)] + pub(crate) fn get_chunk(&self, idx: usize) -> &Chunk { + let committed = self.segments.len() * CHUNKS_PER_SEGMENT; + if idx < committed { + &self.segments[idx / CHUNKS_PER_SEGMENT][idx % CHUNKS_PER_SEGMENT] + } else { + &self.pending[idx - committed] + } + } +} diff --git a/crates/rotortree/src/error.rs b/crates/rotortree/src/error.rs new file mode 100644 index 0000000..34bca47 --- /dev/null +++ b/crates/rotortree/src/error.rs @@ -0,0 +1,66 @@ +use core::fmt; + +/// Errors produced by tree operations. +#[derive(Debug, PartialEq, Eq)] +pub enum TreeError { + /// The tree would exceed its compile-time maximum depth. + MaxDepthExceeded { max_depth: usize }, + /// The requested leaf index is outside `[0, size)`. + IndexOutOfRange { index: u64, size: u64 }, + /// `insert_many` was called with an empty slice. + EmptyBatch, + /// The tree's `u64` size cannot be represented as `usize` + CapacityExceeded, + /// Math error + MathError, + /// Consistency proof depth mismatch + InvalidProofDepth { expected: usize, actual: usize }, + /// The proof update is a no-op (old_size == new_size) + NoUpdateNeeded, +} + +impl fmt::Display for TreeError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::MaxDepthExceeded { max_depth } => { + write!( + f, + "tree depth would exceed maximum of \ + {max_depth}" + ) + } + Self::IndexOutOfRange { index, size } => { + write!( + f, + "leaf index {index} out of range for tree \ + of size {size}" + ) + } + Self::EmptyBatch => { + write!(f, "insert_many called with empty batch") + } + Self::CapacityExceeded => { + write!( + f, + "tree capacity exceeded (u64 to usize \ + overflow)" + ) + } + Self::MathError => { + write!(f, "math error") + } + Self::InvalidProofDepth { expected, actual } => { + write!( + f, + "consistency proof depth {actual} does not match \ + expected {expected}" + ) + } + Self::NoUpdateNeeded => { + write!(f, "proof update is a no-op (same size)") + } + } + } +} + +impl core::error::Error for TreeError {} diff --git a/crates/rotortree/src/hash.rs b/crates/rotortree/src/hash.rs new file mode 100644 index 0000000..6ee2197 --- /dev/null +++ b/crates/rotortree/src/hash.rs @@ -0,0 +1,46 @@ +pub type Hash = [u8; 32]; + +/// Streaming hash state +pub trait HashState { + /// Feed bytes into the state + fn update(&mut self, data: &[u8]); + + /// Finalize and return the digest + fn finalize(self) -> Hash; +} + +/// Raw hash primitive supporting streaming +pub trait Hasher: Clone + Send + Sync + 'static { + type State: HashState; + + /// Create a fresh hashing state + fn new_state(&self) -> Self::State; + + /// Hash the concatenation of child node hashes into their parent. + /// + /// Children are hashed verbatim, with no domain tag and no length + /// prefix, so an internal node is `hash(child_0 || child_1 || ...)`. + /// This matches the canonical Lean IMT node hash and keeps the + /// digest compatible with other implementations. + #[inline] + fn hash_children(&self, children: &[Hash]) -> Hash { + let mut state = self.new_state(); + state.update(children.as_flattened()); + state.finalize() + } + + /// Hash a batch of independent parent groups into `out`. + /// + /// Element `i` of `out` receives `hash_children(groups[i])`. The + /// default implementation is a scalar loop; hashers backed by a + /// multi-input SIMD kernel (e.g. BLAKE3) override this to hash several + /// groups at once while staying byte-identical to `hash_children`. + /// + /// `out.len()` must be at least `groups.len()`. + #[inline] + fn hash_many_into(&self, groups: &[&[Hash]], out: &mut [Hash]) { + for (g, o) in groups.iter().zip(out.iter_mut()) { + *o = self.hash_children(g); + } + } +} diff --git a/crates/rotortree/src/lib.rs b/crates/rotortree/src/lib.rs new file mode 100644 index 0000000..0ccb4a6 --- /dev/null +++ b/crates/rotortree/src/lib.rs @@ -0,0 +1,79 @@ +#![cfg_attr(feature = "docs", doc = include_utils::include_md!("README.md:intro"))] +#![cfg_attr(feature = "docs", doc = include_utils::include_md!("README.md:design"))] +#![cfg_attr(feature = "docs", doc = include_utils::include_md!("README.md:usage"))] +#![cfg_attr(feature = "docs", doc = include_utils::include_md!("README.md:devnote"))] +#![cfg_attr(not(test), deny(clippy::cast_possible_truncation))] +#![cfg_attr(docsrs, feature(doc_cfg))] +#![deny(unused_crate_dependencies)] +#![deny(warnings)] +#![cfg_attr(not(feature = "std"), no_std)] + +#[cfg(not(feature = "std"))] +#[cfg_attr(docsrs, doc(cfg(not(feature = "std"))))] +extern crate alloc; + +// only used in benches/examples +#[cfg(test)] +use { + crabtime as _, + criterion as _, + proptest as _, + rotortree as _, + tempfile as _, + tikv_jemallocator as _, +}; + +pub(crate) mod chunked_level; +mod error; +mod hash; +mod proof; +#[cfg(any(test, feature = "test-helpers"))] +pub mod test_util; +mod tree; + +pub mod adapters; + +#[cfg(feature = "storage")] +#[cfg_attr(docsrs, doc(cfg(feature = "storage")))] +pub mod storage; + +pub use error::TreeError; +pub use hash::{ + Hash, + HashState, + Hasher, +}; +pub use proof::{ + ConsistencyLevel, + ConsistencyProof, + NaryProof, + ProofLevel, +}; +pub use tree::{ + LeanIMT, + TreeSnapshot, +}; + +#[cfg(feature = "blake3")] +#[cfg_attr(docsrs, doc(cfg(feature = "blake3")))] +pub use adapters::blake3::Blake3Hasher; + +#[cfg(feature = "storage")] +#[cfg_attr(docsrs, doc(cfg(feature = "storage")))] +pub use storage::{ + CheckpointPolicy, + DurabilityToken, + FlushPolicy, + RotorTree, + RotorTreeConfig, + RotorTreeError, + StorageError, + TieringConfig, +}; + +#[cfg(feature = "storage")] +#[doc(hidden)] +pub use storage::checkpoint::write_test_meta; + +#[cfg(feature = "storage")] +pub use storage::checkpoint::CheckpointMeta; diff --git a/crates/rotortree/src/proof.rs b/crates/rotortree/src/proof.rs new file mode 100644 index 0000000..ba425b2 --- /dev/null +++ b/crates/rotortree/src/proof.rs @@ -0,0 +1,1178 @@ +use crate::{ + Hash, + Hasher, + TreeError, + tree::{ + TreeSnapshot, + ceil_log_n, + u64_to_usize, + }, +}; + +fn to_u8(v: usize) -> Result { + u8::try_from(v).map_err(|_| TreeError::MathError) +} + +/// One level of an N-ary Merkle proof. +/// +/// Uses a fixed-size `[Hash; N]` array (at most `N-1` siblings +/// are valid, indicated by `sibling_count`). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[cfg_attr(feature = "wincode", derive(wincode::SchemaWrite, wincode::SchemaRead))] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +pub struct ProofLevel { + /// Child position within the group (`0..N-1`) + pub position: u8, + /// Number of valid siblings in `siblings` (0 for lifted) + pub sibling_count: u8, + /// Sibling hashes + #[cfg_attr( + feature = "serde", + serde(with = "serde_with::As::<[serde_with::Same; N]>") + )] + pub siblings: [Hash; N], +} + +impl ProofLevel { + const EMPTY: Self = Self { + position: 0, + sibling_count: 0, + siblings: [[0u8; 32]; N], + }; +} + +/// An N-ary Merkle inclusion proof +#[derive(Debug, Clone, PartialEq, Eq)] +#[cfg_attr(feature = "wincode", derive(wincode::SchemaWrite, wincode::SchemaRead))] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +pub struct NaryProof { + /// The Merkle root this proof verifies against + pub root: Hash, + /// The leaf hash being proved + pub leaf: Hash, + /// The 0-based index of the leaf + pub leaf_index: u64, + /// Number of leaves in the tree this proof was generated from + pub tree_size: u64, + /// Number of valid levels in `levels` + pub level_count: usize, + /// Proof levels from leaf to root + #[cfg_attr( + feature = "serde", + serde(with = "serde_with::As::<[serde_with::Same; MAX_DEPTH]>") + )] + pub levels: [ProofLevel; MAX_DEPTH], +} + +impl NaryProof { + /// Verify this proof against the given hasher. + pub fn verify(&self, hasher: &H) -> Result { + if self.tree_size == 0 || self.leaf_index >= self.tree_size { + return Err(TreeError::MathError); + } + + let expected_depth = ceil_log_n(self.tree_size, N); + if self.level_count != expected_depth || expected_depth > MAX_DEPTH { + return Err(TreeError::InvalidProofDepth { + expected: expected_depth, + actual: self.level_count, + }); + } + + let mut expected_index = u64_to_usize(self.leaf_index)?; + let mut current = self.leaf; + let mut children = [[0u8; 32]; N]; + + for level in &self.levels[..self.level_count] { + let expected_pos = expected_index % N; + if level.position as usize != expected_pos { + return Err(TreeError::MathError); + } + + if level.sibling_count > 0 { + let total = (level.sibling_count as usize) + 1; + let pos = level.position as usize; + if total > N || pos >= total { + return Err(TreeError::MathError); + } + children[..pos].copy_from_slice(&level.siblings[..pos]); + children[pos] = current; + let rest = total - pos - 1; + children[pos + 1..total] + .copy_from_slice(&level.siblings[pos..pos + rest]); + current = hasher.hash_children(&children[..total]); + } + + expected_index /= N; + } + + Ok(current == self.root) + } + + /// Verify this inclusion proof against an externally-trusted root. + pub fn verify_against( + &self, + hasher: &H, + trusted_root: Hash, + ) -> Result { + if self.root != trusted_root { + return Ok(false); + } + self.verify(hasher) + } +} + +/// Per-level data in a consistency proof +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[cfg_attr(feature = "wincode", derive(wincode::SchemaWrite, wincode::SchemaRead))] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +pub struct ConsistencyLevel { + /// Number of shared (left) siblings in `hashes` + pub shared_count: u8, + /// Number of new-only (right) siblings in `hashes`. + pub new_count: u8, + /// `[shared..., new...]` hashes + #[cfg_attr( + feature = "serde", + serde(with = "serde_with::As::<[serde_with::Same; N]>") + )] + pub hashes: [Hash; N], +} + +impl ConsistencyLevel { + const EMPTY: Self = Self { + shared_count: 0, + new_count: 0, + hashes: [[0u8; 32]; N], + }; + + fn new_hashes(&self) -> &[Hash] { + let sc = self.shared_count as usize; + let nc = self.new_count as usize; + &self.hashes[sc..sc + nc] + } +} + +/// Consistency proof +/// +/// Proves that a tree of `old_size` leaves is a prefix of a tree of +/// `new_size` leaves +#[derive(Debug, Clone, PartialEq, Eq)] +#[cfg_attr(feature = "wincode", derive(wincode::SchemaWrite, wincode::SchemaRead))] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +pub struct ConsistencyProof { + /// Root of the old (smaller) tree + pub old_root: Hash, + /// Root of the new (larger) tree + pub new_root: Hash, + /// Number of leaves in the old tree + pub old_size: u64, + /// Number of leaves in the new tree + pub new_size: u64, + /// Number of valid levels in `levels` + pub level_count: usize, + /// Proof levels from leaves toward root + #[cfg_attr( + feature = "serde", + serde(with = "serde_with::As::<[serde_with::Same; MAX_DEPTH]>") + )] + pub levels: [ConsistencyLevel; MAX_DEPTH], +} + +impl ConsistencyProof { + /// Verify this consistency proof against the given hasher. + /// + /// Returns `Ok(true)` if both reconstructed roots match. + pub fn verify(&self, hasher: &H) -> Result { + if self.old_size == 0 + || self.new_size == 0 + || self.old_size > self.new_size + || self.level_count > MAX_DEPTH + { + return Err(TreeError::MathError); + } + + if self.old_size == self.new_size { + return Ok(self.old_root == self.new_root); + } + if self.level_count == 0 { + return Ok(false); + } + + let mut old_level_size = u64_to_usize(self.old_size)?; + let mut new_level_size = u64_to_usize(self.new_size)?; + + let expected_depth = { + let (mut os, mut ns) = (old_level_size, new_level_size); + let mut d = 0usize; + while os > 1 || ns > 1 { + d += 1; + os = os.div_ceil(N); + ns = ns.div_ceil(N); + } + d + }; + if self.level_count != expected_depth { + return Err(TreeError::InvalidProofDepth { + expected: expected_depth, + actual: self.level_count, + }); + } + + let mut old_hash = [0u8; 32]; + let mut new_hash = [0u8; 32]; + let mut children = [[0u8; 32]; N]; + + for (i, level) in self.levels[..self.level_count].iter().enumerate() { + let shared = level.shared_count as usize; + let new_only = level.new_count as usize; + let boundary = old_level_size - 1; + let boundary_pos = boundary % N; + let group_start = boundary - boundary_pos; + + let expected_shared = if i == 0 { + boundary_pos + 1 + } else { + boundary_pos + }; + let expected_new_count = new_level_size.min(group_start + N) - boundary - 1; + if shared != expected_shared || new_only != expected_new_count { + return Err(TreeError::MathError); + } + + let (old_mid, new_mid) = if i == 0 { + (level.hashes[boundary_pos], level.hashes[boundary_pos]) + } else { + (old_hash, new_hash) + }; + + let new_child_count = boundary_pos + 1 + new_only; + children[..boundary_pos].copy_from_slice(&level.hashes[..boundary_pos]); + children[boundary_pos] = old_mid; + old_hash = if boundary_pos > 0 { + hasher.hash_children(&children[..boundary_pos + 1]) + } else { + old_mid + }; + children[boundary_pos] = new_mid; + children[boundary_pos + 1..new_child_count] + .copy_from_slice(level.new_hashes()); + new_hash = if new_child_count > 1 { + hasher.hash_children(&children[..new_child_count]) + } else { + new_mid + }; + + old_level_size = old_level_size.div_ceil(N); + new_level_size = new_level_size.div_ceil(N); + } + + Ok(old_hash == self.old_root && new_hash == self.new_root) + } + + /// Verify this consistency proof against externally-trusted roots. + pub fn verify_against( + &self, + hasher: &H, + trusted_old_root: Hash, + trusted_new_root: Hash, + ) -> Result { + if self.old_root != trusted_old_root || self.new_root != trusted_new_root { + return Ok(false); + } + self.verify(hasher) + } + + /// Verify this consistency proof against a trusted old root. + /// + /// If verification succeeds, the caller can trust `self.new_root`. + pub fn verify_transition( + &self, + hasher: &H, + trusted_old_root: Hash, + ) -> Result { + self.verify_against(hasher, trusted_old_root, self.new_root) + } + + /// Update an existing inclusion proof from `old_root` to `new_root`. + /// + /// Given a valid `NaryProof` for some leaf against `old_root`, produces + /// a valid `NaryProof` for the same leaf against `new_root` using only + /// the data in this consistency proof + pub fn update_inclusion_proof( + &self, + old_proof: &NaryProof, + hasher: &H, + ) -> Result, TreeError> { + if old_proof.root != self.old_root { + return Err(TreeError::MathError); + } + if !self.verify(hasher)? { + return Err(TreeError::MathError); + } + if old_proof.leaf_index >= self.old_size { + return Err(TreeError::IndexOutOfRange { + index: old_proof.leaf_index, + size: self.old_size, + }); + } + + if self.old_size == self.new_size { + return Err(TreeError::NoUpdateNeeded); + } + + let mut new_levels = [ProofLevel::::EMPTY; MAX_DEPTH]; + let mut old_level_size = u64_to_usize(self.old_size)?; + let mut member_idx = u64_to_usize(old_proof.leaf_index)?; + let mut new_boundary_hash = [0u8; 32]; + let mut children = [[0u8; 32]; N]; + + for (tree_level, level) in self.levels[..self.level_count].iter().enumerate() { + let boundary = old_level_size - 1; + let boundary_pos = boundary % N; + let member_pos = member_idx % N; + + let new_only = level.new_count as usize; + + if member_idx / N != boundary / N { + if tree_level < old_proof.level_count { + new_levels[tree_level] = old_proof.levels[tree_level]; + } + } else { + let mut group = [[0u8; 32]; N]; + let old_group_size = if tree_level < old_proof.level_count { + let old_level = &old_proof.levels[tree_level]; + let group_size = (old_level.sibling_count as usize) + 1; + group[..member_pos] + .copy_from_slice(&old_level.siblings[..member_pos]); + if member_pos + 1 < group_size { + group[member_pos + 1..group_size].copy_from_slice( + &old_level.siblings[member_pos..group_size - 1], + ); + } + group_size + } else { + 1 + }; + + if tree_level > 0 && member_pos != boundary_pos { + group[boundary_pos] = new_boundary_hash; + } + + let new_group_size = old_group_size + new_only; + if new_only > 0 { + group[old_group_size..new_group_size] + .copy_from_slice(level.new_hashes()); + } + + let mut siblings = [[0u8; 32]; N]; + siblings[..member_pos].copy_from_slice(&group[..member_pos]); + if member_pos + 1 < new_group_size { + siblings[member_pos..new_group_size - 1] + .copy_from_slice(&group[member_pos + 1..new_group_size]); + } + + new_levels[tree_level] = ProofLevel { + position: to_u8(member_pos)?, + sibling_count: to_u8(new_group_size - 1)?, + siblings, + }; + } + + let mid = if tree_level == 0 { + level.hashes[boundary_pos] + } else { + new_boundary_hash + }; + let child_count = boundary_pos + 1 + new_only; + new_boundary_hash = if child_count == 1 { + mid + } else { + children[..boundary_pos].copy_from_slice(&level.hashes[..boundary_pos]); + children[boundary_pos] = mid; + children[boundary_pos + 1..child_count] + .copy_from_slice(level.new_hashes()); + hasher.hash_children(&children[..child_count]) + }; + + member_idx /= N; + old_level_size = old_level_size.div_ceil(N); + } + + Ok(NaryProof { + root: self.new_root, + leaf: old_proof.leaf, + leaf_index: old_proof.leaf_index, + tree_size: self.new_size, + level_count: self.level_count, + levels: new_levels, + }) + } +} + +impl TreeSnapshot { + /// Generate an inclusion proof for the leaf at `leaf_index` + pub fn generate_proof( + &self, + leaf_index: u64, + ) -> Result, TreeError> { + if leaf_index >= self.size { + return Err(TreeError::IndexOutOfRange { + index: leaf_index, + size: self.size, + }); + } + + let leaf_idx = u64_to_usize(leaf_index)?; + let mut index = leaf_idx; + let mut levels = [ProofLevel::::EMPTY; MAX_DEPTH]; + + #[allow(clippy::needless_range_loop)] + for level in 0..self.depth { + let child_pos = index % N; + let group_start = index - child_pos; + let group_end = core::cmp::min(group_start + N, self.levels[level].len()); + let group_size = group_end - group_start; + + let mut group = [[0u8; 32]; N]; + self.levels[level].get_group(group_start, group_size, &mut group); + let mut siblings = [[0u8; 32]; N]; + siblings[..child_pos].copy_from_slice(&group[..child_pos]); + let rest = group_size - child_pos - 1; + siblings[child_pos..child_pos + rest] + .copy_from_slice(&group[child_pos + 1..group_size]); + let sib_count = group_size - 1; + levels[level] = ProofLevel { + position: to_u8(child_pos)?, + sibling_count: to_u8(sib_count)?, + siblings, + }; + + index /= N; + } + + Ok(NaryProof { + root: self.root.expect("set prev; qed"), + leaf: self.levels[0].get(leaf_idx)?, + leaf_index, + tree_size: self.size, + level_count: self.depth, + levels, + }) + } + + /// Generate a consistency proof proving the tree at `old_size` is a prefix + /// of the current tree + pub fn generate_consistency_proof( + &self, + old_size: u64, + old_root: Hash, + ) -> Result, TreeError> { + if old_size == 0 || self.size == 0 || old_size > self.size { + return Err(TreeError::IndexOutOfRange { + index: old_size, + size: self.size, + }); + } + + let new_root = self.root.expect("size > 0; qed"); + let mut levels = [ConsistencyLevel::::EMPTY; MAX_DEPTH]; + + if old_size == self.size { + return Ok(ConsistencyProof { + old_root, + new_root, + old_size, + new_size: self.size, + level_count: 0, + levels, + }); + } + + let mut old_level_size = u64_to_usize(old_size)?; + let mut new_level_size = u64_to_usize(self.size)?; + let mut depth = 0usize; + + while old_level_size > 1 || new_level_size > 1 { + if depth >= MAX_DEPTH { + return Err(TreeError::MathError); + } + + let boundary = old_level_size - 1; + let boundary_pos = boundary % N; + let group_start = boundary - boundary_pos; + let shared = if depth == 0 { + boundary_pos + 1 + } else { + boundary_pos + }; + let new_count = new_level_size.min(group_start + N) - boundary - 1; + + let mut hashes = [[0u8; 32]; N]; + self.levels[depth].get_group(group_start, shared, &mut hashes[..shared]); + self.levels[depth].get_group( + boundary + 1, + new_count, + &mut hashes[shared..shared + new_count], + ); + levels[depth] = ConsistencyLevel { + shared_count: to_u8(shared)?, + new_count: to_u8(new_count)?, + hashes, + }; + + depth += 1; + old_level_size = old_level_size.div_ceil(N); + new_level_size = new_level_size.div_ceil(N); + } + + Ok(ConsistencyProof { + old_root, + new_root, + old_size, + new_size: self.size, + level_count: depth, + levels, + }) + } +} + +#[cfg(test)] +#[cfg_attr(feature = "concurrent", allow(unused_mut))] +mod tests { + #[cfg(not(feature = "std"))] + use alloc::vec::Vec; + #[cfg(feature = "std")] + use std::vec::Vec; + + use super::*; + use crate::{ + LeanIMT, + test_util::*, + }; + + fn build_snapshots( + leaves: &[Hash], + ) -> Vec<(u64, Hash, TreeSnapshot)> { + let mut tree = LeanIMT::::new(XorHasher); + let mut snapshots = Vec::new(); + for &l in leaves { + tree.insert(l).unwrap(); + let snap = tree.snapshot(); + let root = snap.root().unwrap(); + let size = snap.size(); + snapshots.push((size, root, snap)); + } + snapshots + } + + /// Verify generate + verify for every (old, new) pair up to `count` leaves. + fn verify_consistency_all_pairs(count: u32) { + let leaves: Vec = (1..=count).map(leaf).collect(); + let snaps = build_snapshots::(&leaves); + for i in 0..snaps.len() { + for j in i..snaps.len() { + let proof = snaps[j] + .2 + .generate_consistency_proof(snaps[i].0, snaps[i].1) + .unwrap(); + assert!( + proof.verify(&XorHasher).unwrap(), + "N={} consistency failed for {} -> {}", + N, + i + 1, + j + 1 + ); + } + } + } + + /// Verify update_inclusion_proof for every (old, new, member) triple, + /// and assert the result equals a freshly generated proof. + fn verify_update_all_pairs(count: u32) { + let leaves: Vec = (1..=count).map(leaf).collect(); + let snaps = build_snapshots::(&leaves); + for i in 0..snaps.len() { + for j in i..snaps.len() { + let cp = snaps[j] + .2 + .generate_consistency_proof(snaps[i].0, snaps[i].1) + .unwrap(); + if i == j { + // Same-size update returns NoUpdateNeeded + let ip = snaps[i].2.generate_proof(0).unwrap(); + let err = cp.update_inclusion_proof(&ip, &XorHasher).unwrap_err(); + assert_eq!(err, TreeError::NoUpdateNeeded); + continue; + } + for m in 0..=i { + let old_ip = snaps[i].2.generate_proof(m as u64).unwrap(); + let updated = cp.update_inclusion_proof(&old_ip, &XorHasher).unwrap(); + assert!( + updated.verify(&XorHasher).unwrap(), + "N={} update failed: member {} from {} -> {}", + N, + m, + i + 1, + j + 1, + ); + assert_eq!(updated.root, snaps[j].1); + assert_eq!(updated.leaf, old_ip.leaf); + assert_eq!(updated.leaf_index, old_ip.leaf_index); + let fresh = snaps[j].2.generate_proof(m as u64).unwrap(); + assert_eq!( + updated, + fresh, + "N={} update mismatch: member {} from {} -> {}", + N, + m, + i + 1, + j + 1, + ); + } + } + } + } + + #[test] + fn proof_single_leaf() { + let th = XorHasher; + let mut tree = LeanIMT::::new(XorHasher); + let l = leaf(1); + tree.insert(l).unwrap(); + let snap = tree.snapshot(); + + let proof = snap.generate_proof(0).unwrap(); + assert_eq!(proof.leaf, l); + assert_eq!(proof.root, l); + assert_eq!(proof.leaf_index, 0); + assert_eq!(proof.level_count, 0); + assert!(proof.verify(&th).unwrap()); + } + + #[test] + fn proof_two_leaves_binary() { + let h = XorHasher; + let th = XorHasher; + let mut tree = LeanIMT::::new(h.clone()); + let l0 = leaf(1); + let l1 = leaf(2); + tree.insert(l0).unwrap(); + tree.insert(l1).unwrap(); + let snap = tree.snapshot(); + + let p0 = snap.generate_proof(0).unwrap(); + assert_eq!(p0.leaf, l0); + assert_eq!(p0.level_count, 1); + assert_eq!(p0.levels[0].position, 0); + assert_eq!(p0.levels[0].sibling_count, 1); + assert_eq!(p0.levels[0].siblings[0], l1); + assert!(p0.verify(&th).unwrap()); + + let p1 = snap.generate_proof(1).unwrap(); + assert_eq!(p1.leaf, l1); + assert_eq!(p1.levels[0].position, 1); + assert_eq!(p1.levels[0].sibling_count, 1); + assert_eq!(p1.levels[0].siblings[0], l0); + assert!(p1.verify(&th).unwrap()); + } + + #[test] + fn proof_three_leaves_binary_lifted() { + let h = XorHasher; + let th = XorHasher; + let mut tree = LeanIMT::::new(h.clone()); + let l0 = leaf(1); + let l1 = leaf(2); + let l2 = leaf(3); + tree.insert(l0).unwrap(); + tree.insert(l1).unwrap(); + tree.insert(l2).unwrap(); + let snap = tree.snapshot(); + + let p = snap.generate_proof(2).unwrap(); + assert_eq!(p.level_count, 2); + assert_eq!(p.levels[0].sibling_count, 0); + let h01 = th.hash_children(&[l0, l1]); + assert_eq!(p.levels[1].position, 1); + assert_eq!(p.levels[1].sibling_count, 1); + assert_eq!(p.levels[1].siblings[0], h01); + assert!(p.verify(&th).unwrap()); + } + + #[test] + fn proof_four_leaves_binary() { + let h = XorHasher; + let th = XorHasher; + let mut tree = LeanIMT::::new(h.clone()); + let leaves: Vec = (1..=4).map(leaf).collect(); + for &l in &leaves { + tree.insert(l).unwrap(); + } + let snap = tree.snapshot(); + + for i in 0..4u64 { + let p = snap.generate_proof(i).unwrap(); + assert!(p.verify(&th).unwrap()); + assert_eq!(p.leaf, leaves[i as usize]); + } + } + + #[test] + fn proof_ternary() { + let h = XorHasher; + let th = XorHasher; + let mut tree = LeanIMT::::new(h.clone()); + let leaves: Vec = (1..=4).map(leaf).collect(); + for &l in &leaves { + tree.insert(l).unwrap(); + } + let snap = tree.snapshot(); + + for i in 0..4u64 { + let p = snap.generate_proof(i).unwrap(); + assert!(p.verify(&th).unwrap()); + } + } + + #[test] + fn proof_quaternary() { + let h = XorHasher; + let th = XorHasher; + let mut tree = LeanIMT::::new(h.clone()); + let leaves: Vec = (1..=5).map(leaf).collect(); + for &l in &leaves { + tree.insert(l).unwrap(); + } + let snap = tree.snapshot(); + + for i in 0..5u64 { + let p = snap.generate_proof(i).unwrap(); + assert!(p.verify(&th).unwrap()); + } + } + + #[test] + fn verify_rejects_wrong_leaf() { + let h = XorHasher; + let th = XorHasher; + let mut tree = LeanIMT::::new(h.clone()); + tree.insert(leaf(1)).unwrap(); + tree.insert(leaf(2)).unwrap(); + let snap = tree.snapshot(); + + let mut proof = snap.generate_proof(0).unwrap(); + proof.leaf = leaf(99); + assert!(!proof.verify(&th).unwrap()); + } + + #[test] + fn verify_rejects_wrong_root() { + let h = XorHasher; + let th = XorHasher; + let mut tree = LeanIMT::::new(h.clone()); + tree.insert(leaf(1)).unwrap(); + tree.insert(leaf(2)).unwrap(); + let snap = tree.snapshot(); + + let mut proof = snap.generate_proof(0).unwrap(); + proof.root = [0xFF; 32]; + assert!(!proof.verify(&th).unwrap()); + } + + #[test] + fn verify_rejects_wrong_sibling() { + let h = XorHasher; + let th = XorHasher; + let mut tree = LeanIMT::::new(h.clone()); + tree.insert(leaf(1)).unwrap(); + tree.insert(leaf(2)).unwrap(); + let snap = tree.snapshot(); + + let mut proof = snap.generate_proof(0).unwrap(); + proof.levels[0].siblings[0] = [0xFF; 32]; + assert!(!proof.verify(&th).unwrap()); + } + + #[test] + fn proof_index_out_of_range() { + let mut tree = LeanIMT::::new(XorHasher); + tree.insert(leaf(1)).unwrap(); + let snap = tree.snapshot(); + + let err = snap.generate_proof(1).unwrap_err(); + assert_eq!(err, TreeError::IndexOutOfRange { index: 1, size: 1 }); + } + + #[cfg(feature = "blake3")] + #[test] + fn proof_blake3_round_trip() { + use crate::Blake3Hasher; + let h = Blake3Hasher; + let th = Blake3Hasher; + let mut tree = LeanIMT::::new(h); + + for i in 0u8..20 { + tree.insert(*::blake3::hash(&[i]).as_bytes()).unwrap(); + } + let snap = tree.snapshot(); + + for i in 0..20u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&th).unwrap()); + } + } + + #[cfg(feature = "wincode")] + fn verify_wincode_proof_round_trip(count: u32) { + let leaves: Vec = (1..=count).map(leaf).collect(); + let mut tree = LeanIMT::::new(XorHasher); + for &l in &leaves { + tree.insert(l).unwrap(); + } + let snap = tree.snapshot(); + for i in 0..count as u64 { + let proof = snap.generate_proof(i).unwrap(); + let bytes = wincode::serialize(&proof).unwrap(); + let decoded: NaryProof = wincode::deserialize(&bytes).unwrap(); + assert_eq!(decoded, proof); + assert!(decoded.verify(&XorHasher).unwrap()); + } + } + + #[cfg(feature = "wincode")] + #[test] + fn wincode_round_trip_binary() { + verify_wincode_proof_round_trip::<2>(10); + } + + #[cfg(feature = "wincode")] + #[test] + fn wincode_round_trip_ternary() { + verify_wincode_proof_round_trip::<3>(10); + } + + #[cfg(feature = "wincode")] + #[test] + fn wincode_round_trip_quaternary() { + verify_wincode_proof_round_trip::<4>(10); + } + + #[test] + fn consistency_same_size_trivial() { + let mut tree = LeanIMT::::new(XorHasher); + tree.insert(leaf(1)).unwrap(); + tree.insert(leaf(2)).unwrap(); + let snap = tree.snapshot(); + let root = snap.root().unwrap(); + let proof = snap.generate_consistency_proof(2, root).unwrap(); + assert_eq!(proof.level_count, 0); + assert!(proof.verify(&XorHasher).unwrap()); + } + + #[test] + fn consistency_binary_all_pairs_small() { + verify_consistency_all_pairs::<2>(8); + } + + #[test] + fn consistency_ternary_all_pairs_small() { + verify_consistency_all_pairs::<3>(9); + } + + #[test] + fn consistency_quaternary_all_pairs_small() { + verify_consistency_all_pairs::<4>(8); + } + + #[test] + fn consistency_rejects_wrong_old_root() { + let leaves: Vec = (1..=4).map(leaf).collect(); + let snaps = build_snapshots::<2, 32>(&leaves); + let proof = snaps[3] + .2 + .generate_consistency_proof(snaps[1].0, snaps[1].1) + .unwrap(); + let mut tampered = proof; + tampered.old_root = [0xFF; 32]; + assert!(!tampered.verify(&XorHasher).unwrap()); + } + + #[test] + fn consistency_rejects_wrong_new_root() { + let leaves: Vec = (1..=4).map(leaf).collect(); + let snaps = build_snapshots::<2, 32>(&leaves); + let proof = snaps[3] + .2 + .generate_consistency_proof(snaps[1].0, snaps[1].1) + .unwrap(); + let mut tampered = proof; + tampered.new_root = [0xFF; 32]; + assert!(!tampered.verify(&XorHasher).unwrap()); + } + + #[test] + fn consistency_rejects_tampered_hash() { + let leaves: Vec = (1..=4).map(leaf).collect(); + let snaps = build_snapshots::<2, 32>(&leaves); + let proof = snaps[3] + .2 + .generate_consistency_proof(snaps[1].0, snaps[1].1) + .unwrap(); + let mut tampered = proof; + if tampered.level_count > 0 { + tampered.levels[0].hashes[0] = [0xFF; 32]; + } + assert!(!tampered.verify(&XorHasher).unwrap()); + } + + #[test] + fn consistency_rejects_invalid_sizes() { + let mut tree = LeanIMT::::new(XorHasher); + tree.insert(leaf(1)).unwrap(); + let snap = tree.snapshot(); + assert!(snap.generate_consistency_proof(0, [0u8; 32]).is_err()); + assert!(snap.generate_consistency_proof(2, [0u8; 32]).is_err()); + } + + #[test] + fn consistency_verify_rejects_zero_sizes() { + let proof = ConsistencyProof::<2, 32> { + old_root: [0u8; 32], + new_root: [0u8; 32], + old_size: 0, + new_size: 1, + level_count: 0, + levels: [ConsistencyLevel::EMPTY; 32], + }; + assert!(proof.verify(&XorHasher).is_err()); + } + + #[test] + fn consistency_verify_rejects_old_gt_new() { + let proof = ConsistencyProof::<2, 32> { + old_root: [0u8; 32], + new_root: [0u8; 32], + old_size: 5, + new_size: 3, + level_count: 0, + levels: [ConsistencyLevel::EMPTY; 32], + }; + assert!(proof.verify(&XorHasher).is_err()); + } + + #[test] + fn update_same_size() { + let mut tree = LeanIMT::::new(XorHasher); + tree.insert(leaf(1)).unwrap(); + tree.insert(leaf(2)).unwrap(); + let snap = tree.snapshot(); + let root = snap.root().unwrap(); + let cp = snap.generate_consistency_proof(2, root).unwrap(); + let ip = snap.generate_proof(0).unwrap(); + let err = cp.update_inclusion_proof(&ip, &XorHasher).unwrap_err(); + assert_eq!(err, TreeError::NoUpdateNeeded); + } + + #[test] + fn update_binary_all_pairs() { + verify_update_all_pairs::<2>(8); + } + + #[test] + fn update_ternary_all_pairs() { + verify_update_all_pairs::<3>(9); + } + + #[test] + fn update_quaternary_all_pairs() { + verify_update_all_pairs::<4>(8); + } + + #[test] + fn update_rejects_wrong_root() { + let leaves: Vec = (1..=4).map(leaf).collect(); + let snaps = build_snapshots::<2, 32>(&leaves); + let cp = snaps[3] + .2 + .generate_consistency_proof(snaps[1].0, snaps[1].1) + .unwrap(); + let mut bad_proof = snaps[1].2.generate_proof(0).unwrap(); + bad_proof.root = [0xFF; 32]; + assert!(cp.update_inclusion_proof(&bad_proof, &XorHasher).is_err()); + } + + #[cfg(feature = "wincode")] + fn verify_wincode_consistency_round_trip() { + let leaves: Vec = (1..=6).map(leaf).collect(); + let snaps = build_snapshots::(&leaves); + let proof = snaps[5] + .2 + .generate_consistency_proof(snaps[2].0, snaps[2].1) + .unwrap(); + let bytes = wincode::serialize(&proof).unwrap(); + let decoded: ConsistencyProof = wincode::deserialize(&bytes).unwrap(); + assert_eq!(decoded, proof); + assert!(decoded.verify(&XorHasher).unwrap()); + } + + #[cfg(feature = "wincode")] + #[test] + fn consistency_wincode_binary() { + verify_wincode_consistency_round_trip::<2>(); + } + + #[cfg(feature = "wincode")] + #[test] + fn consistency_wincode_ternary() { + verify_wincode_consistency_round_trip::<3>(); + } + + #[cfg(feature = "wincode")] + #[test] + fn consistency_wincode_quaternary() { + verify_wincode_consistency_round_trip::<4>(); + } + + #[test] + fn verify_rejects_level_count_exceeds_max_depth() { + let proof = NaryProof::<2, 4> { + root: [0u8; 32], + leaf: [0u8; 32], + leaf_index: 0, + tree_size: 2, + level_count: 5, + levels: [ProofLevel::EMPTY; 4], + }; + assert!(proof.verify(&XorHasher).is_err()); + } + + #[test] + fn verify_rejects_position_exceeds_n() { + let mut tree = LeanIMT::::new(XorHasher); + tree.insert(leaf(1)).unwrap(); + tree.insert(leaf(2)).unwrap(); + let snap = tree.snapshot(); + let mut proof = snap.generate_proof(0).unwrap(); + proof.levels[0].position = 2; + assert!(proof.verify(&XorHasher).is_err()); + } + + #[test] + fn verify_rejects_sibling_count_exceeds_n() { + let mut tree = LeanIMT::::new(XorHasher); + tree.insert(leaf(1)).unwrap(); + tree.insert(leaf(2)).unwrap(); + let snap = tree.snapshot(); + let mut proof = snap.generate_proof(0).unwrap(); + proof.levels[0].sibling_count = 2; + assert!(proof.verify(&XorHasher).is_err()); + } + + #[test] + fn consistency_verify_rejects_level_count_exceeds_max_depth() { + let proof = ConsistencyProof::<2, 4> { + old_root: [0u8; 32], + new_root: [0u8; 32], + old_size: 1, + new_size: 2, + level_count: 5, + levels: [ConsistencyLevel::EMPTY; 4], + }; + assert!(proof.verify(&XorHasher).is_err()); + } + + #[test] + fn verify_rejects_spoofed_leaf_index() { + let h = XorHasher; + let th = XorHasher; + let mut tree = LeanIMT::::new(h.clone()); + tree.insert(leaf(1)).unwrap(); + tree.insert(leaf(2)).unwrap(); + let snap = tree.snapshot(); + + let mut proof = snap.generate_proof(0).unwrap(); + proof.leaf_index = 1; // spoof: claim index 1 instead of 0 + assert!(proof.verify(&th).is_err()); + } + + #[test] + fn verify_rejects_padded_levels() { + let h = XorHasher; + let th = XorHasher; + let mut tree = LeanIMT::::new(h.clone()); + tree.insert(leaf(1)).unwrap(); + tree.insert(leaf(2)).unwrap(); + let snap = tree.snapshot(); + + let mut proof = snap.generate_proof(0).unwrap(); + // Pad with an extra zero-sibling level + proof.level_count = 2; + proof.levels[1] = ProofLevel::EMPTY; + assert!(proof.verify(&th).is_err()); + } + + #[test] + fn verify_rejects_truncated_levels() { + let h = XorHasher; + let th = XorHasher; + let mut tree = LeanIMT::::new(h.clone()); + for i in 1..=4u32 { + tree.insert(leaf(i)).unwrap(); + } + let snap = tree.snapshot(); + + let mut proof = snap.generate_proof(2).unwrap(); + // Truncate: claim fewer levels than the tree requires + proof.level_count = 1; + assert!(proof.verify(&th).is_err()); + } + + #[test] + fn verify_rejects_tree_size_zero() { + let proof = NaryProof::<2, 4> { + root: [0u8; 32], + leaf: [0u8; 32], + leaf_index: 0, + tree_size: 0, + level_count: 0, + levels: [ProofLevel::EMPTY; 4], + }; + assert!(proof.verify(&XorHasher).is_err()); + } + + #[test] + fn verify_rejects_leaf_index_ge_tree_size() { + let h = XorHasher; + let th = XorHasher; + let mut tree = LeanIMT::::new(h.clone()); + tree.insert(leaf(1)).unwrap(); + tree.insert(leaf(2)).unwrap(); + let snap = tree.snapshot(); + + let mut proof = snap.generate_proof(0).unwrap(); + proof.leaf_index = 2; // >= tree_size of 2 + assert!(proof.verify(&th).is_err()); + } + + #[test] + fn verify_rejects_tree_size_exceeding_max_depth() { + // tree_size=32 with N=2 requires depth 5, but MAX_DEPTH=4. + // Must not panic on self.levels[..5] slice. + let proof = NaryProof::<2, 4> { + root: [0u8; 32], + leaf: [0u8; 32], + leaf_index: 0, + tree_size: 32, + level_count: 5, + levels: [ProofLevel::EMPTY; 4], + }; + match proof.verify(&XorHasher) { + Err(TreeError::InvalidProofDepth { + expected: 5, + actual: 5, + }) => {} + other => panic!("expected InvalidProofDepth, got {other:?}"), + } + } +} diff --git a/crates/rotortree/src/storage/checkpoint.rs b/crates/rotortree/src/storage/checkpoint.rs new file mode 100644 index 0000000..67521d4 --- /dev/null +++ b/crates/rotortree/src/storage/checkpoint.rs @@ -0,0 +1,371 @@ +use std::{ + fs, + io::{ + self, + Seek, + Write, + }, + path::{ + Path, + PathBuf, + }, + sync::Arc, +}; + +use crate::{ + Hash, + chunked_level::CHUNK_SIZE, +}; + +use super::{ + data::MmapRegion, + error::StorageError, + frame, +}; + +const HEADER_MAGIC: [u8; 4] = *b"RTRD"; +const META_MAGIC: [u8; 4] = *b"RTMD"; + +pub(crate) const CHUNK_BYTE_SIZE: usize = CHUNK_SIZE * 32; +pub(crate) const CHUNKS_PER_SHARD: usize = 65_536; + +/// Returns `(shard_index, byte_offset_within_shard)` for a given chunk index. +#[inline] +pub(crate) fn shard_address(chunk_idx: usize) -> (usize, usize) { + let shard_idx = chunk_idx / CHUNKS_PER_SHARD; + let offset_in_shard = (chunk_idx % CHUNKS_PER_SHARD) * CHUNK_BYTE_SIZE; + (shard_idx, offset_in_shard) +} + +/// Controls when checkpoints are triggered +#[derive(Default)] +pub enum CheckpointPolicy { + /// Caller calls `checkpoint()` explicitly + #[default] + Manual, + /// Auto-checkpoint after every N WAL entries + EveryNEntries(u64), + /// Auto-checkpoint when in-memory chunks exceed N bytes + MemoryThreshold(usize), + /// Checkpoint only on graceful close + OnClose, +} + +/// Controls which tree levels are kept in memory vs mmap'd +pub struct TieringConfig { + /// Levels below this value have their committed chunks mmap'd after checkpoint. + /// Set to `usize::MAX` to mmap all levels (default), `0` to keep everything in memory + pub pin_above_level: usize, +} + +impl Default for TieringConfig { + fn default() -> Self { + Self { + pin_above_level: usize::MAX, + } + } +} + +/// Versioned header written once at data directory creation +#[derive(Debug, wincode::SchemaWrite, wincode::SchemaRead)] +enum HeaderFrame { + V1 { + magic: [u8; 4], + n: u32, + max_depth: u32, + chunk_size: u32, + }, +} + +/// Versioned checkpoint metadata (atomically written at each checkpoint) +#[derive(Debug, wincode::SchemaWrite, wincode::SchemaRead)] +enum MetaFrame { + V1 { + magic: [u8; 4], + n: u32, + max_depth: u32, + last_wal_seq: u64, + leaf_count: u64, + depth: u32, + root_hash: Hash, + }, +} + +/// Deserialized header.bin contents +#[derive(Debug, Clone, Copy)] +pub(crate) struct HeaderData { + pub(crate) n: u32, + pub(crate) max_depth: u32, + pub(crate) chunk_size: u32, +} + +/// Deserialized checkpoint.meta contents +#[doc(hidden)] +#[derive(Debug, Clone, Copy)] +pub struct CheckpointMeta { + pub n: u32, + pub max_depth: u32, + pub last_wal_seq: u64, + pub leaf_count: u64, + pub depth: u32, + pub root_hash: Hash, +} + +/// Write `header.bin` +pub(crate) fn write_header(data_dir: &Path, n: u32, max_depth: u32) -> io::Result<()> { + #[allow(clippy::cast_possible_truncation)] + let header = HeaderFrame::V1 { + magic: HEADER_MAGIC, + n, + max_depth, + chunk_size: CHUNK_SIZE as u32, + }; + let buf = frame::serialize_frame(&header); + atomic_write(&data_dir.join("header.bin"), &buf) +} + +/// Read and validate `header.bin`. Returns `None` if missing or corrupt +pub(crate) fn read_header(data_dir: &Path) -> Result, StorageError> { + let path = data_dir.join("header.bin"); + let header: HeaderFrame = match frame::read_frame_file(&path)? { + Some(h) => h, + None => return Ok(None), + }; + + let HeaderFrame::V1 { + magic, + n, + max_depth, + chunk_size, + } = header; + + if magic != HEADER_MAGIC { + return Ok(None); + } + + Ok(Some(HeaderData { + n, + max_depth, + chunk_size, + })) +} + +/// Write checkpoint metadata atomically +pub(crate) fn write_meta(data_dir: &Path, meta: &CheckpointMeta) -> io::Result<()> { + let frame = MetaFrame::V1 { + magic: META_MAGIC, + n: meta.n, + max_depth: meta.max_depth, + last_wal_seq: meta.last_wal_seq, + leaf_count: meta.leaf_count, + depth: meta.depth, + root_hash: meta.root_hash, + }; + let buf = frame::serialize_frame(&frame); + atomic_write(&data_dir.join("checkpoint.meta"), &buf) +} + +/// Read and validate checkpoint metadata. Returns `None` if missing or corrupt +pub(crate) fn read_meta(data_dir: &Path) -> Result, StorageError> { + let path = data_dir.join("checkpoint.meta"); + let meta_frame: MetaFrame = match frame::read_frame_file(&path)? { + Some(m) => m, + None => return Ok(None), + }; + + let MetaFrame::V1 { + magic, + n, + max_depth, + last_wal_seq, + leaf_count, + depth, + root_hash, + } = meta_frame; + + if magic != META_MAGIC { + return Ok(None); + } + + Ok(Some(CheckpointMeta { + n, + max_depth, + last_wal_seq, + leaf_count, + depth, + root_hash, + })) +} + +/// Write all level tails atomically (tmp -> fsync -> rename). +pub(crate) fn write_tails( + data_dir: &Path, + tails: &[[Hash; CHUNK_SIZE]], + max_depth: usize, +) -> io::Result<()> { + let total_size = max_depth * CHUNK_BYTE_SIZE; + let mut buf = vec![0u8; total_size]; + + for (i, tail) in tails.iter().enumerate() { + let base = i * CHUNK_BYTE_SIZE; + buf[base..base + CHUNK_BYTE_SIZE].copy_from_slice(tail.as_flattened()); + } + + atomic_write(&data_dir.join("tails.bin"), &buf) +} + +/// Read all tails from disk. Returns `None` if the file is missing or wrong size +pub(crate) fn read_tails( + data_dir: &Path, + max_depth: usize, +) -> io::Result>> { + let path = data_dir.join("tails.bin"); + let data = match fs::read(&path) { + Ok(d) => d, + Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(None), + Err(e) => return Err(e), + }; + let expected = max_depth * CHUNK_BYTE_SIZE; + if data.len() != expected { + return Ok(None); + } + + let mut tails = Vec::with_capacity(max_depth); + for i in 0..max_depth { + let base = i * CHUNK_BYTE_SIZE; + let chunk = &data[base..base + CHUNK_BYTE_SIZE]; + let mut tail = [[0u8; 32]; CHUNK_SIZE]; + tail.as_flattened_mut().copy_from_slice(chunk); + tails.push(tail); + } + + Ok(Some(tails)) +} + +pub(crate) fn level_dir_path(data_dir: &Path, level_idx: usize) -> PathBuf { + data_dir.join(format!("level_{level_idx}")) +} + +pub(crate) fn shard_file_path( + data_dir: &Path, + level_idx: usize, + shard_idx: usize, +) -> PathBuf { + data_dir.join(format!("level_{level_idx}/shard_{shard_idx:04}.dat")) +} + +pub(crate) fn append_chunks_to_level<'a>( + data_dir: &Path, + level_idx: usize, + from_chunk: usize, + chunks: impl Iterator, +) -> io::Result> { + let chunks: Vec<_> = chunks.collect(); + if chunks.is_empty() { + return Ok(Vec::new()); + } + + let first_shard = from_chunk / CHUNKS_PER_SHARD; + let last_shard = (from_chunk + chunks.len() - 1) / CHUNKS_PER_SHARD; + + (first_shard..=last_shard) + .map(|shard_idx| { + let shard_start = shard_idx * CHUNKS_PER_SHARD; + let local_start = shard_start.saturating_sub(from_chunk); + let local_end = ((shard_idx + 1) * CHUNKS_PER_SHARD) + .saturating_sub(from_chunk) + .min(chunks.len()); + let (_, offset) = shard_address(from_chunk + local_start); + + let path = shard_file_path(data_dir, level_idx, shard_idx); + let mut file = fs::OpenOptions::new() + .create(true) + .write(true) + .read(true) + .truncate(false) + .open(&path)?; + file.seek(io::SeekFrom::Start(offset as u64))?; + + for chunk in &chunks[local_start..local_end] { + file.write_all(chunk.as_flattened())?; + } + + Ok(file) + }) + .collect() +} + +pub(crate) fn mmap_level_shards( + data_dir: &Path, + level_idx: usize, + valid_chunks: usize, +) -> io::Result>> { + if valid_chunks == 0 { + return Ok(Vec::new()); + } + + let shard_count = valid_chunks.div_ceil(CHUNKS_PER_SHARD); + let mut regions = Vec::with_capacity(shard_count); + + for shard_idx in 0..shard_count { + let path = shard_file_path(data_dir, level_idx, shard_idx); + let file = match fs::File::open(&path) { + Ok(f) => f, + Err(e) if e.kind() == io::ErrorKind::NotFound && shard_idx == 0 => { + return Ok(Vec::new()); + } + Err(e) if e.kind() == io::ErrorKind::NotFound => { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "level_{level_idx}/shard_{shard_idx:04}.dat missing \ + (shards 0..{shard_idx} exist)" + ), + )); + } + Err(e) => return Err(e), + }; + + let chunks_in_shard = + CHUNKS_PER_SHARD.min(valid_chunks - shard_idx * CHUNKS_PER_SHARD); + let valid_len = chunks_in_shard * CHUNK_BYTE_SIZE; + let file_len = usize::try_from(file.metadata()?.len()).unwrap_or(usize::MAX); + + if file_len < valid_len { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "level_{level_idx}/shard_{shard_idx:04}.dat too small: \ + {file_len} bytes < {valid_len} expected" + ), + )); + } + + let mmap = unsafe { memmap2::MmapOptions::new().len(file_len).map_copy(&file)? }; + regions.push(Arc::new(MmapRegion::new(mmap, valid_len))); + } + + Ok(regions) +} + +pub(crate) fn atomic_write(path: &Path, data: &[u8]) -> io::Result<()> { + let mut s = path.as_os_str().to_owned(); + s.push(".tmp"); + let tmp = PathBuf::from(s); + { + let mut file = fs::File::create(&tmp)?; + file.write_all(data)?; + file.sync_all()?; + } + fs::rename(&tmp, path)?; + if let Some(parent) = path.parent() { + let dir = fs::File::open(parent)?; + dir.sync_all()?; + } + Ok(()) +} + +#[doc(hidden)] +pub fn write_test_meta(data_dir: &Path, meta: &CheckpointMeta) -> io::Result<()> { + write_meta(data_dir, meta) +} diff --git a/crates/rotortree/src/storage/config.rs b/crates/rotortree/src/storage/config.rs new file mode 100644 index 0000000..d72ffd5 --- /dev/null +++ b/crates/rotortree/src/storage/config.rs @@ -0,0 +1,37 @@ +use std::{ + path::PathBuf, + time::Duration, +}; + +use super::checkpoint::{ + CheckpointPolicy, + TieringConfig, +}; + +/// Configuration for opening a `RotorTree` +pub struct RotorTreeConfig { + /// Directory path where the WAL and data files are stored + pub path: PathBuf, + /// Controls when WAL entries are fsynced to disk + pub flush_policy: FlushPolicy, + /// Controls when checkpoints are triggered + pub checkpoint_policy: CheckpointPolicy, + /// Controls which tree levels are kept in memory vs mmap'd + pub tiering: TieringConfig, + /// Recompute Merkle root on recovery to detect corruption beyond CRC + pub verify_checkpoint: bool, +} + +/// Controls when buffered WAL entries are fsynced to disk +pub enum FlushPolicy { + /// Fsync on a periodic interval (default: 10ms) + Interval(Duration), + /// Caller controls flushing via `flush()` + Manual, +} + +impl Default for FlushPolicy { + fn default() -> Self { + Self::Interval(Duration::from_millis(10)) + } +} diff --git a/crates/rotortree/src/storage/data.rs b/crates/rotortree/src/storage/data.rs new file mode 100644 index 0000000..4d83a4d --- /dev/null +++ b/crates/rotortree/src/storage/data.rs @@ -0,0 +1,19 @@ +/// mmap wrapper +pub(crate) struct MmapRegion { + mmap: memmap2::MmapMut, + valid_len: usize, +} + +impl MmapRegion { + pub(crate) fn new(mmap: memmap2::MmapMut, valid_len: usize) -> Self { + Self { mmap, valid_len } + } + + pub(crate) fn as_ptr(&self) -> *const u8 { + self.mmap.as_ptr() + } + + pub(crate) fn valid_len(&self) -> usize { + self.valid_len + } +} diff --git a/crates/rotortree/src/storage/error.rs b/crates/rotortree/src/storage/error.rs new file mode 100644 index 0000000..0d79748 --- /dev/null +++ b/crates/rotortree/src/storage/error.rs @@ -0,0 +1,163 @@ +use std::{ + fmt, + sync::Arc, +}; + +use crate::TreeError; + +#[derive(Debug)] +pub(crate) enum BackgroundError { + FlushFailed(Arc), + CheckpointFailed(String), +} + +/// Errors from WAL and storage operations +#[derive(Debug)] +pub enum StorageError { + /// An I/O error occurred + Io(std::io::Error), + /// CRC32C checksum mismatch at a given file offset + CrcMismatch { + offset: u64, + expected: u32, + actual: u32, + }, + /// WAL file is corrupted at the given offset (not a tail truncation) + WalCorrupted { offset: u64 }, + /// WAL file header has different N or MAX_DEPTH than expected + ConfigMismatch { + expected_n: u32, + actual_n: u32, + expected_max_depth: u32, + actual_max_depth: u32, + }, + /// Another process holds an exclusive lock on the WAL file + FileLocked, + /// The tree has been closed; no further operations are allowed + Closed, + /// A tree operation failed during WAL recovery + Tree(TreeError), + /// The background flush thread encountered an I/O error + FlushFailed(Arc), + /// The background checkpoint thread encountered an error + CheckpointFailed(String), + /// Math error + MathError, + /// Data corruption detected (e.g., root recomputation mismatch) + DataCorruption { detail: String }, + /// Frame deserialization failed despite valid CRC (schema/version mismatch) + SerdeFailed { path: String }, +} + +impl fmt::Display for StorageError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Io(e) => write!(f, "storage I/O error: {e}"), + Self::CrcMismatch { + offset, + expected, + actual, + } => write!( + f, + "CRC mismatch at offset {offset}: expected {expected:#010x}, \ + got {actual:#010x}" + ), + Self::WalCorrupted { offset } => { + write!(f, "WAL corrupted at offset {offset}") + } + Self::ConfigMismatch { + expected_n, + actual_n, + expected_max_depth, + actual_max_depth, + } => write!( + f, + "WAL config mismatch: expected N={expected_n}, \ + MAX_DEPTH={expected_max_depth}; found N={actual_n}, \ + MAX_DEPTH={actual_max_depth}" + ), + Self::FileLocked => { + write!(f, "WAL file is locked by another process") + } + Self::Closed => write!(f, "tree has been closed"), + Self::Tree(e) => write!(f, "tree error during recovery: {e}"), + Self::FlushFailed(e) => { + write!(f, "background flush failed: {e}") + } + Self::CheckpointFailed(detail) => { + write!(f, "background checkpoint failed: {detail}") + } + Self::MathError => { + write!(f, "math error") + } + Self::DataCorruption { detail } => { + write!(f, "data corruption: {detail}") + } + Self::SerdeFailed { path } => { + write!(f, "frame deserialization failed for {path}") + } + } + } +} + +impl std::error::Error for StorageError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Io(e) => Some(e), + Self::Tree(e) => Some(e), + Self::FlushFailed(e) => Some(e.as_ref()), + _ => None, + } + } +} + +impl From for StorageError { + fn from(e: std::io::Error) -> Self { + Self::Io(e) + } +} + +/// Unified error type for `RotorTree` operations +#[derive(Debug)] +pub enum RotorTreeError { + /// A tree-level error (depth exceeded, capacity, etc.) + Tree(TreeError), + /// A storage-level error (I/O, corruption, etc.) + Storage(StorageError), +} + +impl fmt::Display for RotorTreeError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Tree(e) => write!(f, "{e}"), + Self::Storage(e) => write!(f, "{e}"), + } + } +} + +impl std::error::Error for RotorTreeError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Tree(e) => Some(e), + Self::Storage(e) => Some(e), + } + } +} + +impl From for RotorTreeError { + fn from(e: TreeError) -> Self { + Self::Tree(e) + } +} + +impl From for RotorTreeError { + fn from(e: StorageError) -> Self { + Self::Storage(e) + } +} + +impl From for RotorTreeError { + fn from(e: std::io::Error) -> Self { + Self::Storage(StorageError::Io(e)) + } +} diff --git a/crates/rotortree/src/storage/frame.rs b/crates/rotortree/src/storage/frame.rs new file mode 100644 index 0000000..d2fcd7c --- /dev/null +++ b/crates/rotortree/src/storage/frame.rs @@ -0,0 +1,118 @@ +use super::error::StorageError; + +pub(super) const MAX_FRAME_PAYLOAD: u32 = 128 * 1024 * 1024; + +/// Write a length-prefixed, CRC32C-checksummed frame into `buf` +pub(super) fn write_frame< + T: wincode::SchemaWrite, +>( + buf: &mut Vec, + value: &T, +) { + let size = wincode::serialized_size(value).expect("serialized_size cannot fail"); + assert!( + size <= MAX_FRAME_PAYLOAD as u64, + "frame exceeds maximum payload size ({size} > {MAX_FRAME_PAYLOAD})" + ); + + #[allow(clippy::cast_possible_truncation)] + buf.reserve(size as usize + 8); + let start = buf.len(); + buf.extend_from_slice( + &(u32::try_from(size).expect("size fits in u32; qed")).to_le_bytes(), + ); + wincode::serialize_into(&mut *buf, value).expect("serialize_into cannot fail"); + let crc = crc_fast::crc32_iscsi(&buf[start..]); + buf.extend_from_slice(&crc.to_le_bytes()); +} + +/// Serialize a value into a self-contained frame buffer +pub(super) fn serialize_frame< + T: wincode::SchemaWrite, +>( + value: &T, +) -> Vec { + let mut buf = Vec::new(); + write_frame(&mut buf, value); + buf +} + +/// Read a frame at `offset`. Returns the payload slice and total frame size +pub(super) fn read_frame( + data: &[u8], + offset: usize, +) -> Result, StorageError> { + let remaining = data.len().saturating_sub(offset); + if remaining < 4 { + return Ok(None); + } + + let end = offset.checked_add(4).ok_or(StorageError::MathError)?; + let len = u32::from_le_bytes(data[offset..end].try_into().unwrap()); + if len > MAX_FRAME_PAYLOAD { + return Ok(None); + } + + let len_usize = len as usize; + let frame_size = 4usize + .checked_add(len_usize) + .and_then(|v| v.checked_add(4)) + .ok_or(StorageError::MathError)?; + if remaining < frame_size { + return Ok(None); + } + + let crc_offset = offset + .checked_add(4) + .and_then(|v| v.checked_add(len_usize)) + .ok_or(StorageError::MathError)?; + let crc_end = crc_offset.checked_add(4).ok_or(StorageError::MathError)?; + let stored_crc = u32::from_le_bytes(data[crc_offset..crc_end].try_into().unwrap()); + let computed_crc = crc_fast::crc32_iscsi(&data[offset..crc_offset]); + + if stored_crc != computed_crc { + if offset + .checked_add(frame_size) + .ok_or(StorageError::MathError)? + >= data.len() + { + return Ok(None); // tail truncation + } + return Err(StorageError::CrcMismatch { + offset: offset as u64, + expected: stored_crc, + actual: computed_crc, + }); + } + + let payload_start = offset.checked_add(4).ok_or(StorageError::MathError)?; + Ok(Some((&data[payload_start..crc_offset], frame_size))) +} + +/// Deserialize a single frame file (e.g., header.bin, checkpoint.meta) +pub(super) fn read_frame_file( + path: &std::path::Path, +) -> Result, StorageError> +where + T: for<'a> wincode::SchemaRead<'a, wincode::config::DefaultConfig, Dst = T>, +{ + let data = match std::fs::read(path) { + Ok(d) => d, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(e) => return Err(StorageError::Io(e)), + }; + + let (payload, _) = match read_frame(&data, 0) { + Ok(Some(r)) => r, + Ok(None) => return Ok(None), + Err(StorageError::CrcMismatch { .. }) => return Ok(None), + Err(e) => return Err(e), + }; + + match wincode::deserialize(payload) { + Ok(v) => Ok(Some(v)), + Err(_) => Err(StorageError::SerdeFailed { + path: path.display().to_string(), + }), + } +} diff --git a/crates/rotortree/src/storage/mod.rs b/crates/rotortree/src/storage/mod.rs new file mode 100644 index 0000000..1481b9a --- /dev/null +++ b/crates/rotortree/src/storage/mod.rs @@ -0,0 +1,816 @@ +pub(crate) mod checkpoint; +mod config; +pub(crate) mod data; +mod error; +mod frame; +mod recovery; +mod token; +mod wal; + +use std::{ + io::{ + Seek, + SeekFrom, + Write, + }, + path::PathBuf, + sync::{ + Arc, + atomic::{ + AtomicBool, + AtomicUsize, + Ordering, + }, + }, + time::Duration, +}; + +use arc_swap::ArcSwap; +use parking_lot::Mutex; + +use crate::{ + Hash, + Hasher, + LeanIMT, + chunked_level::{ + CHUNK_SIZE, + Chunk, + }, + tree::{ + TreeInner, + TreeSnapshot, + }, +}; + +pub use checkpoint::{ + CheckpointPolicy, + TieringConfig, +}; +pub use config::{ + FlushPolicy, + RotorTreeConfig, +}; +pub use error::{ + RotorTreeError, + StorageError, +}; +pub use token::DurabilityToken; + +use recovery::RecoveryResult; +use token::DurabilityTracker; + +/// level ordered checkpoint during snapshot +struct LevelCheckpointData { + /// unwritten + new_chunks: Vec, + /// start from + from_chunk: usize, + total_chunks: usize, + tail: [Hash; CHUNK_SIZE], +} + +/// checkpoint snapshot +struct CheckpointSnap { + depth: usize, + leaf_count: u64, + last_seq: u64, + root_hash: Hash, + level_data: Vec, +} + +/// Internal mutable state +struct DurableState { + inner: TreeInner, + buffer: Vec, + next_seq: u64, + checkpointed_chunks: Vec, +} + +struct CheckpointCoord { + requested: bool, + completed: u64, +} + +/// Shared state +struct Shared { + hasher: H, + state: Mutex>, + wal_file: Mutex, + snapshot: ArcSwap>, + durability: DurabilityTracker, + closed: AtomicBool, + bg_error: ArcSwap>, + data_dir: PathBuf, + checkpoint_policy: CheckpointPolicy, + tiering: TieringConfig, + entries_since_checkpoint: AtomicUsize, + uncheckpointed_memory_bytes: AtomicUsize, + checkpoint: (Mutex, parking_lot::Condvar), +} + +impl Shared { + fn flush_inner(&self) -> Result<(), StorageError> { + let wal_file = self.wal_file.lock(); + let (buf, last_seq) = { + let mut state = self.state.lock(); + if state.buffer.is_empty() { + return Ok(()); + } + let buf = std::mem::take(&mut state.buffer); + let last_seq = state.next_seq.saturating_sub(1); + (buf, last_seq) + }; + + let pos_before = (&*wal_file).stream_position().map_err(StorageError::Io)?; + + match (&*wal_file).write_all(&buf) { + Ok(()) => { + let sync_result = wal_file.sync_data(); + + let mut state = self.state.lock(); + if state.buffer.is_empty() { + let mut returned = buf; + returned.clear(); + state.buffer = returned; + } + + match sync_result { + Ok(()) => { + self.durability.mark_flushed(last_seq); + Ok(()) + } + Err(e) => { + let err = Arc::new(e); + self.bg_error.store(Arc::new(Some( + error::BackgroundError::FlushFailed(Arc::clone(&err)), + ))); + Err(StorageError::FlushFailed(err)) + } + } + } + Err(e) => { + let _ = (&*wal_file).seek(SeekFrom::Start(pos_before)); + let _ = wal_file.set_len(pos_before); + + let mut state = self.state.lock(); + if state.buffer.is_empty() { + state.buffer = buf; + } else { + let mut combined = buf; + combined.extend_from_slice(&state.buffer); + state.buffer = combined; + } + let err = Arc::new(e); + self.bg_error + .store(Arc::new(Some(error::BackgroundError::FlushFailed( + Arc::clone(&err), + )))); + Err(StorageError::FlushFailed(err)) + } + } + } + + fn check_closed(&self) -> Result<(), StorageError> { + if self.closed.load(Ordering::Acquire) { + return Err(StorageError::Closed); + } + if let Some(ref err) = **self.bg_error.load() { + return Err(match err { + error::BackgroundError::FlushFailed(e) => { + StorageError::FlushFailed(Arc::clone(e)) + } + error::BackgroundError::CheckpointFailed(s) => { + StorageError::CheckpointFailed(s.clone()) + } + }); + } + Ok(()) + } + + fn flush_buffer_locked( + &self, + mut wal_file: &std::fs::File, + state: &mut DurableState, + ) -> Result<(), StorageError> { + if !state.buffer.is_empty() { + wal_file.write_all(&state.buffer)?; + wal_file.sync_data()?; + self.durability + .mark_flushed(state.next_seq.saturating_sub(1)); + state.buffer.clear(); + } + Ok(()) + } + + fn checkpoint_inner(&self) -> Result<(), StorageError> { + let mut wal_file = self.wal_file.lock(); + + let snap = { + let mut state = self.state.lock(); + + self.flush_buffer_locked(&wal_file, &mut state)?; + + let depth = state.inner.depth; + let leaf_count = state.inner.size; + let last_seq = state.next_seq.saturating_sub(1); + + if leaf_count == 0 { + return Ok(()); + } + + let active_levels = depth.min(MAX_DEPTH - 1) + 1; + let mut level_data = Vec::with_capacity(active_levels); + + for level_idx in 0..active_levels { + let total_chunks = state.inner.levels[level_idx].chunk_count(); + let already = if level_idx < state.checkpointed_chunks.len() { + state.checkpointed_chunks[level_idx] + } else { + 0 + }; + + let new_chunks: Vec = + state.inner.levels[level_idx].chunks_since(already); + + let tail = *state.inner.levels[level_idx].tail_data(); + + level_data.push(LevelCheckpointData { + new_chunks, + from_chunk: already, + total_chunks, + tail, + }); + } + + let root_hash = state.inner.root.unwrap_or([0u8; 32]); + + CheckpointSnap { + depth, + leaf_count, + last_seq, + root_hash, + level_data, + } + }; + + std::fs::create_dir_all(&self.data_dir)?; + #[allow(clippy::cast_possible_truncation)] + checkpoint::write_header(&self.data_dir, N as u32, MAX_DEPTH as u32)?; + + let mut files_to_sync = Vec::new(); + + for (level_idx, ld) in snap.level_data.iter().enumerate() { + if !ld.new_chunks.is_empty() { + std::fs::create_dir_all(checkpoint::level_dir_path( + &self.data_dir, + level_idx, + ))?; + let chunks_iter = ld.new_chunks.iter().map(|c| c.as_slice()); + let mut shard_files = checkpoint::append_chunks_to_level( + &self.data_dir, + level_idx, + ld.from_chunk, + chunks_iter, + )?; + files_to_sync.append(&mut shard_files); + } + } + + for file in &files_to_sync { + file.sync_data()?; + } + + let tails: Vec<[Hash; CHUNK_SIZE]> = + snap.level_data.iter().map(|ld| ld.tail).collect(); + checkpoint::write_tails(&self.data_dir, &tails, MAX_DEPTH)?; + + #[allow(clippy::cast_possible_truncation)] + checkpoint::write_meta( + &self.data_dir, + &checkpoint::CheckpointMeta { + n: N as u32, + max_depth: MAX_DEPTH as u32, + last_wal_seq: snap.last_seq, + leaf_count: snap.leaf_count, + depth: snap.depth as u32, + root_hash: snap.root_hash, + }, + )?; + + { + let mut state = self.state.lock(); + + state.checkpointed_chunks.resize(snap.level_data.len(), 0); + + for (level_idx, ld) in snap.level_data.iter().enumerate() { + let snapshot_total = ld.total_chunks; + + if snapshot_total > 0 && level_idx < self.tiering.pin_above_level { + let regions = checkpoint::mmap_level_shards( + &self.data_dir, + level_idx, + snapshot_total, + )?; + if !regions.is_empty() { + state.inner.levels[level_idx] + .remap_chunks(snapshot_total, ®ions); + } + } + + state.checkpointed_chunks[level_idx] = snapshot_total; + } + + // truncate wal + #[allow(clippy::cast_possible_truncation)] + { + let header_buf = wal::serialize_header(N as u32, MAX_DEPTH as u32); + let wal = &mut *wal_file; + wal.seek(SeekFrom::Start(0))?; + wal.write_all(&header_buf)?; + wal.set_len(header_buf.len() as u64)?; + wal.sync_data()?; + } + + self.entries_since_checkpoint.store(0, Ordering::Relaxed); + self.uncheckpointed_memory_bytes.store(0, Ordering::Relaxed); + + self.flush_buffer_locked(&wal_file, &mut state)?; + + let new_snap = state.inner.snapshot(); + self.snapshot.store(Arc::new(new_snap)); + } + + { + let (lock, cvar) = &self.checkpoint; + let mut coord = lock.lock(); + coord.completed += 1; + cvar.notify_all(); + } + + Ok(()) + } + + /// depends on configured checkpoint policy + fn should_auto_checkpoint(&self) -> bool { + match &self.checkpoint_policy { + CheckpointPolicy::Manual | CheckpointPolicy::OnClose => false, + CheckpointPolicy::EveryNEntries(n) => { + self.entries_since_checkpoint.load(Ordering::Relaxed) as u64 >= *n + } + CheckpointPolicy::MemoryThreshold(bytes) => { + self.uncheckpointed_memory_bytes.load(Ordering::Relaxed) >= *bytes + } + } + } + + /// offload checkpoint request + fn request_checkpoint(&self) { + let (lock, cvar) = &self.checkpoint; + let mut coord = lock.lock(); + coord.requested = true; + cvar.notify_all(); + } +} + +/// Handle for the background flush thread +struct FlushHandle { + handle: std::thread::JoinHandle<()>, + shutdown: Arc<(Mutex, parking_lot::Condvar)>, +} + +/// Handle for the background checkpoint thread +struct CheckpointHandle { + handle: std::thread::JoinHandle<()>, +} + +/// Wraps the core tree with persistence via WAL and optional checkpointing. +/// +/// Insertions go to an in-memory WAL buffer, but the state is available for +/// computing on immediately. When `flush()` is called (or triggered by +/// `FlushPolicy`), the buffer is written to disk. Checkpoints materialize +/// the tree state to data files, allowing WAL truncation and mmap-backed reads. +pub struct RotorTree { + shared: Arc>, + flush_handle: Option, + checkpoint_handle: Option, +} + +impl RotorTree { + const _ASSERT_N: () = assert!(N >= 2, "branching factor must be at least 2"); + const _ASSERT_DEPTH: () = assert!(MAX_DEPTH >= 1, "max depth must be at least 1"); + + /// Open or create a RotorTree at the given path. + /// + /// If data files from a previous checkpoint exist, loads state from them + /// and replays only the WAL delta. Otherwise, replays the full WAL. + pub fn open(hasher: H, config: RotorTreeConfig) -> Result { + let () = Self::_ASSERT_N; + let () = Self::_ASSERT_DEPTH; + + std::fs::create_dir_all(&config.path)?; + + let data_dir = config.path.join("data"); + let wal_path = config.path.join("wal"); + let file = std::fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(&wal_path)?; + + use fs4::fs_std::FileExt; + if !file + .try_lock_exclusive() + .map_err(|_| StorageError::FileLocked)? + { + return Err(StorageError::FileLocked.into()); + } + + if file.metadata()?.len() == 0 + && let Some(parent) = wal_path.parent() + { + let dir = std::fs::File::open(parent)?; + dir.sync_all()?; + } + + let mut file = file; + + // recover from checkpoint first + let has_data_dir = data_dir.exists(); + let RecoveryResult { inner, next_seq } = if has_data_dir { + recovery::recover_with_checkpoint::( + &mut file, + &hasher, + &data_dir, + config.verify_checkpoint, + )? + } else { + recovery::recover::(&mut file, &hasher)? + }; + + // Compute initial level file state from checkpoint metadata + let checkpointed_chunks = compute_initial_level_files::(&data_dir)?; + + let snap = inner.snapshot(); + let durability = DurabilityTracker::new(); + + if next_seq > 0 { + durability.mark_flushed(next_seq.checked_sub(1).expect("next_seq > 0; qed")); + } + + let state = DurableState { + inner, + buffer: Vec::new(), + next_seq, + checkpointed_chunks, + }; + + let shared = Arc::new(Shared { + hasher, + state: Mutex::new(state), + wal_file: Mutex::new(file), + snapshot: ArcSwap::from_pointee(snap), + durability, + closed: AtomicBool::new(false), + bg_error: ArcSwap::from_pointee(None), + data_dir, + checkpoint_policy: config.checkpoint_policy, + tiering: config.tiering, + entries_since_checkpoint: AtomicUsize::new(0), + uncheckpointed_memory_bytes: AtomicUsize::new(0), + checkpoint: ( + Mutex::new(CheckpointCoord { + requested: false, + completed: 0, + }), + parking_lot::Condvar::new(), + ), + }); + + let flush_handle = start_flush_thread(&shared, &config.flush_policy)?; + let checkpoint_handle = start_checkpoint_thread(&shared)?; + + Ok(Self { + shared, + flush_handle, + checkpoint_handle, + }) + } + + /// Insert a single leaf + pub fn insert(&self, leaf: Hash) -> Result<(Hash, DurabilityToken), RotorTreeError> { + self.shared.check_closed()?; + + let (root, token) = { + let mut state = self.shared.state.lock(); + let root = LeanIMT::::_insert( + &mut state.inner, + &self.shared.hasher, + leaf, + )?; + let seq = state.next_seq; + wal::serialize_entry(&mut state.buffer, seq, wal::WalPayload::Single(leaf)); + state.next_seq = state.next_seq.checked_add(1).expect("seq overflow; qed"); + self.shared + .entries_since_checkpoint + .fetch_add(1, Ordering::Relaxed); + self.shared + .uncheckpointed_memory_bytes + .fetch_add(32, Ordering::Relaxed); + let snap = state.inner.snapshot(); + self.shared.snapshot.store(Arc::new(snap)); + let token = self.shared.durability.token(seq); + (root, token) + }; + + self.maybe_auto_checkpoint(); + Ok((root, token)) + } + + /// Insert multiple leaves in a batch + pub fn insert_many( + &self, + leaves: &[Hash], + ) -> Result<(Hash, DurabilityToken), RotorTreeError> { + self.shared.check_closed()?; + + let (root, token) = { + let mut state = self.shared.state.lock(); + let root = LeanIMT::::_insert_many( + &mut state.inner, + &self.shared.hasher, + leaves, + )?; + let seq = state.next_seq; + wal::serialize_entry( + &mut state.buffer, + seq, + wal::WalPayload::Batch(wal::NewCow::Borrowed(leaves)), + ); + state.next_seq = state.next_seq.checked_add(1).expect("seq overflow; qed"); + self.shared + .entries_since_checkpoint + .fetch_add(1, Ordering::Relaxed); + self.shared + .uncheckpointed_memory_bytes + .fetch_add(leaves.len() * 32, Ordering::Relaxed); + let snap = state.inner.snapshot(); + self.shared.snapshot.store(Arc::new(snap)); + let token = self.shared.durability.token(seq); + (root, token) + }; + + self.maybe_auto_checkpoint(); + Ok((root, token)) + } + + /// Insert a single leaf and block until it is durable (fsynced) + pub fn insert_durable(&self, leaf: Hash) -> Result { + let (root, token) = self.insert(leaf)?; + self.flush().map_err(RotorTreeError::Storage)?; + token.wait(); + Ok(root) + } + + /// The current Merkle root, or `None` if the tree is empty + pub fn root(&self) -> Option { + self.shared.snapshot.load().root + } + + /// Number of leaves in the tree + pub fn size(&self) -> u64 { + self.shared.snapshot.load().size + } + + /// Current depth (hash layers above the leaf level) + pub fn depth(&self) -> usize { + self.shared.snapshot.load().depth + } + + /// Get a lock-free snapshot for proof generation + pub fn snapshot(&self) -> Arc> { + self.shared.snapshot.load_full() + } + + /// Flush all buffered WAL entries to disk + pub fn flush(&self) -> Result<(), StorageError> { + self.shared.check_closed()?; + self.shared.flush_inner() + } + + /// Write a checkpoint + pub fn checkpoint(&self) -> Result<(), StorageError> { + self.shared.check_closed()?; + self.shared.checkpoint_inner() + } + + /// Block until a background checkpoint completes (or timeout) + pub fn wait_for_checkpoint(&self, timeout: Duration) -> bool { + let (lock, cvar) = &self.shared.checkpoint; + let mut coord = lock.lock(); + let initial = coord.completed; + let deadline = std::time::Instant::now() + timeout; + loop { + if coord.completed > initial { + return true; + } + let remaining = deadline.saturating_duration_since(std::time::Instant::now()); + if remaining.is_zero() { + return false; + } + cvar.wait_for(&mut coord, remaining); + } + } + + /// Close the tree + pub fn close(mut self) -> Result<(), StorageError> { + self.close_inner() + } + + fn close_inner(&mut self) -> Result<(), StorageError> { + if self.shared.closed.swap(true, Ordering::AcqRel) { + return Ok(()); + } + + if self.checkpoint_handle.is_some() { + self.shared.request_checkpoint(); // wake it so it sees `closed` + if let Some(handle) = self.checkpoint_handle.take() { + let _ = handle.handle.join(); + } + } + + if let Some(handle) = self.flush_handle.take() { + { + let (lock, cvar) = &*handle.shutdown; + let mut shutdown = lock.lock(); + *shutdown = true; + cvar.notify_one(); + } + let _ = handle.handle.join(); + } + + self.shared.flush_inner()?; + + // checkpoint on close if configured + if matches!(self.shared.checkpoint_policy, CheckpointPolicy::OnClose) { + self.shared.checkpoint_inner()?; + } + + { + let file = self.shared.wal_file.lock(); + let _ = fs4::fs_std::FileExt::unlock(&*file); + } + + Ok(()) + } + + fn maybe_auto_checkpoint(&self) { + if self.shared.should_auto_checkpoint() { + self.shared.request_checkpoint(); + } + } +} + +fn compute_initial_level_files( + data_dir: &std::path::Path, +) -> Result, StorageError> { + let meta = match checkpoint::read_meta(data_dir) { + Ok(Some(m)) => m, + Ok(None) => return Ok(Vec::new()), + Err(e) => return Err(e), + }; + + let leaf_count = + usize::try_from(meta.leaf_count).map_err(|_| StorageError::MathError)?; + let depth = meta.depth as usize; + + if leaf_count == 0 { + return Ok(Vec::new()); + } + + let active_levels = depth.min(MAX_DEPTH - 1) + 1; + let mut result = Vec::with_capacity(active_levels); + let mut level_len = leaf_count; + + for _ in 0..active_levels { + let chunks = level_len / CHUNK_SIZE; + result.push(chunks); + level_len = level_len.div_ceil(N); + } + + Ok(result) +} + +fn start_flush_thread( + shared: &Arc>, + policy: &FlushPolicy, +) -> Result, RotorTreeError> { + match policy { + FlushPolicy::Manual => Ok(None), + FlushPolicy::Interval(duration) => { + let duration = *duration; + let shared = Arc::clone(shared); + spawn_flush_thread(move |shutdown| { + let (lock, cvar) = &*shutdown; + loop { + let mut stop = lock.lock(); + cvar.wait_for(&mut stop, duration); + if *stop { + break; + } + drop(stop); + let _ = shared.flush_inner(); + } + }) + } + } +} + +fn start_checkpoint_thread( + shared: &Arc>, +) -> Result, RotorTreeError> { + if matches!( + shared.checkpoint_policy, + CheckpointPolicy::Manual | CheckpointPolicy::OnClose + ) { + return Ok(None); + } + + let shared = Arc::clone(shared); + let handle = std::thread::Builder::new() + .name("rotortree-checkpoint".to_string()) + .spawn(move || { + let (lock, cvar) = &shared.checkpoint; + loop { + { + let mut coord = lock.lock(); + while !coord.requested && !shared.closed.load(Ordering::Acquire) { + cvar.wait(&mut coord); + } + if shared.closed.load(Ordering::Acquire) { + break; + } + coord.requested = false; + } + + if let Err(e) = shared.checkpoint_inner() { + shared.bg_error.store(Arc::new(Some( + error::BackgroundError::CheckpointFailed(e.to_string()), + ))); + break; + } + } + }) + .map_err(|e| RotorTreeError::Storage(StorageError::Io(e)))?; + + Ok(Some(CheckpointHandle { handle })) +} + +fn spawn_flush_thread(body: F) -> Result, RotorTreeError> +where + F: FnOnce(Arc<(Mutex, parking_lot::Condvar)>) + Send + 'static, +{ + let shutdown = Arc::new((Mutex::new(false), parking_lot::Condvar::new())); + let shutdown_clone = Arc::clone(&shutdown); + + let handle = std::thread::Builder::new() + .name("rotortree-flush".to_string()) + .spawn(move || { + body(shutdown_clone); + }) + .map_err(|e| RotorTreeError::Storage(StorageError::Io(e)))?; + + Ok(Some(FlushHandle { handle, shutdown })) +} + +impl Drop + for RotorTree +{ + fn drop(&mut self) { + if self.shared.closed.load(Ordering::Acquire) { + return; + } + if std::thread::panicking() { + self.shared.closed.store(true, Ordering::Release); + self.shared.request_checkpoint(); + if let Some(handle) = self.checkpoint_handle.take() { + let _ = handle.handle.join(); + } + if let Some(handle) = self.flush_handle.take() { + let (lock, cvar) = &*handle.shutdown; + let mut shutdown = lock.lock(); + *shutdown = true; + cvar.notify_one(); + drop(shutdown); + let _ = handle.handle.join(); + } + if let Some(file) = self.shared.wal_file.try_lock() { + let _ = fs4::fs_std::FileExt::unlock(&*file); + } + return; + } + if let Err(e) = self.close_inner() { + eprintln!("rotortree: error during drop: {e}"); + } + } +} diff --git a/crates/rotortree/src/storage/recovery.rs b/crates/rotortree/src/storage/recovery.rs new file mode 100644 index 0000000..7110517 --- /dev/null +++ b/crates/rotortree/src/storage/recovery.rs @@ -0,0 +1,558 @@ +use std::{ + io::{ + Read, + Seek, + SeekFrom, + Write, + }, + sync::Arc, +}; + +use crate::chunked_level::{ + CHUNK_SIZE, + Chunk, +}; + +use super::checkpoint; + +use crate::{ + Hash, + Hasher, + tree::TreeInner, +}; + +use super::{ + error::StorageError, + wal::{ + WalPayload, + deserialize_entry, + deserialize_header, + serialize_header, + }, +}; + +/// Result of WAL recovery +pub(crate) struct RecoveryResult { + pub inner: TreeInner, + pub next_seq: u64, +} + +/// Abstraction over file-like objects for recovery +pub(crate) trait WalFile: Read + Write + Seek { + fn file_len(&self) -> Result; + fn truncate_at(&self, len: u64) -> Result<(), std::io::Error>; + fn sync(&self) -> Result<(), std::io::Error>; +} + +impl WalFile for std::fs::File { + fn file_len(&self) -> Result { + Ok(self.metadata()?.len()) + } + + fn truncate_at(&self, len: u64) -> Result<(), std::io::Error> { + self.set_len(len) + } + + fn sync(&self) -> Result<(), std::io::Error> { + self.sync_data() + } +} + +/// Recover tree state from a WAL file. +/// +/// If the file is empty, writes a fresh header. If it contains +/// entries, replays them into a new `TreeInner`. Truncates any +/// incomplete tail entries. +pub(crate) fn recover( + file: &mut F, + hasher: &H, +) -> Result, StorageError> +where + H: Hasher, + F: WalFile, +{ + let file_len = file.file_len()?; + + if file_len == 0 { + #[allow(clippy::cast_possible_truncation)] + let buf = serialize_header(N as u32, MAX_DEPTH as u32); + file.write_all(&buf)?; + file.sync()?; + return Ok(RecoveryResult { + inner: TreeInner::new(), + next_seq: 0, + }); + } + + let file_len_usize = + usize::try_from(file_len).map_err(|_| StorageError::MathError)?; + let mut all_data = vec![0u8; file_len_usize]; + file.seek(SeekFrom::Start(0))?; + file.read_exact(&mut all_data)?; + + let (n, max_depth, header_size) = deserialize_header(&all_data)?; + + #[allow(clippy::cast_possible_truncation)] + if n != N as u32 || max_depth != MAX_DEPTH as u32 { + return Err(StorageError::ConfigMismatch { + expected_n: N as u32, + actual_n: n, + expected_max_depth: MAX_DEPTH as u32, + actual_max_depth: max_depth, + }); + } + + let entry_data = &all_data[header_size..]; + let mut inner = TreeInner::::new(); + let (last_seq, valid_offset) = + replay_wal_entries::(entry_data, hasher, &mut inner, None)?; + + let valid_end = (header_size as u64) + .checked_add(valid_offset as u64) + .ok_or(StorageError::MathError)?; + if valid_end < file_len { + file.truncate_at(valid_end)?; + file.sync()?; + } + file.seek(SeekFrom::Start(valid_end))?; + + let next_seq = match last_seq { + Some(s) => s.checked_add(1).ok_or(StorageError::MathError)?, + None => 0, + }; + Ok(RecoveryResult { inner, next_seq }) +} + +/// Recover from checkpoint data files, falling back to full wal replay +pub(crate) fn recover_with_checkpoint( + wal_file: &mut F, + hasher: &H, + data_dir: &std::path::Path, + verify: bool, +) -> Result, StorageError> +where + H: Hasher, + F: WalFile, +{ + let meta = match checkpoint::read_meta(data_dir)? { + Some(m) => m, + None => return recover(wal_file, hasher), + }; + + #[allow(clippy::cast_possible_truncation)] + if meta.n != N as u32 || meta.max_depth != MAX_DEPTH as u32 { + return recover(wal_file, hasher); + } + + #[allow(clippy::cast_possible_truncation)] + match checkpoint::read_header(data_dir)? { + Some(h) + if h.n == N as u32 + && h.max_depth == MAX_DEPTH as u32 + && h.chunk_size == CHUNK_SIZE as u32 => + { + // valid + } + _ => return recover(wal_file, hasher), + } + + let leaf_count = + usize::try_from(meta.leaf_count).map_err(|_| StorageError::MathError)?; + let depth = meta.depth as usize; + + let mut level_lens = [0usize; MAX_DEPTH]; + if leaf_count > 0 { + level_lens[0] = leaf_count; + for k in 1..=depth.min(MAX_DEPTH - 1) { + level_lens[k] = level_lens[k - 1].div_ceil(N); + } + } + + let tails = match checkpoint::read_tails(data_dir, MAX_DEPTH)? { + Some(t) => t, + None => return recover(wal_file, hasher), + }; + + let mut inner = TreeInner::::new(); + + for level_idx in 0..=depth.min(MAX_DEPTH - 1) { + let len = level_lens[level_idx]; + if len == 0 { + continue; + } + + let num_chunks = len / CHUNK_SIZE; + let tail_len = len % CHUNK_SIZE; + + let regions = if num_chunks > 0 { + checkpoint::mmap_level_shards(data_dir, level_idx, num_chunks)? + } else { + Vec::new() + }; + + let chunks: Vec = if regions.is_empty() { + Vec::new() + } else { + (0..num_chunks) + .map(|chunk_idx| { + let (shard_idx, offset_in_shard) = + checkpoint::shard_address(chunk_idx); + Chunk::new_mapped(Arc::clone(®ions[shard_idx]), offset_in_shard) + }) + .collect() + }; + + inner.set_level_from_parts(level_idx, chunks, tails[level_idx], tail_len, len); + } + + inner.root = if leaf_count > 0 { + Some(meta.root_hash) + } else { + None + }; + inner.size = meta.leaf_count; + inner.depth = depth; + + if verify && leaf_count > 0 { + let computed = inner.recompute_root(hasher); + if computed != inner.root { + return Err(StorageError::DataCorruption { + detail: format!( + "root mismatch: stored {:?}, recomputed {:?}", + inner.root, computed + ), + }); + } + } + + let file_len = wal_file.file_len()?; + + let next_seq = if file_len == 0 { + #[allow(clippy::cast_possible_truncation)] + let buf = serialize_header(N as u32, MAX_DEPTH as u32); + wal_file.write_all(&buf)?; + wal_file.sync()?; + meta.last_wal_seq + .checked_add(1) + .ok_or(StorageError::MathError)? + } else { + let file_len_usize = + usize::try_from(file_len).map_err(|_| StorageError::MathError)?; + let mut all_data = vec![0u8; file_len_usize]; + wal_file.seek(SeekFrom::Start(0))?; + wal_file.read_exact(&mut all_data)?; + + let (n, max_depth, header_size) = deserialize_header(&all_data)?; + #[allow(clippy::cast_possible_truncation)] + if n != N as u32 || max_depth != MAX_DEPTH as u32 { + return Err(StorageError::ConfigMismatch { + expected_n: N as u32, + actual_n: n, + expected_max_depth: MAX_DEPTH as u32, + actual_max_depth: max_depth, + }); + } + + let entry_data = &all_data[header_size..]; + let (last_seq, valid_offset) = replay_wal_entries::( + entry_data, + hasher, + &mut inner, + Some(meta.last_wal_seq), + )?; + + let valid_end = (header_size as u64) + .checked_add(valid_offset as u64) + .ok_or(StorageError::MathError)?; + if valid_end < file_len { + wal_file.truncate_at(valid_end)?; + wal_file.sync()?; + } + wal_file.seek(SeekFrom::Start(valid_end))?; + + match last_seq { + Some(s) => s.checked_add(1).ok_or(StorageError::MathError)?, + None => meta + .last_wal_seq + .checked_add(1) + .ok_or(StorageError::MathError)?, + } + }; + + Ok(RecoveryResult { inner, next_seq }) +} + +/// replay wal +fn replay_wal_entries( + entry_data: &[u8], + hasher: &H, + inner: &mut TreeInner, + skip_until_seq: Option, +) -> Result<(Option, usize), StorageError> { + let mut offset: usize = 0; + let mut last_seq: Option = None; + let mut pending_singles: Vec = Vec::new(); + + loop { + match deserialize_entry(entry_data, offset, last_seq) { + Ok(Some((entry, consumed))) => { + let should_replay = match skip_until_seq { + Some(skip) => entry.seq > skip, + None => true, + }; + if should_replay { + match entry.payload { + WalPayload::Single(leaf) => { + pending_singles.push(leaf); + } + WalPayload::Batch(cow) => { + flush_pending::( + inner, + hasher, + &mut pending_singles, + )?; + crate::LeanIMT::::_insert_many( + inner, + hasher, + cow.as_slice(), + ) + .map_err(StorageError::Tree)?; + } + } + } + last_seq = Some(entry.seq); + offset = offset + .checked_add(consumed) + .ok_or(StorageError::MathError)?; + } + Ok(None) => break, + Err(e) => return Err(e), + } + } + + flush_pending::(inner, hasher, &mut pending_singles)?; + Ok((last_seq, offset)) +} + +/// Flush accumulated single-insert leaves as a batch +fn flush_pending( + inner: &mut TreeInner, + hasher: &H, + pending: &mut Vec, +) -> Result<(), StorageError> { + if pending.is_empty() { + return Ok(()); + } + crate::LeanIMT::::_insert_many(inner, hasher, pending) + .map_err(StorageError::Tree)?; + pending.clear(); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::{ + super::wal, + *, + }; + use crate::Hash; + use std::{ + cell::RefCell, + io::Cursor, + }; + + /// In-memory file for testing recovery without disk I/O + struct MemFile { + inner: RefCell>>, + } + + impl MemFile { + fn new() -> Self { + Self { + inner: RefCell::new(Cursor::new(Vec::new())), + } + } + + fn from_bytes(bytes: Vec) -> Self { + Self { + inner: RefCell::new(Cursor::new(bytes)), + } + } + + fn data(&self) -> Vec { + self.inner.borrow().get_ref().clone() + } + } + + impl Read for MemFile { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + self.inner.borrow_mut().read(buf) + } + } + + impl Write for MemFile { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.inner.borrow_mut().write(buf) + } + fn flush(&mut self) -> std::io::Result<()> { + self.inner.borrow_mut().flush() + } + } + + impl Seek for MemFile { + fn seek(&mut self, pos: SeekFrom) -> std::io::Result { + self.inner.borrow_mut().seek(pos) + } + } + + impl WalFile for MemFile { + fn file_len(&self) -> Result { + Ok(self.inner.borrow().get_ref().len() as u64) + } + + fn truncate_at(&self, len: u64) -> Result<(), std::io::Error> { + self.inner.borrow_mut().get_mut().truncate(len as usize); + Ok(()) + } + + fn sync(&self) -> Result<(), std::io::Error> { + Ok(()) + } + } + + use crate::test_util::*; + + #[test] + fn recover_empty_wal() { + let mut file = MemFile::new(); + let result = recover::(&mut file, &XorHasher).unwrap(); + assert_eq!(result.inner.size, 0); + assert_eq!(result.inner.root, None); + assert_eq!(result.next_seq, 0); + + let header_size = wal::serialize_header(2, 32).len(); + assert_eq!(file.data().len(), header_size); + } + + #[test] + fn recover_single_entries() { + let mut buf = wal::serialize_header(2, 32); + for i in 0..5u64 { + wal::serialize_entry(&mut buf, i, wal::WalPayload::Single(leaf(i as u32))); + } + + let mut file = MemFile::from_bytes(buf); + let result = recover::(&mut file, &XorHasher).unwrap(); + assert_eq!(result.inner.size, 5); + assert_eq!(result.next_seq, 5); + assert!(result.inner.root.is_some()); + } + + #[test] + fn recover_batch_entry() { + let leaves: Vec = (0..10u32).map(leaf).collect(); + let mut buf = wal::serialize_header(2, 32); + wal::serialize_entry( + &mut buf, + 0, + wal::WalPayload::Batch(wal::NewCow::Borrowed(&leaves)), + ); + + let mut file = MemFile::from_bytes(buf); + let result = recover::(&mut file, &XorHasher).unwrap(); + assert_eq!(result.inner.size, 10); + assert_eq!(result.next_seq, 1); + } + + #[test] + fn recover_truncated_tail() { + let mut buf = wal::serialize_header(2, 32); + wal::serialize_entry(&mut buf, 0, wal::WalPayload::Single(leaf(1))); + wal::serialize_entry(&mut buf, 1, wal::WalPayload::Single(leaf(2))); + buf.truncate(buf.len() - 10); + + let mut file = MemFile::from_bytes(buf); + let result = recover::(&mut file, &XorHasher).unwrap(); + assert_eq!(result.inner.size, 1); + assert_eq!(result.next_seq, 1); + } + + #[test] + fn recover_config_mismatch() { + let buf = wal::serialize_header(4, 32); + + let mut file = MemFile::from_bytes(buf); + let result = recover::(&mut file, &XorHasher); + assert!(matches!(result, Err(StorageError::ConfigMismatch { .. }))); + } + + #[test] + fn recover_matches_sequential_inserts() { + let mut buf = wal::serialize_header(2, 32); + for i in 0..20u64 { + wal::serialize_entry(&mut buf, i, wal::WalPayload::Single(leaf(i as u32))); + } + + let mut file = MemFile::from_bytes(buf); + let recovered = recover::(&mut file, &XorHasher).unwrap(); + + let mut inner = TreeInner::<2, 32>::new(); + for i in 0..20u32 { + crate::LeanIMT::::_insert(&mut inner, &XorHasher, leaf(i)) + .unwrap(); + } + + assert_eq!(recovered.inner.root, inner.root); + assert_eq!(recovered.inner.size, inner.size); + assert_eq!(recovered.inner.depth, inner.depth); + } + + #[test] + fn recover_header_only_wal() { + let buf = wal::serialize_header(2, 32); + + let mut file = MemFile::from_bytes(buf); + let result = recover::(&mut file, &XorHasher).unwrap(); + assert_eq!(result.inner.size, 0); + assert_eq!(result.inner.root, None); + assert_eq!(result.next_seq, 0); + } + + #[test] + fn recover_mixed_single_and_batch() { + let mut buf = wal::serialize_header(2, 32); + + for i in 0..3u64 { + wal::serialize_entry(&mut buf, i, wal::WalPayload::Single(leaf(i as u32))); + } + let batch: Vec = (3..8u32).map(leaf).collect(); + wal::serialize_entry( + &mut buf, + 3, + wal::WalPayload::Batch(wal::NewCow::Borrowed(&batch)), + ); + for i in 4..6u64 { + wal::serialize_entry( + &mut buf, + i, + wal::WalPayload::Single(leaf((i + 4) as u32)), + ); + } + + let mut file = MemFile::from_bytes(buf); + let recovered = recover::(&mut file, &XorHasher).unwrap(); + assert_eq!(recovered.inner.size, 10); + assert_eq!(recovered.next_seq, 6); + + // Verify against reference + let mut inner = TreeInner::<2, 32>::new(); + let all_leaves: Vec = (0..10u32).map(leaf).collect(); + crate::LeanIMT::::_insert_many( + &mut inner, + &XorHasher, + &all_leaves, + ) + .unwrap(); + assert_eq!(recovered.inner.root, inner.root); + } +} diff --git a/crates/rotortree/src/storage/token.rs b/crates/rotortree/src/storage/token.rs new file mode 100644 index 0000000..d8c99c8 --- /dev/null +++ b/crates/rotortree/src/storage/token.rs @@ -0,0 +1,63 @@ +use std::sync::Arc; + +use parking_lot::{ + Condvar, + Mutex, +}; + +/// Shared state for tracking the last durably flushed WAL sequence number +pub(crate) struct DurabilityTracker { + inner: Arc<(Mutex>, Condvar)>, +} + +impl DurabilityTracker { + pub(crate) fn new() -> Self { + Self { + inner: Arc::new((Mutex::new(None), Condvar::new())), + } + } + + /// Create a token that can wait for the given sequence number to be flushed + pub(crate) fn token(&self, seq: u64) -> DurabilityToken { + DurabilityToken { + seq, + inner: Arc::clone(&self.inner), + } + } + + /// Mark all entries up to `seq` as durably flushed and wake waiters + pub(crate) fn mark_flushed(&self, seq: u64) { + let (lock, cvar) = &*self.inner; + let mut flushed = lock.lock(); + match *flushed { + Some(prev) if seq <= prev => return, + _ => *flushed = Some(seq), + } + cvar.notify_all(); + } +} + +/// A token returned from `RotorTree::insert` that tracks whether the +/// corresponding WAL entry has been flushed to disk +#[derive(Debug)] +pub struct DurabilityToken { + seq: u64, + inner: Arc<(Mutex>, Condvar)>, +} + +impl DurabilityToken { + /// Block until the WAL entry for this insert has been fsynced + pub fn wait(&self) { + let (lock, cvar) = &*self.inner; + let mut flushed = lock.lock(); + while !matches!(*flushed, Some(f) if f >= self.seq) { + cvar.wait(&mut flushed); + } + } + + /// Non-blocking check if entry was fsynced + pub fn is_durable(&self) -> bool { + let (lock, _) = &*self.inner; + matches!(*lock.lock(), Some(f) if f >= self.seq) + } +} diff --git a/crates/rotortree/src/storage/wal.rs b/crates/rotortree/src/storage/wal.rs new file mode 100644 index 0000000..09bb330 --- /dev/null +++ b/crates/rotortree/src/storage/wal.rs @@ -0,0 +1,231 @@ +use crate::Hash; + +use super::{ + error::StorageError, + frame, +}; + +pub(crate) const FILE_MAGIC: [u8; 4] = [0x52, 0x4F, 0x54, 0x52]; + +#[derive(Debug, wincode::SchemaWrite, wincode::SchemaRead)] +pub(crate) enum WalHeader { + V1 { + magic: [u8; 4], + n: u32, + max_depth: u32, + }, +} + +/// new cow moo +#[derive(Debug, wincode::SchemaWrite, wincode::SchemaRead)] +pub(crate) enum NewCow<'a> { + Owned(Vec), + Borrowed(&'a [Hash]), +} + +impl NewCow<'_> { + pub fn as_slice(&self) -> &[Hash] { + match self { + NewCow::Owned(v) => v, + NewCow::Borrowed(s) => s, + } + } +} + +impl PartialEq for NewCow<'_> { + fn eq(&self, other: &Self) -> bool { + self.as_slice() == other.as_slice() + } +} + +impl Eq for NewCow<'_> {} + +#[derive(Debug, PartialEq, Eq, wincode::SchemaWrite, wincode::SchemaRead)] +pub(crate) enum WalEntry<'a> { + V1(WalEntryV1<'a>), +} + +#[derive(Debug, PartialEq, Eq, wincode::SchemaWrite, wincode::SchemaRead)] +pub(crate) struct WalEntryV1<'a> { + pub seq: u64, + pub payload: WalPayload<'a>, +} + +#[derive(Debug, PartialEq, Eq, wincode::SchemaWrite, wincode::SchemaRead)] +pub(crate) enum WalPayload<'a> { + Single(Hash), + Batch(NewCow<'a>), +} + +pub(crate) fn serialize_header(n: u32, max_depth: u32) -> Vec { + let header = WalHeader::V1 { + magic: FILE_MAGIC, + n, + max_depth, + }; + frame::serialize_frame(&header) +} + +/// Returns (n, max_depth, header_frame_size). +pub(crate) fn deserialize_header(data: &[u8]) -> Result<(u32, u32, usize), StorageError> { + let (payload, frame_size) = + frame::read_frame(data, 0)?.ok_or(StorageError::WalCorrupted { offset: 0 })?; + let header: WalHeader = wincode::deserialize(payload) + .map_err(|_| StorageError::WalCorrupted { offset: 0 })?; + let WalHeader::V1 { + magic, + n, + max_depth, + } = header; + if magic != FILE_MAGIC { + return Err(StorageError::WalCorrupted { offset: 0 }); + } + Ok((n, max_depth, frame_size)) +} + +pub(crate) fn serialize_entry(buf: &mut Vec, seq: u64, payload: WalPayload<'_>) { + let entry = WalEntry::V1(WalEntryV1 { seq, payload }); + frame::write_frame(buf, &entry); +} + +pub(crate) fn deserialize_entry<'a>( + data: &'a [u8], + offset: usize, + last_seq: Option, +) -> Result, usize)>, StorageError> { + let (payload, frame_size) = match frame::read_frame(data, offset)? { + Some(r) => r, + None => return Ok(None), + }; + + let entry: WalEntry<'a> = match wincode::deserialize(payload) { + Ok(e) => e, + Err(_) => { + return Err(StorageError::WalCorrupted { + offset: offset as u64, + }); + } + }; + + let WalEntry::V1(v1) = entry; + + if let Some(last) = last_seq + && v1.seq != last.checked_add(1).ok_or(StorageError::MathError)? + { + return Ok(None); + } + + Ok(Some((v1, frame_size))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn header_round_trip() { + let buf = serialize_header(4, 32); + let (n, max_depth, consumed) = deserialize_header(&buf).unwrap(); + assert_eq!(n, 4); + assert_eq!(max_depth, 32); + assert_eq!(consumed, buf.len()); + } + + #[test] + fn header_bad_magic() { + let mut buf = serialize_header(2, 32); + buf[4] ^= 0xFF; + assert!(deserialize_header(&buf).is_err()); + } + + #[test] + fn header_bad_crc() { + let mut buf = serialize_header(2, 32); + let last = buf.len() - 1; + buf[last] ^= 0xFF; + assert!(deserialize_header(&buf).is_err()); + } + + #[test] + fn single_entry_round_trip() { + let leaf = [42u8; 32]; + let mut buf = Vec::new(); + serialize_entry(&mut buf, 0, WalPayload::Single(leaf)); + + let (entry, consumed) = deserialize_entry(&buf, 0, None).unwrap().unwrap(); + assert_eq!(consumed, buf.len()); + assert_eq!(entry.seq, 0); + assert_eq!(entry.payload, WalPayload::Single(leaf)); + } + + #[test] + fn batch_entry_round_trip() { + let leaves: Vec = (0..10u8) + .map(|i| { + let mut h = [0u8; 32]; + h[0] = i; + h + }) + .collect(); + let mut buf = Vec::new(); + serialize_entry(&mut buf, 5, WalPayload::Batch(NewCow::Borrowed(&leaves))); + + let (entry, consumed) = deserialize_entry(&buf, 0, Some(4)).unwrap().unwrap(); + assert_eq!(consumed, buf.len()); + assert_eq!(entry.seq, 5); + assert_eq!(entry.payload, WalPayload::Batch(NewCow::Borrowed(&leaves))); + } + + #[test] + fn crc_corruption_mid_file() { + let leaf = [1u8; 32]; + let mut buf = Vec::new(); + serialize_entry(&mut buf, 0, WalPayload::Single(leaf)); + buf[6] ^= 0xFF; + // Append extra bytes so this is not at the tail. + buf.extend_from_slice(&[0u8; 100]); + let result = deserialize_entry(&buf, 0, None); + assert!(matches!(result, Err(StorageError::CrcMismatch { .. }))); + } + + #[test] + fn truncated_tail_returns_none() { + let mut buf = Vec::new(); + serialize_entry(&mut buf, 0, WalPayload::Single([1u8; 32])); + buf.pop(); + assert!(deserialize_entry(&buf, 0, None).unwrap().is_none()); + } + + #[test] + fn seq_gap_returns_none() { + let mut buf = Vec::new(); + serialize_entry(&mut buf, 5, WalPayload::Single([1u8; 32])); + assert!(deserialize_entry(&buf, 0, Some(3)).unwrap().is_none()); + } + + #[test] + fn multiple_entries_sequential() { + let mut buf = Vec::new(); + for i in 0..5u64 { + let mut leaf = [0u8; 32]; + leaf[0] = i as u8; + serialize_entry(&mut buf, i, WalPayload::Single(leaf)); + } + + let mut offset = 0; + let mut last_seq = None; + let mut count = 0u64; + + while let Some((entry, consumed)) = + deserialize_entry(&buf, offset, last_seq).unwrap() + { + assert_eq!(entry.seq, count); + last_seq = Some(entry.seq); + offset += consumed; + count += 1; + } + + assert_eq!(count, 5); + assert_eq!(offset, buf.len()); + } +} diff --git a/crates/rotortree/src/test_util.rs b/crates/rotortree/src/test_util.rs new file mode 100644 index 0000000..2196035 --- /dev/null +++ b/crates/rotortree/src/test_util.rs @@ -0,0 +1,43 @@ +use crate::{ + Hash, + HashState, + Hasher, +}; + +pub struct XorState { + hash: [u8; 32], + pos: usize, +} + +impl HashState for XorState { + fn update(&mut self, data: &[u8]) { + for &b in data { + self.hash[self.pos % 32] ^= b; + self.pos += 1; + } + } + + fn finalize(self) -> Hash { + self.hash + } +} + +#[derive(Clone)] +pub struct XorHasher; + +impl Hasher for XorHasher { + type State = XorState; + + fn new_state(&self) -> Self::State { + XorState { + hash: [0u8; 32], + pos: 0, + } + } +} + +pub fn leaf(n: u32) -> Hash { + let mut h = [0u8; 32]; + h[..4].copy_from_slice(&n.to_le_bytes()); + h +} diff --git a/crates/rotortree/src/tree.rs b/crates/rotortree/src/tree.rs new file mode 100644 index 0000000..af2f408 --- /dev/null +++ b/crates/rotortree/src/tree.rs @@ -0,0 +1,1046 @@ +use crate::{ + Hash, + Hasher, + TreeError, + chunked_level::ChunkedLevel, +}; + +/// Number of parents per rayon task +#[cfg(feature = "parallel")] +const PAR_CHUNK_SIZE: usize = 64; + +#[cfg(feature = "parallel")] +pub(crate) fn parallel_threshold() -> usize { + static THRESHOLD: std::sync::OnceLock = std::sync::OnceLock::new(); + *THRESHOLD.get_or_init(|| { + std::env::var("ROTORTREE_PARALLEL_THRESHOLD") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(1024) + }) +} + +/// Returns the number of hash layers above the leaf level. +/// +/// - `size <= 1` → 0 (root IS the leaf or tree is empty) +/// - `size == N` → 1 +/// - `size == N^k` → k +#[allow(clippy::arithmetic_side_effects)] +#[inline(always)] +pub(crate) fn ceil_log_n(size: u64, n: usize) -> usize { + if size <= 1 { + return 0; + } + if n.is_power_of_two() { + let k = n.trailing_zeros(); + let bits = u64::BITS - (size - 1).leading_zeros(); + return bits.div_ceil(k) as usize; + } + (size - 1).ilog(n as u64) as usize + 1 +} + +/// Convert `u64` to `usize`, returning `CapacityExceeded` on +/// failure (relevant on 32-bit platforms). +#[inline] +pub(crate) fn u64_to_usize(val: u64) -> Result { + usize::try_from(val).map_err(|_| TreeError::CapacityExceeded) +} + +/// Immutable snapshot of the tree for lock-free reads and proof +/// generation. +pub struct TreeSnapshot { + pub(crate) levels: [ChunkedLevel; MAX_DEPTH], + pub(crate) root: Option, + pub(crate) size: u64, + pub(crate) depth: usize, +} + +impl TreeSnapshot { + /// The Merkle root, or `None` if the tree is empty. + pub fn root(&self) -> Option { + self.root + } + + /// Number of leaves in the snapshot. + pub fn size(&self) -> u64 { + self.size + } + + /// Current depth (hash layers above the leaf level). + pub fn depth(&self) -> usize { + self.depth + } + + /// Number of nodes at the given level. + pub fn level_len(&self, level: usize) -> usize { + if level > self.depth { + return 0; + } + self.levels[level].len() + } + + /// Retrieve the hash of a specific node by level and index. + pub fn get_node(&self, level: usize, index: usize) -> Result { + if level > self.depth { + return Err(TreeError::IndexOutOfRange { + index: index as u64, + size: 0, + }); + } + self.levels[level].get(index) + } +} + +/// Mutable tree state. +pub(crate) struct TreeInner { + /// Levels 0..depth-1 + pub(crate) levels: [ChunkedLevel; MAX_DEPTH], + /// The root hash + pub(crate) root: Option, + pub(crate) size: u64, + pub(crate) depth: usize, +} + +impl TreeInner { + pub(crate) fn new() -> Self { + Self { + levels: core::array::from_fn(|_| ChunkedLevel::new()), + root: None, + size: 0, + depth: 0, + } + } + + /// Replace a level's contents from checkpoint data. + #[cfg(feature = "storage")] + pub(crate) fn set_level_from_parts( + &mut self, + level_idx: usize, + chunks: Vec, + tail: [Hash; crate::chunked_level::CHUNK_SIZE], + tail_len: usize, + len: usize, + ) { + self.levels[level_idx] = ChunkedLevel::from_parts(chunks, tail, tail_len, len); + } + + /// Recompute the root hash from level 0 data bottom-up + #[cfg(feature = "storage")] + pub(crate) fn recompute_root(&self, hasher: &H) -> Option { + use std::vec::Vec; + if self.size == 0 { + return None; + } + if self.size == 1 { + return self.levels[0].get(0).ok(); + } + + let depth = self.depth; + let level0_len = self.levels[0].len(); + let mut current: Vec = Vec::with_capacity(level0_len); + for i in 0..level0_len { + current.push(self.levels[0].get(i).ok()?); + } + + for _level in 0..depth { + let len = current.len(); + let num_parents = len.div_ceil(N); + + #[cfg(feature = "parallel")] + let parents = { + use rayon::prelude::*; + if num_parents >= parallel_threshold() { + let mut buf = vec![[0u8; 32]; num_parents]; + buf.par_chunks_mut(PAR_CHUNK_SIZE).enumerate().for_each( + |(ci, chunk)| { + let base = ci * PAR_CHUNK_SIZE; + for (i, slot) in chunk.iter_mut().enumerate() { + *slot = + Self::_hash_group(¤t, base + i, len, hasher); + } + }, + ); + buf + } else { + (0..num_parents) + .map(|parent_idx| { + Self::_hash_group(¤t, parent_idx, len, hasher) + }) + .collect() + } + }; + + #[cfg(not(feature = "parallel"))] + let parents: Vec = (0..num_parents) + .map(|parent_idx| Self::_hash_group(¤t, parent_idx, len, hasher)) + .collect(); + + current = parents; + } + + current.first().copied() + } + + #[cfg(feature = "storage")] + #[inline] + fn _hash_group( + current: &[Hash], + parent_idx: usize, + len: usize, + hasher: &H, + ) -> Hash { + let start = parent_idx * N; + let end = core::cmp::min(start + N, len); + let count = end - start; + if count == 1 { + current[start] + } else { + hasher.hash_children(¤t[start..end]) + } + } + + pub(crate) fn snapshot(&self) -> TreeSnapshot { + let mut levels = core::array::from_fn(|_| ChunkedLevel::new()); + let snap_count = core::cmp::min(self.depth.saturating_add(1), MAX_DEPTH); + for (dst, src) in levels.iter_mut().zip(self.levels.iter()).take(snap_count) { + *dst = src.clone(); + } + TreeSnapshot { + levels, + root: self.root, + size: self.size, + depth: self.depth, + } + } +} + +/// An N-ary Lean Incremental Merkle Tree. +/// +/// # Type Parameters +/// +/// - `H`: Hash function ([`Hasher`]) +/// - `N`: Branching factor (compile-time, must be >= 2) +/// - `MAX_DEPTH`: Maximum tree depth (must be >= 1) +pub struct LeanIMT { + hasher: H, + #[cfg(not(feature = "concurrent"))] + #[cfg_attr(docsrs, doc(cfg(not(feature = "concurrent"))))] + inner: TreeInner, + #[cfg(feature = "concurrent")] + #[cfg_attr(docsrs, doc(cfg(feature = "concurrent")))] + inner: parking_lot::RwLock>, +} + +impl LeanIMT { + const _ASSERT_N: () = assert!(N >= 2, "branching factor must be at least 2"); + const _ASSERT_DEPTH: () = assert!(MAX_DEPTH >= 1, "max depth must be at least 1"); + + /// Create a new empty tree. + #[cfg(not(feature = "concurrent"))] + pub fn new(hasher: H) -> Self { + let () = Self::_ASSERT_N; + let () = Self::_ASSERT_DEPTH; + Self { + hasher, + inner: TreeInner::new(), + } + } + + /// Create a new empty tree + #[cfg(feature = "concurrent")] + #[cfg_attr(docsrs, doc(cfg(feature = "concurrent")))] + pub fn new(hasher: H) -> Self { + let () = Self::_ASSERT_N; + let () = Self::_ASSERT_DEPTH; + Self { + hasher, + inner: parking_lot::RwLock::new(TreeInner::new()), + } + } + + /// Insert a single leaf. Returns the new Merkle root. + #[cfg(not(feature = "concurrent"))] + pub fn insert(&mut self, leaf: Hash) -> Result { + Self::_insert(&mut self.inner, &self.hasher, leaf) + } + + /// Insert a single leaf. Returns the new Merkle root. + #[cfg(feature = "concurrent")] + #[cfg_attr(docsrs, doc(cfg(feature = "concurrent")))] + pub fn insert(&self, leaf: Hash) -> Result { + Self::_insert(&mut self.inner.write(), &self.hasher, leaf) + } + + /// Insert multiple leaves in a batch. Returns the new root. + #[cfg(not(feature = "concurrent"))] + pub fn insert_many(&mut self, leaves: &[Hash]) -> Result { + Self::_insert_many(&mut self.inner, &self.hasher, leaves) + } + + /// Insert multiple leaves in a batch. Returns the new root. + #[cfg(feature = "concurrent")] + #[cfg_attr(docsrs, doc(cfg(feature = "concurrent")))] + pub fn insert_many(&self, leaves: &[Hash]) -> Result { + Self::_insert_many(&mut self.inner.write(), &self.hasher, leaves) + } + + /// The current Merkle root, or `None` if the tree is empty. + #[cfg(not(feature = "concurrent"))] + pub fn root(&self) -> Option { + self.inner.root + } + + /// The current Merkle root, or `None` if the tree is empty. + #[cfg(feature = "concurrent")] + #[cfg_attr(docsrs, doc(cfg(feature = "concurrent")))] + pub fn root(&self) -> Option { + self.inner.read().root + } + + /// Number of leaves in the tree. + #[cfg(not(feature = "concurrent"))] + pub fn size(&self) -> u64 { + self.inner.size + } + + /// Number of leaves in the tree. + #[cfg(feature = "concurrent")] + #[cfg_attr(docsrs, doc(cfg(feature = "concurrent")))] + pub fn size(&self) -> u64 { + self.inner.read().size + } + + /// Current depth (hash layers above the leaf level). + #[cfg(not(feature = "concurrent"))] + pub fn depth(&self) -> usize { + self.inner.depth + } + + /// Current depth (hash layers above the leaf level). + #[cfg(feature = "concurrent")] + #[cfg_attr(docsrs, doc(cfg(feature = "concurrent")))] + pub fn depth(&self) -> usize { + self.inner.read().depth + } + + /// Create an immutable snapshot for proof generation. + #[cfg(not(feature = "concurrent"))] + pub fn snapshot(&self) -> TreeSnapshot { + self.inner.snapshot() + } + + /// Create an immutable snapshot for proof generation. + #[cfg(feature = "concurrent")] + #[cfg_attr(docsrs, doc(cfg(feature = "concurrent")))] + pub fn snapshot(&self) -> TreeSnapshot { + self.inner.read().snapshot() + } + + #[inline] + pub(crate) fn _insert( + inner: &mut TreeInner, + hasher: &H, + leaf: Hash, + ) -> Result { + let new_size = inner + .size + .checked_add(1) + .ok_or(TreeError::CapacityExceeded)?; + let depth = ceil_log_n(new_size, N); + if depth > MAX_DEPTH { + return Err(TreeError::MaxDepthExceeded { + max_depth: MAX_DEPTH, + }); + } + let index = u64_to_usize(inner.size)?; + + let mut node = leaf; + let mut idx = index; + for level in 0..depth { + inner.levels[level].set(idx, node)?; + + let child_pos = idx % N; + if child_pos != 0 { + let group_start = idx - child_pos; + let count = child_pos + 1; + let mut children = [[0u8; 32]; N]; + if child_pos > 0 { + inner.levels[level].get_group(group_start, child_pos, &mut children); + } + children[child_pos] = node; + node = hasher.hash_children(&children[..count]); + } + idx /= N; + } + + if depth < MAX_DEPTH { + inner.levels[depth].set(0, node)?; + } + inner.root = Some(node); + inner.size = new_size; + inner.depth = depth; + Ok(node) + } + + /// Compute the parent hash for a group at `parent_idx` + /// within a single level. + #[inline(always)] + fn _compute_parent( + child_level: &ChunkedLevel, + parent_idx: usize, + level_len: usize, + hasher: &H, + ) -> Result { + let group_start = parent_idx * N; + let group_end = core::cmp::min(group_start + N, level_len); + let count = group_end - group_start; + if count == 1 { + child_level.get(group_start) + } else { + let mut children = [[0u8; 32]; N]; + child_level.get_group(group_start, count, &mut children); + Ok(hasher.hash_children(&children[..count])) + } + } + + /// Sequential inner loop for one level of `_insert_many`. + /// + /// Consecutive FULL groups (each exactly `N` children) at an eligible + /// arity are gathered into windows and hashed together via + /// [`Hasher::hash_many_into`], which lets a SIMD-capable hasher hash + /// several parents at once. The trailing partial group and the count==1 + /// lift fall back to the scalar [`Self::_compute_parent`] path. For + /// hashers without a batch override this is identical to the old per-group + /// loop (the default `hash_many_into` is a scalar loop). + #[inline] + #[allow(clippy::too_many_arguments)] + fn _insert_many_level_seq( + levels: &mut [ChunkedLevel], + level: usize, + start_parent: usize, + num_parents: usize, + level_len: usize, + is_root_level: bool, + hasher: &H, + root: &mut Hash, + ) -> Result<(), TreeError> { + let next_level = level + 1; + + if next_level < levels.len() { + // Split so the child level is borrowed immutably while the parent + // level is written mutably. + let (head, tail) = levels.split_at_mut(next_level); + let child = &head[level]; + let parent = &mut tail[0]; + + // Leading full groups (count == N) batch through hash_many_into. + // The final group is full iff `level_len` is a multiple of N. + let full_parents = (level_len / N).max(start_parent); + Self::_batch_full_groups(child, parent, start_parent, full_parents, hasher); + + // Trailing partial group / lift (empty range when level_len % N == 0). + for parent_idx in full_parents..num_parents { + let p = Self::_compute_parent(child, parent_idx, level_len, hasher)?; + parent.set_preallocated(parent_idx, p); + } + + // At the root level there is exactly one parent; it is the root. + if is_root_level { + *root = parent.get(num_parents - 1)?; + } + } else { + // No parent level to write into: compute scalar, track root only. + for parent_idx in start_parent..num_parents { + let p = + Self::_compute_parent(&levels[level], parent_idx, level_len, hasher)?; + if is_root_level { + *root = p; + } + } + } + Ok(()) + } + + /// Hash the full groups `start_parent..full_parents` of `child` into + /// `parent`, batching eligible runs through [`Hasher::hash_many_into`]. + #[inline] + fn _batch_full_groups( + child: &ChunkedLevel, + parent: &mut ChunkedLevel, + start_parent: usize, + full_parents: usize, + hasher: &H, + ) { + /// Parents gathered per `hash_many_into` call. The Blake3 override + /// re-splits this into `simd_degree`-sized SIMD calls internally. + const WINDOW: usize = 16; + + let mut parent_idx = start_parent; + let mut refs: [&[Hash]; WINDOW] = [&[]; WINDOW]; + let mut out: [Hash; WINDOW] = [[0u8; 32]; WINDOW]; + while parent_idx < full_parents { + let take = (full_parents - parent_idx).min(WINDOW); + let mut filled = 0; + for slot in refs.iter_mut().take(take) { + let start = (parent_idx + filled) * N; + match child.group_slice(start, N) { + Some(g) => *slot = g, + // Group straddles a boundary (cannot happen for eligible N + // since CHUNK_SIZE % N == 0, but stay correct anyway): + // stop the window here and let the scalar tail finish it. + None => break, + } + filled += 1; + } + if filled == 0 { + // Could not borrow even one group contiguously; scalar. + let p = hasher.hash_children(&Self::_copy_group(child, parent_idx)); + parent.set_preallocated(parent_idx, p); + parent_idx += 1; + continue; + } + hasher.hash_many_into(&refs[..filled], &mut out[..filled]); + for (i, &h) in out[..filled].iter().enumerate() { + parent.set_preallocated(parent_idx + i, h); + } + parent_idx += filled; + } + } + + /// Copy a full N-child group into a stack buffer (boundary-straddling + /// fallback for `_batch_full_groups`). + #[inline] + fn _copy_group(child: &ChunkedLevel, parent_idx: usize) -> [Hash; N] { + let mut buf = [[0u8; 32]; N]; + child.get_group(parent_idx * N, N, &mut buf); + buf + } + + pub(crate) fn _insert_many( + inner: &mut TreeInner, + hasher: &H, + leaves: &[Hash], + ) -> Result { + if leaves.is_empty() { + return Err(TreeError::EmptyBatch); + } + + let batch_len = u64::try_from(leaves.len()).unwrap_or(u64::MAX); + let new_size = inner + .size + .checked_add(batch_len) + .ok_or(TreeError::CapacityExceeded)?; + let depth = ceil_log_n(new_size, N); + if depth > MAX_DEPTH { + return Err(TreeError::MaxDepthExceeded { + max_depth: MAX_DEPTH, + }); + } + + inner.levels[0].extend(leaves)?; + + // allocate upfront + { + let mut level_len = inner.levels[0].len(); + for level in 0..depth { + let num_parents = level_len.div_ceil(N); + if level + 1 < MAX_DEPTH { + inner.levels[level + 1].ensure_len(num_parents)?; + } + level_len = num_parents; + } + } + + let old_size_usize = u64_to_usize(inner.size)?; + let mut start_parent = old_size_usize / N; + + let mut root = if depth == 0 { + inner.levels[0].get(0)? + } else { + [0u8; 32] + }; + + #[cfg(feature = "parallel")] + let mut par_buf: std::vec::Vec = std::vec::Vec::new(); + #[cfg(feature = "parallel")] + let par_threshold = parallel_threshold(); + + for level in 0..depth { + let level_len = inner.levels[level].len(); + let num_parents = level_len.div_ceil(N); + let is_root_level = level + 1 == depth; + + #[cfg(feature = "parallel")] + { + let work = num_parents - start_parent; + if work >= par_threshold { + use rayon::prelude::*; + + let split_at = level + 1; + let (child_levels, parent_levels) = + inner.levels.split_at_mut(split_at); + let child_level = &child_levels[level]; + + par_buf.clear(); + par_buf.reserve(work); + + let spare = &mut par_buf.spare_capacity_mut()[..work]; + spare.par_chunks_mut(PAR_CHUNK_SIZE).enumerate().for_each( + |(ci, chunk)| { + let base = start_parent + ci * PAR_CHUNK_SIZE; + for (i, slot) in chunk.iter_mut().enumerate() { + slot.write( + Self::_compute_parent( + child_level, + base + i, + level_len, + hasher, + ) + .expect("ensure_len guarantees valid indices"), + ); + } + }, + ); + // SAFETY: the parallel loop above initialised every + // element in `spare[..work]` via `MaybeUninit::write`. + unsafe { par_buf.set_len(work) }; + + let parent_level = &mut parent_levels[0]; + for (i, &parent) in par_buf.iter().enumerate() { + let parent_idx = start_parent + i; + if split_at < MAX_DEPTH { + parent_level.set_preallocated(parent_idx, parent); + } + if is_root_level { + root = parent; + } + } + } else { + Self::_insert_many_level_seq( + &mut inner.levels, + level, + start_parent, + num_parents, + level_len, + is_root_level, + hasher, + &mut root, + )?; + } + } + + #[cfg(not(feature = "parallel"))] + Self::_insert_many_level_seq( + &mut inner.levels, + level, + start_parent, + num_parents, + level_len, + is_root_level, + hasher, + &mut root, + )?; + + start_parent /= N; + } + + inner.root = Some(root); + inner.size = new_size; + inner.depth = depth; + Ok(root) + } +} + +#[cfg(test)] +#[cfg_attr(feature = "concurrent", allow(unused_mut))] +mod tests { + #[cfg(not(feature = "std"))] + use alloc::vec::Vec; + #[cfg(feature = "std")] + use std::vec::Vec; + + use super::*; + use crate::{ + chunked_level::{ + CHUNK_SIZE, + Chunk, + }, + test_util::*, + }; + + #[test] + fn ceil_log_n_empty() { + assert_eq!(ceil_log_n(0, 2), 0); + } + + #[test] + fn ceil_log_n_one() { + assert_eq!(ceil_log_n(1, 2), 0); + } + + #[test] + fn ceil_log_n_binary_full() { + assert_eq!(ceil_log_n(4, 2), 2); + } + + #[test] + fn ceil_log_n_binary_partial() { + assert_eq!(ceil_log_n(3, 2), 2); + } + + #[test] + fn ceil_log_n_ternary() { + assert_eq!(ceil_log_n(4, 3), 2); + } + + #[test] + fn ceil_log_n_ternary_exact() { + assert_eq!(ceil_log_n(9, 3), 2); + } + + #[test] + fn ceil_log_n_quaternary() { + assert_eq!(ceil_log_n(16, 4), 2); + assert_eq!(ceil_log_n(17, 4), 3); + } + + #[test] + fn ceil_log_n_large_n() { + assert_eq!(ceil_log_n(256, 16), 2); + assert_eq!(ceil_log_n(257, 16), 3); + } + + #[test] + fn chunked_level_push_and_get() { + let mut level = ChunkedLevel::new(); + for i in 0u32..10 { + level.push(leaf(i)).unwrap(); + } + assert_eq!(level.len(), 10); + for i in 0u32..10 { + assert_eq!(level.get(i as usize).unwrap(), leaf(i)); + } + } + + #[test] + fn chunked_level_promotes_at_chunk_size() { + let mut level = ChunkedLevel::new(); + for i in 0..CHUNK_SIZE { + level.push(leaf(i as u32)).unwrap(); + } + assert_eq!(level.chunk_count(), 1); + assert_eq!(level.tail_len(), 0); + assert_eq!(level.len(), CHUNK_SIZE); + + // One more goes into the new tail. + level.push(leaf(0xFF)).unwrap(); + assert_eq!(level.chunk_count(), 1); + assert_eq!(level.tail_len(), 1); + assert_eq!(level.len(), CHUNK_SIZE + 1); + } + + #[test] + fn chunked_level_clone_shares_arcs() { + let mut level = ChunkedLevel::new(); + for i in 0..CHUNK_SIZE + 5 { + level.push(leaf(i as u32)).unwrap(); + } + let snap = level.clone(); + assert_eq!(snap.len(), level.len()); + // The completed chunk Arc is shared. + assert!(Chunk::ptr_eq(level.get_chunk(0), snap.get_chunk(0))); + // Data matches. + for i in 0..level.len() { + assert_eq!(level.get(i).unwrap(), snap.get(i).unwrap()); + } + } + + #[test] + fn empty_tree() { + let tree = LeanIMT::::new(XorHasher); + assert_eq!(tree.root(), None); + assert_eq!(tree.size(), 0); + assert_eq!(tree.depth(), 0); + } + + #[test] + fn insert_single_leaf_binary() { + let mut tree = LeanIMT::::new(XorHasher); + let l = leaf(1); + let root = tree.insert(l).unwrap(); + assert_eq!(root, l); // single leaf stored as-is + assert_eq!(tree.size(), 1); + assert_eq!(tree.depth(), 0); + } + + #[test] + fn insert_two_leaves_binary() { + let th = XorHasher; + let mut tree = LeanIMT::::new(XorHasher); + let l0 = leaf(1); + let l1 = leaf(2); + tree.insert(l0).unwrap(); + let root = tree.insert(l1).unwrap(); + + let expected = th.hash_children(&[l0, l1]); + assert_eq!(root, expected); + assert_eq!(tree.size(), 2); + assert_eq!(tree.depth(), 1); + } + + #[test] + fn insert_three_leaves_binary() { + let th = XorHasher; + let mut tree = LeanIMT::::new(XorHasher); + let l0 = leaf(1); + let l1 = leaf(2); + let l2 = leaf(3); + tree.insert(l0).unwrap(); + tree.insert(l1).unwrap(); + let root = tree.insert(l2).unwrap(); + + // Level 0: [l0, l1, l2] + // Level 1: [H(l0,l1), l2_lifted] + // Level 2: [H(H(l0,l1), l2)] + let h01 = th.hash_children(&[l0, l1]); + let expected = th.hash_children(&[h01, l2]); + assert_eq!(root, expected); + assert_eq!(tree.depth(), 2); + } + + #[test] + fn insert_four_leaves_binary() { + let th = XorHasher; + let mut tree = LeanIMT::::new(XorHasher); + let leaves: Vec = (1..=4).map(leaf).collect(); + for &l in &leaves { + tree.insert(l).unwrap(); + } + + // Level 0: [l0, l1, l2, l3] + // Level 1: [H(l0,l1), H(l2,l3)] + // Level 2: [H(H(l0,l1), H(l2,l3))] + let h01 = th.hash_children(&[leaves[0], leaves[1]]); + let h23 = th.hash_children(&[leaves[2], leaves[3]]); + let expected = th.hash_children(&[h01, h23]); + assert_eq!(tree.root(), Some(expected)); + assert_eq!(tree.depth(), 2); + } + + #[test] + fn insert_four_leaves_ternary() { + let th = XorHasher; + let mut tree = LeanIMT::::new(XorHasher); + let leaves: Vec = (1..=4).map(leaf).collect(); + for &l in &leaves { + tree.insert(l).unwrap(); + } + + // Level 0: [l0, l1, l2, l3] + // Level 1: [H(l0,l1,l2), l3_lifted] + // Level 2: [H(H(l0,l1,l2), l3)] + let h012 = th.hash_children(&[leaves[0], leaves[1], leaves[2]]); + let expected = th.hash_children(&[h012, leaves[3]]); + assert_eq!(tree.root(), Some(expected)); + assert_eq!(tree.depth(), 2); + } + + #[test] + fn insert_two_leaves_ternary() { + let th = XorHasher; + let mut tree = LeanIMT::::new(XorHasher); + let l0 = leaf(1); + let l1 = leaf(2); + tree.insert(l0).unwrap(); + let root = tree.insert(l1).unwrap(); + + let expected = th.hash_children(&[l0, l1]); + assert_eq!(root, expected); + assert_eq!(tree.depth(), 1); + } + + #[test] + fn insert_five_leaves_quaternary() { + let th = XorHasher; + let mut tree = LeanIMT::::new(XorHasher); + let leaves: Vec = (1..=5).map(leaf).collect(); + for &l in &leaves { + tree.insert(l).unwrap(); + } + + // Level 0: [l0..l4] + // Level 1: [H(l0,l1,l2,l3), l4_lifted] + // Level 2: [H(H(l0,l1,l2,l3), l4)] + let h0123 = th.hash_children(&[leaves[0], leaves[1], leaves[2], leaves[3]]); + let expected = th.hash_children(&[h0123, leaves[4]]); + assert_eq!(tree.root(), Some(expected)); + assert_eq!(tree.depth(), 2); + } + + #[test] + fn insert_many_matches_sequential_binary() { + let h = XorHasher; + let leaves: Vec = (1..=7).map(leaf).collect(); + + let mut seq = LeanIMT::::new(h.clone()); + for &l in &leaves { + seq.insert(l).unwrap(); + } + + let mut batch = LeanIMT::::new(h.clone()); + batch.insert_many(&leaves).unwrap(); + + assert_eq!(seq.root(), batch.root()); + assert_eq!(seq.size(), batch.size()); + } + + #[test] + fn insert_many_matches_sequential_ternary() { + let h = XorHasher; + let leaves: Vec = (1..=10).map(leaf).collect(); + + let mut seq = LeanIMT::::new(h.clone()); + for &l in &leaves { + seq.insert(l).unwrap(); + } + + let mut batch = LeanIMT::::new(h.clone()); + batch.insert_many(&leaves).unwrap(); + + assert_eq!(seq.root(), batch.root()); + } + + #[test] + fn insert_many_incremental() { + let h = XorHasher; + let leaves: Vec = (1..=10).map(leaf).collect(); + + let mut seq = LeanIMT::::new(h.clone()); + for &l in &leaves { + seq.insert(l).unwrap(); + } + + let mut mixed = LeanIMT::::new(h.clone()); + for &l in &leaves[..3] { + mixed.insert(l).unwrap(); + } + mixed.insert_many(&leaves[3..]).unwrap(); + + assert_eq!(seq.root(), mixed.root()); + } + + #[test] + fn insert_many_empty_batch_error() { + let mut tree = LeanIMT::::new(XorHasher); + assert_eq!(tree.insert_many(&[]), Err(TreeError::EmptyBatch)); + } + + #[test] + fn insert_many_chunk_boundary() { + let leaves: Vec = (0..CHUNK_SIZE) + .map(|i| { + let mut h = [0u8; 32]; + let bytes = (i as u64).to_le_bytes(); + h[..8].copy_from_slice(&bytes); + h + }) + .collect(); + + let mut seq = LeanIMT::::new(XorHasher); + for &l in &leaves { + seq.insert(l).unwrap(); + } + + let mut batch = LeanIMT::::new(XorHasher); + batch.insert_many(&leaves).unwrap(); + + assert_eq!(seq.root(), batch.root()); + assert_eq!(seq.size(), batch.size()); + assert_eq!(seq.size(), CHUNK_SIZE as u64); + } + + #[test] + fn max_depth_exceeded() { + let mut tree = LeanIMT::::new(XorHasher); + let l = [0u8; 32]; + tree.insert(l).unwrap(); // size=1, depth=0 + tree.insert(l).unwrap(); // size=2, depth=1 + let err = tree.insert(l).unwrap_err(); + assert_eq!(err, TreeError::MaxDepthExceeded { max_depth: 1 }); + } + + #[cfg(feature = "blake3")] + mod blake3_tests { + use super::*; + use crate::Blake3Hasher; + + fn blake3_leaf(n: u8) -> Hash { + *::blake3::hash(&[n]).as_bytes() + } + + #[test] + fn binary_four_leaves_known_vector() { + let th = Blake3Hasher; + let mut tree = LeanIMT::::new(Blake3Hasher); + + let l0 = blake3_leaf(0); + let l1 = blake3_leaf(1); + let l2 = blake3_leaf(2); + let l3 = blake3_leaf(3); + + let r1 = tree.insert(l0).unwrap(); + assert_eq!(r1, l0); + + let r2 = tree.insert(l1).unwrap(); + let h01 = th.hash_children(&[l0, l1]); + assert_eq!(r2, h01); + + let r3 = tree.insert(l2).unwrap(); + let expected3 = th.hash_children(&[h01, l2]); + assert_eq!(r3, expected3); + + let r4 = tree.insert(l3).unwrap(); + let h23 = th.hash_children(&[l2, l3]); + let expected4 = th.hash_children(&[h01, h23]); + assert_eq!(r4, expected4); + } + + #[test] + fn ternary_four_leaves_known_vector() { + let th = Blake3Hasher; + let mut tree = LeanIMT::::new(Blake3Hasher); + + let l0 = blake3_leaf(0); + let l1 = blake3_leaf(1); + let l2 = blake3_leaf(2); + let l3 = blake3_leaf(3); + + tree.insert(l0).unwrap(); + + let r2 = tree.insert(l1).unwrap(); + assert_eq!(r2, th.hash_children(&[l0, l1])); + + let r3 = tree.insert(l2).unwrap(); + assert_eq!(r3, th.hash_children(&[l0, l1, l2])); + + let r4 = tree.insert(l3).unwrap(); + let h012 = th.hash_children(&[l0, l1, l2]); + assert_eq!(r4, th.hash_children(&[h012, l3])); + } + + #[test] + fn quaternary_five_leaves_known_vector() { + let th = Blake3Hasher; + let mut tree = LeanIMT::::new(Blake3Hasher); + + let leaves: Vec = (0..5).map(blake3_leaf).collect(); + for &l in &leaves { + tree.insert(l).unwrap(); + } + + let h0123 = th.hash_children(&[leaves[0], leaves[1], leaves[2], leaves[3]]); + let expected = th.hash_children(&[h0123, leaves[4]]); + assert_eq!(tree.root(), Some(expected)); + } + } +} diff --git a/crates/rotortree/tests/concurrent.rs b/crates/rotortree/tests/concurrent.rs new file mode 100644 index 0000000..a2711d6 --- /dev/null +++ b/crates/rotortree/tests/concurrent.rs @@ -0,0 +1,172 @@ +#![cfg(feature = "concurrent")] + +use std::{ + sync::Arc, + thread, +}; + +use rotortree::{ + Hash, + LeanIMT, + test_util::{ + XorHasher, + leaf, + }, +}; + +#[test] +fn concurrent_insert_single_thread() { + let tree = LeanIMT::::new(XorHasher); + for i in 1..=10u32 { + tree.insert(leaf(i)).unwrap(); + } + assert_eq!(tree.size(), 10); + + let snap = tree.snapshot(); + for i in 0..10u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } +} + +#[test] +fn concurrent_multi_thread_insert() { + let tree = Arc::new(LeanIMT::::new(XorHasher)); + let num_threads = 4; + let leaves_per_thread = 50; + + let handles: Vec<_> = (0..num_threads) + .map(|t| { + let tree = Arc::clone(&tree); + thread::spawn(move || { + let base = (t * leaves_per_thread) as u8; + for i in 0..leaves_per_thread as u8 { + let mut l = [0u8; 32]; + l[0] = base.wrapping_add(i); + l[1] = t as u8; + tree.insert(l).unwrap(); + } + }) + }) + .collect(); + + for h in handles { + h.join().unwrap(); + } + + let total = num_threads * leaves_per_thread; + assert_eq!(tree.size(), total as u64); + + let snap = tree.snapshot(); + for i in 0..total as u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } +} + +#[test] +fn concurrent_reader_writer() { + let tree = Arc::new(LeanIMT::::new(XorHasher)); + let num_inserts = 200u64; + let num_readers = 3; + + let writer_tree = Arc::clone(&tree); + let writer = thread::spawn(move || { + for i in 0..num_inserts { + let mut l = [0u8; 32]; + l[0] = i as u8; + l[1] = (i >> 8) as u8; + writer_tree.insert(l).unwrap(); + } + }); + + let readers: Vec<_> = (0..num_readers) + .map(|_| { + let tree = Arc::clone(&tree); + thread::spawn(move || { + for _ in 0..100 { + let snap = tree.snapshot(); + let size = snap.size(); + if size == 0 { + continue; + } + for i in 0..size { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + } + }) + }) + .collect(); + + writer.join().unwrap(); + for r in readers { + r.join().unwrap(); + } + + assert_eq!(tree.size(), num_inserts); +} + +#[test] +fn snapshot_isolation() { + let tree = LeanIMT::::new(XorHasher); + for i in 1..=5u32 { + tree.insert(leaf(i)).unwrap(); + } + + let snap = tree.snapshot(); + let snap_root = snap.root(); + let snap_size = snap.size(); + let snap_depth = snap.depth(); + + // Insert more after taking the snapshot. + for i in 6..=10u32 { + tree.insert(leaf(i)).unwrap(); + } + + // The snapshot must be unchanged. + assert_eq!(snap.root(), snap_root); + assert_eq!(snap.size(), snap_size); + assert_eq!(snap.depth(), snap_depth); + + // But the tree itself has advanced. + assert_eq!(tree.size(), 10); + assert_ne!(tree.root(), snap_root); +} + +#[test] +fn concurrent_insert_many() { + let tree = Arc::new(LeanIMT::::new(XorHasher)); + let num_threads = 4; + let batch_size = 25; + + let handles: Vec<_> = (0..num_threads) + .map(|t| { + let tree = Arc::clone(&tree); + thread::spawn(move || { + let batch: Vec = (0..batch_size) + .map(|i| { + let mut l = [0u8; 32]; + l[0] = i as u8; + l[1] = t as u8; + l + }) + .collect(); + tree.insert_many(&batch).unwrap(); + }) + }) + .collect(); + + for h in handles { + h.join().unwrap(); + } + + let total = (num_threads * batch_size) as u64; + assert_eq!(tree.size(), total); + + let snap = tree.snapshot(); + for i in 0..total { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } +} diff --git a/crates/rotortree/tests/lib.rs b/crates/rotortree/tests/lib.rs new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/crates/rotortree/tests/lib.rs @@ -0,0 +1 @@ + diff --git a/crates/rotortree/tests/parallel.rs b/crates/rotortree/tests/parallel.rs new file mode 100644 index 0000000..88e435e --- /dev/null +++ b/crates/rotortree/tests/parallel.rs @@ -0,0 +1,97 @@ +#![cfg(feature = "parallel")] +#![cfg_attr(feature = "concurrent", allow(unused_mut))] + +use rotortree::{ + Hash, + LeanIMT, + test_util::XorHasher, +}; + +fn make_leaves(count: usize) -> Vec { + (0..count) + .map(|i| { + let mut h = [0u8; 32]; + let bytes = (i as u64).to_le_bytes(); + h[..8].copy_from_slice(&bytes); + h + }) + .collect() +} + +#[test] +fn parallel_insert_many_matches_sequential_binary() { + let leaves = make_leaves(1000); + + let mut seq = LeanIMT::::new(XorHasher); + for &l in &leaves { + seq.insert(l).unwrap(); + } + + let mut batch = LeanIMT::::new(XorHasher); + batch.insert_many(&leaves).unwrap(); + + assert_eq!(seq.root(), batch.root()); + assert_eq!(seq.size(), batch.size()); +} + +#[test] +fn parallel_insert_many_matches_sequential_ternary() { + let leaves = make_leaves(1000); + + let mut seq = LeanIMT::::new(XorHasher); + for &l in &leaves { + seq.insert(l).unwrap(); + } + + let mut batch = LeanIMT::::new(XorHasher); + batch.insert_many(&leaves).unwrap(); + + assert_eq!(seq.root(), batch.root()); +} + +#[test] +fn parallel_insert_many_matches_sequential_quaternary() { + let leaves = make_leaves(1000); + + let mut seq = LeanIMT::::new(XorHasher); + for &l in &leaves { + seq.insert(l).unwrap(); + } + + let mut batch = LeanIMT::::new(XorHasher); + batch.insert_many(&leaves).unwrap(); + + assert_eq!(seq.root(), batch.root()); +} + +#[test] +fn parallel_large_batch_proofs() { + let leaves = make_leaves(2000); + + let mut tree = LeanIMT::::new(XorHasher); + tree.insert_many(&leaves).unwrap(); + + let snap = tree.snapshot(); + for i in 0..2000u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } +} + +#[test] +fn parallel_insert_many_incremental() { + let leaves = make_leaves(1000); + + let mut seq = LeanIMT::::new(XorHasher); + for &l in &leaves { + seq.insert(l).unwrap(); + } + + let mut mixed = LeanIMT::::new(XorHasher); + for &l in &leaves[..100] { + mixed.insert(l).unwrap(); + } + mixed.insert_many(&leaves[100..]).unwrap(); + + assert_eq!(seq.root(), mixed.root()); +} diff --git a/crates/rotortree/tests/proptest_tree.rs b/crates/rotortree/tests/proptest_tree.rs new file mode 100644 index 0000000..b124513 --- /dev/null +++ b/crates/rotortree/tests/proptest_tree.rs @@ -0,0 +1,216 @@ +#![cfg_attr(feature = "concurrent", allow(unused_mut))] + +use proptest::prelude::*; +use rotortree::{ + Hash, + LeanIMT, + test_util::XorHasher, +}; + +fn leaves_strategy(max: usize) -> impl Strategy> { + prop::collection::vec(prop::array::uniform32(any::()), 1..=max) +} + +#[crabtime::function] +fn gen_root_deterministic(n_values: Vec) { + for n in n_values { + crabtime::output! { + proptest! { + #[test] + fn root_deterministic_n{{n}}( + leaves in leaves_strategy(100) + ) { + let mut t1 = + LeanIMT::::new(XorHasher); + let mut t2 = + LeanIMT::::new(XorHasher); + for &leaf in &leaves { + t1.insert(leaf).unwrap(); + t2.insert(leaf).unwrap(); + } + prop_assert_eq!(t1.root(), t2.root()); + } + } + } + } +} + +gen_root_deterministic!([2, 3, 4]); + +#[crabtime::function] +fn gen_proof_round_trip(n_values: Vec) { + for n in n_values { + crabtime::output! { + proptest! { + #[test] + fn proof_round_trip_n{{n}}( + leaves in leaves_strategy(100) + ) { + let mut tree = + LeanIMT::::new(XorHasher); + for &leaf in &leaves { + tree.insert(leaf).unwrap(); + } + let snap = tree.snapshot(); + for i in 0..leaves.len() as u64 { + let proof = snap.generate_proof(i).unwrap(); + prop_assert!(proof.verify(&XorHasher).unwrap()); + } + } + } + } + } +} + +gen_proof_round_trip!([2, 3, 4]); + +#[crabtime::function] +fn gen_insert_many_equivalence(n_values: Vec) { + for n in n_values { + crabtime::output! { + proptest! { + #[test] + fn insert_many_equivalence_n{{n}}( + leaves in leaves_strategy(200) + ) { + let mut seq = + LeanIMT::::new(XorHasher); + for &leaf in &leaves { + seq.insert(leaf).unwrap(); + } + + let mut batch = + LeanIMT::::new(XorHasher); + batch.insert_many(&leaves).unwrap(); + + prop_assert_eq!(seq.root(), batch.root()); + } + } + } + } +} + +gen_insert_many_equivalence!([2, 3, 4]); + +#[crabtime::function] +fn gen_consistency_proof(n_values: Vec) { + for n in n_values { + crabtime::output! { + proptest! { + #[test] + fn consistency_proof_n{{n}}( + leaves in leaves_strategy(100) + ) { + let mut tree = + LeanIMT::::new(XorHasher); + let mut roots = Vec::new(); + let mut sizes = Vec::new(); + for &leaf in &leaves { + tree.insert(leaf).unwrap(); + let snap = tree.snapshot(); + roots.push(snap.root().unwrap()); + sizes.push(snap.size()); + } + let final_snap = tree.snapshot(); + for i in 0..leaves.len() { + let proof = final_snap + .generate_consistency_proof(sizes[i], roots[i]) + .unwrap(); + prop_assert!( + proof.verify(&XorHasher).unwrap(), + "n{{n}} consistency failed for size {} -> {}", + sizes[i], + final_snap.size() + ); + } + } + } + } + } +} + +gen_consistency_proof!([2, 3, 4]); + +#[crabtime::function] +fn gen_consistency_proof_update(n_values: Vec) { + for n in n_values { + crabtime::output! { + proptest! { + #[test] + fn consistency_proof_update_n{{n}}( + leaves in leaves_strategy(50) + ) { + let mut tree = + LeanIMT::::new(XorHasher); + let mut snaps = Vec::new(); + for &leaf in &leaves { + tree.insert(leaf).unwrap(); + let snap = tree.snapshot(); + snaps.push((snap.size(), snap.root().unwrap(), snap)); + } + let last = snaps.len() - 1; + for i in 0..snaps.len() { + let cp = snaps[last].2 + .generate_consistency_proof(snaps[i].0, snaps[i].1) + .unwrap(); + let old_ip = snaps[i].2.generate_proof(0).unwrap(); + if i == last { + let err = cp.update_inclusion_proof(&old_ip, &XorHasher).unwrap_err(); + prop_assert_eq!(err, rotortree::TreeError::NoUpdateNeeded); + continue; + } + let updated = cp.update_inclusion_proof(&old_ip, &XorHasher).unwrap(); + let fresh = snaps[last].2.generate_proof(0).unwrap(); + prop_assert_eq!(updated, fresh, + "n{{n}} update mismatch: size {} -> {}", snaps[i].0, snaps[last].0); + } + } + } + } + } +} + +gen_consistency_proof_update!([2, 3, 4]); + +proptest! { + #[test] + fn insert_many_incremental_binary( + first in leaves_strategy(50), + second in leaves_strategy(50), + ) { + let mut seq = + LeanIMT::::new(XorHasher); + for &leaf in first.iter().chain(second.iter()) { + seq.insert(leaf).unwrap(); + } + + let mut mixed = + LeanIMT::::new(XorHasher); + for &leaf in &first { + mixed.insert(leaf).unwrap(); + } + mixed.insert_many(&second).unwrap(); + + prop_assert_eq!(seq.root(), mixed.root()); + } +} + +#[cfg(feature = "blake3")] +proptest! { + #[test] + fn proof_round_trip_blake3_n2( + leaves in leaves_strategy(100) + ) { + let hasher = rotortree::Blake3Hasher; + let th = rotortree::Blake3Hasher; + let mut tree = LeanIMT::::new(hasher); + for &leaf in &leaves { + tree.insert(leaf).unwrap(); + } + let snap = tree.snapshot(); + for i in 0..leaves.len() as u64 { + let proof = snap.generate_proof(i).unwrap(); + prop_assert!(proof.verify(&th).unwrap()); + } + } +} diff --git a/crates/rotortree/tests/simd_batch.rs b/crates/rotortree/tests/simd_batch.rs new file mode 100644 index 0000000..51b5504 --- /dev/null +++ b/crates/rotortree/tests/simd_batch.rs @@ -0,0 +1,97 @@ +//! SIMD-batched parent hashing equivalence for the Blake3 hasher. +//! +//! The SIMD batch path in `_insert_many` (via `Blake3Hasher::hash_many_into`) +//! must produce a root byte-identical to feeding the same leaves through the +//! scalar one-at-a-time `insert`. The sequential path is the ground truth. + +#![cfg(feature = "blake3")] +#![cfg_attr(feature = "concurrent", allow(unused_mut))] + +use proptest::prelude::*; +use rotortree::{ + Blake3Hasher, + Hash, + LeanIMT, +}; + +fn leaf_at(seed: u64, i: usize) -> Hash { + let mut h = [0u8; 32]; + let v = seed.wrapping_add(i as u64).to_le_bytes(); + h[..8].copy_from_slice(&v); + h +} + +macro_rules! batch_eq_seq { + ($name:ident, $arity:literal) => { + proptest! { + #![proptest_config(ProptestConfig::with_cases(48))] + #[test] + fn $name( + // counts straddling the 128-entry chunk boundary and the + // 256*128 = 32768 chunk-of-chunks boundary, plus the lift and + // single-element edges. + n in prop::sample::select(vec![ + 1usize, 2, 127, 128, 129, 255, 256, 257, 512, + 32767, 32768, 32769, + ]), + seed in any::(), + ) { + let leaves: Vec = (0..n).map(|i| leaf_at(seed, i)).collect(); + + let mut seq = LeanIMT::::new(Blake3Hasher); + for &l in &leaves { + seq.insert(l).unwrap(); + } + + let mut batch = LeanIMT::::new(Blake3Hasher); + batch.insert_many(&leaves).unwrap(); + + prop_assert_eq!(seq.root(), batch.root(), "n={}", n); + prop_assert_eq!(seq.size(), batch.size()); + } + } + }; +} + +batch_eq_seq!(insert_many_matches_sequential_blake3_n2, 2); +batch_eq_seq!(insert_many_matches_sequential_blake3_n4, 4); +batch_eq_seq!(insert_many_matches_sequential_blake3_n8, 8); +batch_eq_seq!(insert_many_matches_sequential_blake3_n16, 16); + +/// Incremental: some leaves via single insert, the rest via batched +/// `insert_many`, must equal the all-sequential tree. Exercises a non-zero +/// `start_parent` and partial first group on the batched path. +macro_rules! incremental_eq { + ($name:ident, $arity:literal) => { + proptest! { + #![proptest_config(ProptestConfig::with_cases(48))] + #[test] + fn $name( + pre in prop::sample::select(vec![1usize, 3, 7, 128, 200]), + post in prop::sample::select(vec![1usize, 5, 129, 300, 1000]), + seed in any::(), + ) { + let total = pre + post; + let leaves: Vec = (0..total).map(|i| leaf_at(seed, i)).collect(); + + let mut seq = LeanIMT::::new(Blake3Hasher); + for &l in &leaves { + seq.insert(l).unwrap(); + } + + let mut mixed = LeanIMT::::new(Blake3Hasher); + for &l in &leaves[..pre] { + mixed.insert(l).unwrap(); + } + mixed.insert_many(&leaves[pre..]).unwrap(); + + prop_assert_eq!(seq.root(), mixed.root(), "pre={} post={}", pre, post); + } + } + }; +} + +incremental_eq!(incremental_matches_sequential_blake3_n2, 2); +incremental_eq!(incremental_matches_sequential_blake3_n4, 4); +incremental_eq!(incremental_matches_sequential_blake3_n8, 8); +incremental_eq!(incremental_matches_sequential_blake3_n16, 16); diff --git a/crates/rotortree/tests/storage_integration.rs b/crates/rotortree/tests/storage_integration.rs new file mode 100644 index 0000000..2c172ba --- /dev/null +++ b/crates/rotortree/tests/storage_integration.rs @@ -0,0 +1,1132 @@ +#![cfg(feature = "storage")] + +use std::sync::Arc; + +use rotortree::{ + CheckpointMeta, + CheckpointPolicy, + FlushPolicy, + Hash, + RotorTree, + RotorTreeConfig, + RotorTreeError, + StorageError, + TieringConfig, + test_util::*, + write_test_meta, +}; + +fn manual_config(dir: &std::path::Path) -> RotorTreeConfig { + RotorTreeConfig { + path: dir.to_path_buf(), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: CheckpointPolicy::default(), + tiering: TieringConfig::default(), + verify_checkpoint: true, + } +} + +#[test] +fn open_empty_close_reopen() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.root(), None); + assert_eq!(tree.size(), 0); + tree.close().unwrap(); + + // when + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // then + assert_eq!(tree.root(), None); + assert_eq!(tree.size(), 0); + tree.close().unwrap(); +} + +#[test] +fn insert_close_reopen() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + let (root, _token) = tree.insert(leaf(1)).unwrap(); + tree.flush().unwrap(); + let size = tree.size(); + // when + tree.close().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // then + assert_eq!(tree.root(), Some(root)); + assert_eq!(tree.size(), size); + tree.close().unwrap(); +} + +#[test] +fn insert_many_close_reopen() { + // given + let dir = tempfile::tempdir().unwrap(); + let leaves: Vec = (0..10u32).map(leaf).collect(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + let (root, _token) = tree.insert_many(&leaves).unwrap(); + tree.flush().unwrap(); + // when + tree.close().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // then + assert_eq!(tree.root(), Some(root)); + assert_eq!(tree.size(), 10); + tree.close().unwrap(); +} + +#[test] +fn many_inserts_close_reopen() { + // given + let dir = tempfile::tempdir().unwrap(); + let n = 50u32; + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + for i in 0..n { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + let root = tree.root(); + let snap = tree.snapshot(); + // Verify all proofs before close. + for i in 0..n as u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + // when + tree.close().unwrap(); + + // Reopen and verify. + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.root(), root); + assert_eq!(tree.size(), n as u64); + let snap = tree.snapshot(); + for i in 0..n as u64 { + let proof = snap.generate_proof(i).unwrap(); + // then + assert!(proof.verify(&XorHasher).unwrap()); + } + tree.close().unwrap(); +} + +#[test] +fn truncated_tail_recovery() { + // given + let dir = tempfile::tempdir().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + for i in 0..5u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + let root_after_5 = tree.root().unwrap(); + let size_after_5 = tree.size(); + + tree.insert(leaf(5)).unwrap(); + tree.flush().unwrap(); + tree.close().unwrap(); + + let wal_path = dir.path().join("wal"); + let data = std::fs::read(&wal_path).unwrap(); + // when: remove last few bytes + std::fs::write(&wal_path, &data[..data.len() - 10]).unwrap(); + + // then: recovery + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.size(), size_after_5); + assert_eq!(tree.root(), Some(root_after_5)); + tree.close().unwrap(); +} + +#[test] +fn crc_corruption_mid_file() { + // given + let dir = tempfile::tempdir().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + for i in 0..3u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + tree.close().unwrap(); + + let wal_path = dir.path().join("wal"); + let mut data = std::fs::read(&wal_path).unwrap(); + // when + let corrupt_offset = data.len() / 2; + data[corrupt_offset] ^= 0xFF; + std::fs::write(&wal_path, &data).unwrap(); + + // then: error + let result = + RotorTree::::open(XorHasher, manual_config(dir.path())); + assert!(result.is_err()); +} + +#[test] +fn config_mismatch() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + tree.close().unwrap(); + + // when, then + let result = + RotorTree::::open(XorHasher, manual_config(dir.path())); + match result { + Err(RotorTreeError::Storage(StorageError::ConfigMismatch { .. })) => {} + Err(e) => panic!("expected ConfigMismatch, got {e:?}"), + Ok(_) => panic!("expected ConfigMismatch, got Ok"), + } +} + +#[test] +fn durability_token_lifecycle() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + let (_root, token) = tree.insert(leaf(1)).unwrap(); + + assert!(!token.is_durable()); + + // when + tree.flush().unwrap(); + + // then + assert!(token.is_durable()); + + tree.close().unwrap(); +} + +#[test] +fn concurrent_insert_recover() { + // given + let dir = tempfile::tempdir().unwrap(); + let n_threads = 4; + let inserts_per_thread = 25u32; + let total = n_threads * inserts_per_thread as usize; + + let tree = Arc::new( + RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(), + ); + + std::thread::scope(|s| { + for t in 0..n_threads { + let tree = Arc::clone(&tree); + // when + s.spawn(move || { + for i in 0..inserts_per_thread { + let val = (t as u32) * inserts_per_thread + i; + tree.insert(leaf(val)).unwrap(); + } + }); + } + }); + + tree.flush().unwrap(); + // then + assert_eq!(tree.size(), total as u64); + + let snap = tree.snapshot(); + for i in 0..total as u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + + let root = tree.root(); + let tree = Arc::try_unwrap(tree) + .ok() + .expect("other Arc references still held"); + tree.close().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.size(), total as u64); + assert_eq!(tree.root(), root); + tree.close().unwrap(); +} + +#[test] +fn recovery_continuation() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + for i in 0..5u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + tree.close().unwrap(); + + // when + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.size(), 5); + for i in 5..10u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + tree.close().unwrap(); + + // then + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.size(), 10); + let snap = tree.snapshot(); + for i in 0..10u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + tree.close().unwrap(); +} + +#[test] +fn depth_change_persistence() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // when + for i in 0..5u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + let depth = tree.depth(); + let root = tree.root(); + assert!( + depth >= 2, + "depth should be at least 2 for 5 leaves with N=2" + ); + tree.close().unwrap(); + + // then + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.depth(), depth); + assert_eq!(tree.root(), root); + tree.close().unwrap(); +} + +#[test] +fn interleaved_single_and_batch() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // when + tree.insert(leaf(0)).unwrap(); + tree.insert(leaf(1)).unwrap(); + let batch: Vec = (2..7u32).map(leaf).collect(); + tree.insert_many(&batch).unwrap(); + tree.insert(leaf(7)).unwrap(); + tree.insert(leaf(8)).unwrap(); + let batch2: Vec = (9..15u32).map(leaf).collect(); + tree.insert_many(&batch2).unwrap(); + + tree.flush().unwrap(); + let root = tree.root(); + let size = tree.size(); + assert_eq!(size, 15); + tree.close().unwrap(); + + // then + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.size(), size); + assert_eq!(tree.root(), root); + + let snap = tree.snapshot(); + for i in 0..size { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + tree.close().unwrap(); +} + +#[test] +fn first_insert_after_recovery() { + // given + let dir = tempfile::tempdir().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + for i in 0..3u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + tree.close().unwrap(); + + // when + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.size(), 3); + + let (root, _) = tree.insert(leaf(3)).unwrap(); + tree.flush().unwrap(); + assert_eq!(tree.size(), 4); + assert_eq!(tree.root(), Some(root)); + + // then + let snap = tree.snapshot(); + for i in 0..4u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + tree.close().unwrap(); +} + +#[test] +fn file_locking() { + // given + let dir = tempfile::tempdir().unwrap(); + let _tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // when, then + let result = + RotorTree::::open(XorHasher, manual_config(dir.path())); + match result { + Err(RotorTreeError::Storage(StorageError::FileLocked)) => {} + Err(e) => panic!("expected FileLocked, got {e:?}"), + Ok(_) => panic!("expected FileLocked, got Ok"), + } +} + +#[test] +fn insert_durable_round_trip() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // when + let root = tree.insert_durable(leaf(1)).unwrap(); + assert_eq!(tree.root(), Some(root)); + assert_eq!(tree.size(), 1); + tree.close().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + // then + assert_eq!(tree.root(), Some(root)); + assert_eq!(tree.size(), 1); + tree.close().unwrap(); +} + +#[test] +fn checkpoint_round_trip() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + let leaves: Vec = (0..200u32).map(leaf).collect(); + let (root, _) = tree.insert_many(&leaves).unwrap(); + tree.flush().unwrap(); + // when + tree.checkpoint().unwrap(); + + assert!(dir.path().join("data").join("header.bin").exists()); + assert!(dir.path().join("data").join("checkpoint.meta").exists()); + assert!(dir.path().join("data").join("tails.bin").exists()); + + tree.close().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.root(), Some(root)); + assert_eq!(tree.size(), 200); + let snap = tree.snapshot(); + // then + for i in 0..200u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + tree.close().unwrap(); +} + +#[test] +fn checkpoint_then_more_inserts() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + let batch1: Vec = (0..50u32).map(leaf).collect(); + tree.insert_many(&batch1).unwrap(); + tree.flush().unwrap(); + tree.checkpoint().unwrap(); + + // when + let batch2: Vec = (50..100u32).map(leaf).collect(); + tree.insert_many(&batch2).unwrap(); + tree.flush().unwrap(); + + let root = tree.root(); + tree.close().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.root(), root); + assert_eq!(tree.size(), 100); + // then + let snap = tree.snapshot(); + for i in 0..100u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + tree.close().unwrap(); +} + +#[test] +fn checkpoint_manual_explicit() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + for i in 0..20u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + + assert!(!dir.path().join("data").exists()); + + // when + tree.checkpoint().unwrap(); + + assert!(dir.path().join("data").join("header.bin").exists()); + assert!(dir.path().join("data").join("checkpoint.meta").exists()); + + let root = tree.root(); + tree.close().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + // then + assert_eq!(tree.root(), root); + assert_eq!(tree.size(), 20); + tree.close().unwrap(); +} + +#[test] +fn checkpoint_on_close_policy() { + // given + let dir = tempfile::tempdir().unwrap(); + let config = RotorTreeConfig { + path: dir.path().to_path_buf(), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: CheckpointPolicy::OnClose, + tiering: TieringConfig::default(), + verify_checkpoint: true, + }; + let tree = RotorTree::::open(XorHasher, config).unwrap(); + + for i in 0..30u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + + assert!(!dir.path().join("data").exists()); + + // when + let root = tree.root(); + tree.close().unwrap(); + + assert!(dir.path().join("data").join("checkpoint.meta").exists()); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.root(), root); + assert_eq!(tree.size(), 30); + let snap = tree.snapshot(); + // then + for i in 0..30u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + tree.close().unwrap(); +} + +#[test] +fn checkpoint_every_n_entries() { + // given + let dir = tempfile::tempdir().unwrap(); + let config = RotorTreeConfig { + path: dir.path().to_path_buf(), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: CheckpointPolicy::EveryNEntries(10), + tiering: TieringConfig::default(), + verify_checkpoint: true, + }; + let tree = RotorTree::::open(XorHasher, config).unwrap(); + + // when + for i in 0..15u32 { + tree.insert(leaf(i)).unwrap(); + } + + assert!(tree.wait_for_checkpoint(std::time::Duration::from_secs(5))); + + let root = tree.root(); + tree.flush().unwrap(); + tree.close().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + // then + assert_eq!(tree.root(), root); + assert_eq!(tree.size(), 15); + tree.close().unwrap(); +} + +#[test] +fn checkpoint_memory_threshold() { + // given + let dir = tempfile::tempdir().unwrap(); + let config = RotorTreeConfig { + path: dir.path().to_path_buf(), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: CheckpointPolicy::MemoryThreshold(1), + tiering: TieringConfig::default(), + verify_checkpoint: true, + }; + let tree = RotorTree::::open(XorHasher, config).unwrap(); + + // when + let leaves: Vec = (0..200u32).map(leaf).collect(); + tree.insert_many(&leaves).unwrap(); + + assert!(tree.wait_for_checkpoint(std::time::Duration::from_secs(5))); + + let root = tree.root(); + tree.flush().unwrap(); + tree.close().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + // then + assert_eq!(tree.root(), root); + assert_eq!(tree.size(), 200); + tree.close().unwrap(); +} + +#[test] +fn checkpoint_empty_tree_noop() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // when + tree.checkpoint().unwrap(); + + // then + assert!(!dir.path().join("data").exists()); + + tree.close().unwrap(); +} + +#[test] +fn checkpoint_idempotent() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + for i in 0..30u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + + // when + tree.checkpoint().unwrap(); + let root_1 = tree.root(); + + tree.checkpoint().unwrap(); + let root_2 = tree.root(); + + assert_eq!(root_1, root_2); + + tree.close().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.root(), root_1); + assert_eq!(tree.size(), 30); + let snap = tree.snapshot(); + // then + for i in 0..30u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + tree.close().unwrap(); +} + +#[test] +fn checkpoint_after_recovery() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + for i in 0..30u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + let root = tree.root(); + tree.close().unwrap(); + + // when + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.size(), 30); + + tree.checkpoint().unwrap(); + assert!(dir.path().join("data").join("checkpoint.meta").exists()); + tree.close().unwrap(); + + // then + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.root(), root); + assert_eq!(tree.size(), 30); + let snap = tree.snapshot(); + for i in 0..30u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + tree.close().unwrap(); +} + +#[test] +fn root_recomputation_catches_bit_rot() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + let leaves: Vec = (0..200u32).map(leaf).collect(); + tree.insert_many(&leaves).unwrap(); + tree.flush().unwrap(); + tree.checkpoint().unwrap(); + tree.close().unwrap(); + + // when + let level_path = dir + .path() + .join("data") + .join("level_0") + .join("shard_0000.dat"); + let mut data = std::fs::read(&level_path).unwrap(); + data[16] ^= 0xFF; + std::fs::write(&level_path, &data).unwrap(); + + let result = + RotorTree::::open(XorHasher, manual_config(dir.path())); + // then + match result { + Err(RotorTreeError::Storage(StorageError::DataCorruption { .. })) => {} + Err(e) => panic!("expected DataCorruption, got {e:?}"), + Ok(_) => panic!("expected DataCorruption, got Ok"), + } +} + +#[test] +fn corrupt_checkpoint_meta_falls_back_to_wal() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + for i in 0..20u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + let root = tree.root(); + tree.close().unwrap(); + + // when + let data_dir = dir.path().join("data"); + std::fs::create_dir_all(&data_dir).unwrap(); + std::fs::write(data_dir.join("checkpoint.meta"), b"garbage").unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + // then + assert_eq!(tree.root(), root); + assert_eq!(tree.size(), 20); + tree.close().unwrap(); +} + +#[test] +fn cross_config_checkpoint_detection() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + for i in 0..10u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + tree.checkpoint().unwrap(); + tree.close().unwrap(); + + // when + let result = + RotorTree::::open(XorHasher, manual_config(dir.path())); + // then + match result { + Err(RotorTreeError::Storage(StorageError::ConfigMismatch { .. })) => {} + Err(e) => panic!("expected ConfigMismatch, got {e:?}"), + Ok(_) => panic!("expected ConfigMismatch, got Ok"), + } +} + +#[test] +fn inflated_leaf_count_meta() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + let leaves: Vec = (0..200u32).map(leaf).collect(); + tree.insert_many(&leaves).unwrap(); + tree.flush().unwrap(); + tree.checkpoint().unwrap(); + tree.close().unwrap(); + + // when + let data_dir = dir.path().join("data"); + write_test_meta( + &data_dir, + &CheckpointMeta { + n: 2, + max_depth: 10, + last_wal_seq: 0, + leaf_count: 500, + depth: 8, + root_hash: [0u8; 32], + }, + ) + .unwrap(); + + // then + let result = + RotorTree::::open(XorHasher, manual_config(dir.path())); + assert!( + result.is_err(), + "expected error from inflated leaf_count, got Ok" + ); +} + +#[test] +fn level_pinning() { + // given + let dir = tempfile::tempdir().unwrap(); + let config = RotorTreeConfig { + path: dir.path().to_path_buf(), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: CheckpointPolicy::default(), + tiering: TieringConfig { pin_above_level: 1 }, + verify_checkpoint: true, + }; + let tree = RotorTree::::open(XorHasher, config).unwrap(); + + let leaves: Vec = (0..200u32).map(leaf).collect(); + tree.insert_many(&leaves).unwrap(); + tree.flush().unwrap(); + // when + tree.checkpoint().unwrap(); + + let more: Vec = (200..250u32).map(leaf).collect(); + tree.insert_many(&more).unwrap(); + tree.flush().unwrap(); + + assert_eq!(tree.size(), 250); + let snap = tree.snapshot(); + let proof_0 = snap.generate_proof(0).unwrap(); + assert!(proof_0.verify(&XorHasher).unwrap()); + let proof_last = snap.generate_proof(249).unwrap(); + assert!(proof_last.verify(&XorHasher).unwrap()); + + let root = tree.root(); + tree.close().unwrap(); + + let config2 = RotorTreeConfig { + path: dir.path().to_path_buf(), + flush_policy: FlushPolicy::Manual, + checkpoint_policy: CheckpointPolicy::default(), + tiering: TieringConfig { pin_above_level: 1 }, + verify_checkpoint: true, + }; + // then + let tree = RotorTree::::open(XorHasher, config2).unwrap(); + assert_eq!(tree.root(), root); + assert_eq!(tree.size(), 250); + tree.close().unwrap(); +} + +#[test] +fn mmap_snapshot_isolation() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + let leaves: Vec = (0..200u32).map(leaf).collect(); + tree.insert_many(&leaves).unwrap(); + tree.flush().unwrap(); + tree.checkpoint().unwrap(); + + // when + let snap1 = tree.snapshot(); + let root1 = snap1.root(); + let size1 = snap1.size(); + + let more: Vec = (200..250u32).map(leaf).collect(); + tree.insert_many(&more).unwrap(); + tree.flush().unwrap(); + tree.checkpoint().unwrap(); + + assert_eq!(tree.size(), 250); + + // then + assert_eq!(snap1.root(), root1); + assert_eq!(snap1.size(), size1); + let proof_0 = snap1.generate_proof(0).unwrap(); + assert!(proof_0.verify(&XorHasher).unwrap()); + let proof_199 = snap1.generate_proof(199).unwrap(); + assert!(proof_199.verify(&XorHasher).unwrap()); + + let snap2 = tree.snapshot(); + let proof_249 = snap2.generate_proof(249).unwrap(); + assert!(proof_249.verify(&XorHasher).unwrap()); + + tree.close().unwrap(); +} + +#[test] +fn wal_truncated_after_checkpoint() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + for i in 0..50u32 { + tree.insert(leaf(i)).unwrap(); + } + tree.flush().unwrap(); + + let wal_path = dir.path().join("wal"); + let wal_size_before = std::fs::metadata(&wal_path).unwrap().len(); + assert!( + wal_size_before > 100, + "WAL should have entries before checkpoint" + ); + + // when + tree.checkpoint().unwrap(); + + let wal_size_after = std::fs::metadata(&wal_path).unwrap().len(); + assert!( + wal_size_after < 100, + "WAL should be truncated to header after checkpoint, got {wal_size_after} bytes" + ); + // then + assert!(wal_size_after < wal_size_before); + + let root = tree.root(); + tree.close().unwrap(); + + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.root(), root); + assert_eq!(tree.size(), 50); + tree.close().unwrap(); +} + +#[test] +fn max_depth_mismatch() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + tree.insert(leaf(1)).unwrap(); + tree.flush().unwrap(); + tree.close().unwrap(); + + // when: reopen with different MAX_DEPTH + let result = + RotorTree::::open(XorHasher, manual_config(dir.path())); + + // then + match result { + Err(RotorTreeError::Storage(StorageError::ConfigMismatch { + expected_max_depth: 20, + actual_max_depth: 10, + .. + })) => {} + Err(e) => panic!("expected ConfigMismatch with max_depth 20 vs 10, got {e:?}"), + Ok(_) => panic!("expected ConfigMismatch, got Ok"), + } +} + +#[test] +fn multiple_checkpoint_cycles() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // Cycle 1: insert 50, flush, checkpoint + let batch1: Vec = (0..50u32).map(leaf).collect(); + tree.insert_many(&batch1).unwrap(); + tree.flush().unwrap(); + tree.checkpoint().unwrap(); + + // Cycle 2: insert 50 more, flush, checkpoint + let batch2: Vec = (50..100u32).map(leaf).collect(); + tree.insert_many(&batch2).unwrap(); + tree.flush().unwrap(); + tree.checkpoint().unwrap(); + + // Cycle 3: insert 50 more, flush, checkpoint + let batch3: Vec = (100..150u32).map(leaf).collect(); + tree.insert_many(&batch3).unwrap(); + tree.flush().unwrap(); + tree.checkpoint().unwrap(); + + let root = tree.root(); + tree.close().unwrap(); + + // when: reopen + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // then + assert_eq!(tree.size(), 150); + assert_eq!(tree.root(), root); + let snap = tree.snapshot(); + for &idx in &[0u64, 49, 50, 99, 100, 149] { + let proof = snap.generate_proof(idx).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + tree.close().unwrap(); +} + +#[test] +fn file_lock_released_after_close() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + tree.insert(leaf(1)).unwrap(); + tree.flush().unwrap(); + tree.close().unwrap(); + + // when: open again in same dir + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // then: second open succeeded + assert_eq!(tree.size(), 1); + tree.close().unwrap(); +} + +#[test] +fn large_batch_storage_segment_freeze() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // 33,000 leaves → 257 chunks → triggers freeze_pending at 256 + let leaves: Vec = (0..33_000u32).map(leaf).collect(); + let (root, _) = tree.insert_many(&leaves).unwrap(); + tree.flush().unwrap(); + tree.checkpoint().unwrap(); + tree.close().unwrap(); + + // when: reopen + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + // then + assert_eq!(tree.size(), 33_000); + assert_eq!(tree.root(), Some(root)); + let snap = tree.snapshot(); + for &idx in &[0u64, 1000, 16383, 32767, 32999] { + let proof = snap.generate_proof(idx).unwrap(); + assert!( + proof.verify(&XorHasher).unwrap(), + "proof failed for idx {idx}" + ); + } + tree.close().unwrap(); +} + +#[test] +fn insert_many_after_checkpoint_recovery() { + // given + let dir = tempfile::tempdir().unwrap(); + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + + let batch1: Vec = (0..100u32).map(leaf).collect(); + tree.insert_many(&batch1).unwrap(); + tree.flush().unwrap(); + tree.checkpoint().unwrap(); + tree.close().unwrap(); + + // when: reopen and insert_many on recovered state + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.size(), 100); + + let batch2: Vec = (100..200u32).map(leaf).collect(); + tree.insert_many(&batch2).unwrap(); + tree.flush().unwrap(); + + let root = tree.root(); + tree.close().unwrap(); + + // then: reopen and verify full state + let tree = RotorTree::::open(XorHasher, manual_config(dir.path())) + .unwrap(); + assert_eq!(tree.size(), 200); + assert_eq!(tree.root(), root); + let snap = tree.snapshot(); + for &idx in &[0u64, 50, 99, 100, 150, 199] { + let proof = snap.generate_proof(idx).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + tree.close().unwrap(); +} diff --git a/crates/rotortree/tests/tree.rs b/crates/rotortree/tests/tree.rs new file mode 100644 index 0000000..6486df0 --- /dev/null +++ b/crates/rotortree/tests/tree.rs @@ -0,0 +1,443 @@ +#![cfg_attr(feature = "concurrent", allow(unused_mut))] + +use rotortree::{ + Hash, + Hasher, + LeanIMT, + TreeError, + test_util::*, +}; + +/// 130 leaves, N=3: group (126,127,128) spans chunk boundary → get_group slow path. +#[test] +fn cross_chunk_get_group_ternary() { + // given + let mut tree = LeanIMT::::new(XorHasher); + for i in 0..130u32 { + tree.insert(leaf(i)).unwrap(); + } + + // when + let snap = tree.snapshot(); + let proof = snap.generate_proof(127).unwrap(); + + // then + assert!(proof.verify(&XorHasher).unwrap()); +} + +/// 10 sequential + 300 batch exercises extend's 3-phase path (fill tail, full chunks, remainder). +#[test] +fn extend_three_phase_via_insert_many() { + // given + let mut seq = LeanIMT::::new(XorHasher); + let mut batch = LeanIMT::::new(XorHasher); + + let first: Vec = (0..10u32).map(leaf).collect(); + let second: Vec = (10..310u32).map(leaf).collect(); + let all: Vec = (0..310u32).map(leaf).collect(); + + for &l in &all { + seq.insert(l).unwrap(); + } + for &l in &first { + batch.insert(l).unwrap(); + } + + // when + batch.insert_many(&second).unwrap(); + + // then + assert_eq!(batch.root(), seq.root()); + assert_eq!(batch.size(), 310); +} + +/// 33,000 leaves → 257 chunks → triggers freeze_pending at 256. +#[test] +fn segment_freeze_large_batch() { + // given + let leaves: Vec = (0..33_000u32).map(leaf).collect(); + let mut seq = LeanIMT::::new(XorHasher); + for &l in &leaves { + seq.insert(l).unwrap(); + } + + // when + let mut batch = LeanIMT::::new(XorHasher); + batch.insert_many(&leaves).unwrap(); + + // then + assert_eq!(batch.root(), seq.root()); + assert_eq!(batch.size(), 33_000); + + let snap = batch.snapshot(); + for &idx in &[0u64, 1000, 16383, 32767, 32999] { + let proof = snap.generate_proof(idx).unwrap(); + assert!( + proof.verify(&XorHasher).unwrap(), + "proof failed for idx {idx}" + ); + } +} + +/// Insert 200, snapshot, insert 200 more → CoW on Arc-shared chunks. +#[test] +fn snapshot_cow_on_shared_chunks() { + // given + let mut tree = LeanIMT::::new(XorHasher); + for i in 0..200u32 { + tree.insert(leaf(i)).unwrap(); + } + let snap = tree.snapshot(); + let snap_root = snap.root(); + let snap_size = snap.size(); + + // when + for i in 200..400u32 { + tree.insert(leaf(i)).unwrap(); + } + + // then + assert_eq!(snap.root(), snap_root); + assert_eq!(snap.size(), snap_size); + assert_eq!(tree.size(), 400); + assert!(tree.root().is_some()); + + let tree_snap = tree.snapshot(); + let p_old = snap.generate_proof(0).unwrap(); + assert!(p_old.verify(&XorHasher).unwrap()); + let p_new = tree_snap.generate_proof(399).unwrap(); + assert!(p_new.verify(&XorHasher).unwrap()); +} + +#[test] +fn snapshot_get_node_and_level_len() { + // given + let mut tree = LeanIMT::::new(XorHasher); + for i in 0..20u32 { + tree.insert(leaf(i)).unwrap(); + } + let snap = tree.snapshot(); + let depth = snap.depth(); + + // then + assert_eq!(snap.level_len(0), 20); + assert_eq!(snap.level_len(1), 10); + assert_eq!(snap.level_len(depth), 1); + assert_eq!(snap.level_len(depth + 1), 0); + + let th = XorHasher; + assert_eq!(snap.get_node(0, 0).unwrap(), leaf(0)); + assert_eq!(snap.get_node(0, 19).unwrap(), leaf(19)); + let expected = th.hash_children(&[leaf(0), leaf(1)]); + assert_eq!(snap.get_node(1, 0).unwrap(), expected); + + match snap.get_node(depth + 1, 0) { + Err(TreeError::IndexOutOfRange { .. }) => {} + other => panic!("expected IndexOutOfRange, got {other:?}"), + } +} + +#[test] +fn proof_for_last_leaf() { + // given + let mut tree = LeanIMT::::new(XorHasher); + for i in 0..100u32 { + tree.insert(leaf(i)).unwrap(); + } + + // when + let snap = tree.snapshot(); + let proof = snap.generate_proof(99).unwrap(); + + // then + assert!(proof.verify(&XorHasher).unwrap()); + assert_eq!(proof.leaf, leaf(99)); +} + +/// 4 leaves (depth=2), 5th leaf increases depth to 3. All proofs must verify. +#[test] +fn proof_after_depth_increase() { + // given + let mut tree = LeanIMT::::new(XorHasher); + for i in 0..4u32 { + tree.insert(leaf(i)).unwrap(); + } + assert_eq!(tree.depth(), 2); + + // when + tree.insert(leaf(4)).unwrap(); + + // then + assert_eq!(tree.depth(), 3); + let snap = tree.snapshot(); + for i in 0..5u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!( + proof.verify(&XorHasher).unwrap(), + "proof failed for leaf {i}" + ); + } +} + +#[test] +fn proof_verify_rejects_bad_level_count() { + // given + let mut tree = LeanIMT::::new(XorHasher); + for i in 0..10u32 { + tree.insert(leaf(i)).unwrap(); + } + let snap = tree.snapshot(); + let mut proof = snap.generate_proof(0).unwrap(); + + // when + proof.level_count = 33; // > MAX_DEPTH + + // then + match proof.verify(&XorHasher) { + Err(TreeError::InvalidProofDepth { .. }) => {} + other => panic!("expected InvalidProofDepth, got {other:?}"), + } +} + +#[test] +fn proof_verify_rejects_bad_sibling_count() { + // given + let mut tree = LeanIMT::::new(XorHasher); + for i in 0..10u32 { + tree.insert(leaf(i)).unwrap(); + } + let snap = tree.snapshot(); + let mut proof = snap.generate_proof(0).unwrap(); + + // when + for level in &mut proof.levels[..proof.level_count] { + if level.sibling_count > 0 { + level.sibling_count = 2; + break; + } + } + + // then + match proof.verify(&XorHasher) { + Err(TreeError::MathError) => {} + other => panic!("expected MathError, got {other:?}"), + } +} + +#[test] +fn proof_verify_rejects_bad_position() { + // given + let mut tree = LeanIMT::::new(XorHasher); + for i in 0..10u32 { + tree.insert(leaf(i)).unwrap(); + } + let snap = tree.snapshot(); + let mut proof = snap.generate_proof(0).unwrap(); + + // when + for level in &mut proof.levels[..proof.level_count] { + if level.sibling_count > 0 { + level.position = level.sibling_count + 1; + break; + } + } + + // then + match proof.verify(&XorHasher) { + Err(TreeError::MathError) => {} + other => panic!("expected MathError, got {other:?}"), + } +} + +#[test] +fn large_branching_factor_n16() { + // given + let mut tree = LeanIMT::::new(XorHasher); + for i in 0..300u32 { + tree.insert(leaf(i)).unwrap(); + } + + // then + assert_eq!(tree.size(), 300); + assert!(tree.root().is_some()); + + let snap = tree.snapshot(); + for &idx in &[0u64, 15, 16, 17, 255, 256, 299] { + let proof = snap.generate_proof(idx).unwrap(); + assert!( + proof.verify(&XorHasher).unwrap(), + "proof failed for idx {idx}" + ); + } +} + +/// N=2, MAX_DEPTH=3 → capacity 8. Insert 8 (ok), 9th → MaxDepthExceeded. +#[test] +fn near_max_depth_binary() { + // given + let mut tree = LeanIMT::::new(XorHasher); + for i in 0..7u32 { + tree.insert(leaf(i)).unwrap(); + } + assert_eq!(tree.depth(), 3); + + // when + tree.insert(leaf(7)).unwrap(); + assert_eq!(tree.depth(), 3); + assert_eq!(tree.size(), 8); + + // then + match tree.insert(leaf(8)) { + Err(TreeError::MaxDepthExceeded { max_depth: 3 }) => {} + other => panic!("expected MaxDepthExceeded {{ max_depth: 3 }}, got {other:?}"), + } +} + +#[test] +fn insert_many_single_leaf() { + // given + let mut single = LeanIMT::::new(XorHasher); + single.insert(leaf(42)).unwrap(); + + // when + let mut batch = LeanIMT::::new(XorHasher); + batch.insert_many(&[leaf(42)]).unwrap(); + + // then + assert_eq!(single.root(), batch.root()); + assert_eq!(single.size(), batch.size()); +} + +/// Two insert_many calls of exactly CHUNK_SIZE (128) each. +#[test] +fn insert_many_exact_chunk_size() { + // given + let chunk1: Vec = (0..128u32).map(leaf).collect(); + let chunk2: Vec = (128..256u32).map(leaf).collect(); + let all: Vec = (0..256u32).map(leaf).collect(); + + let mut seq = LeanIMT::::new(XorHasher); + for &l in &all { + seq.insert(l).unwrap(); + } + + // when + let mut batch = LeanIMT::::new(XorHasher); + batch.insert_many(&chunk1).unwrap(); + batch.insert_many(&chunk2).unwrap(); + + // then + assert_eq!(batch.root(), seq.root()); + assert_eq!(batch.size(), 256); +} + +#[cfg(feature = "parallel")] +#[test] +fn parallel_threshold_boundary() { + // given + let small: Vec = (0..9u32).map(leaf).collect(); + let large: Vec = (0..500u32).map(leaf).collect(); + + // when + let mut seq_small = LeanIMT::::new(XorHasher); + for &l in &small { + seq_small.insert(l).unwrap(); + } + let mut batch_small = LeanIMT::::new(XorHasher); + batch_small.insert_many(&small).unwrap(); + + let mut seq_large = LeanIMT::::new(XorHasher); + for &l in &large { + seq_large.insert(l).unwrap(); + } + let mut batch_large = LeanIMT::::new(XorHasher); + batch_large.insert_many(&large).unwrap(); + + // then + assert_eq!(batch_small.root(), seq_small.root()); + assert_eq!(batch_large.root(), seq_large.root()); +} + +/// 8 writers (100 inserts each) + 4 readers verifying snapshot proofs concurrently. +#[cfg(feature = "concurrent")] +#[test] +fn concurrent_snapshot_proof_stress() { + use std::sync::{ + Arc, + atomic::{ + AtomicBool, + Ordering, + }, + }; + + // given + let tree = Arc::new(LeanIMT::::new(XorHasher)); + let done = Arc::new(AtomicBool::new(false)); + + // when + std::thread::scope(|s| { + for t in 0..8u32 { + let tree = Arc::clone(&tree); + s.spawn(move || { + for i in 0..100u32 { + tree.insert(leaf(t * 100 + i)).unwrap(); + } + }); + } + + for _ in 0..4 { + let tree = Arc::clone(&tree); + let done = Arc::clone(&done); + s.spawn(move || { + while !done.load(Ordering::Relaxed) { + let snap = tree.snapshot(); + let size = snap.size(); + if size == 0 { + continue; + } + for idx in [0, size / 2, size - 1] { + let proof = snap.generate_proof(idx).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } + if size >= 800 { + break; + } + } + }); + } + }); + + // then + assert_eq!(tree.size(), 800); + let snap = tree.snapshot(); + for i in 0..800u64 { + let proof = snap.generate_proof(i).unwrap(); + assert!(proof.verify(&XorHasher).unwrap()); + } +} + +/// Internal nodes are the hash of the concatenation of their children. +#[cfg(feature = "blake3")] +#[test] +fn internal_nodes_hash_children() { + use rotortree::Blake3Hasher; + + // given + let leaves: Vec = (0..4u32).map(leaf).collect(); + let mut tree = LeanIMT::::new(Blake3Hasher); + tree.insert_many(&leaves).unwrap(); + + let snap = tree.snapshot(); + let internal_left = snap.get_node(1, 0).unwrap(); + let internal_right = snap.get_node(1, 1).unwrap(); + + // when + let th = Blake3Hasher; + let h01 = th.hash_children(&[leaves[0], leaves[1]]); + let h23 = th.hash_children(&[leaves[2], leaves[3]]); + + // then + assert_eq!(internal_left, h01); + assert_eq!(internal_right, h23); + assert_eq!(tree.root().unwrap(), th.hash_children(&[h01, h23])); +}