diff --git a/Cargo.lock b/Cargo.lock index 84c72ce8a..682a7197a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3966,11 +3966,19 @@ dependencies = [ "proptest", "rand 0.8.5", "rayon", + "reth-chainspec", + "reth-db", "reth-errors", + "reth-evm", + "reth-evm-ethereum", "reth-execution-errors", + "reth-node-api", + "reth-node-ethereum", "reth-provider", + "reth-revm", "reth-trie", "reth-trie-db", + "reth-trie-parallel", "revm", "rustc-hash 2.1.1", "serde", diff --git a/Cargo.toml b/Cargo.toml index 92ef5359d..c8fdfc229 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -90,6 +90,7 @@ reth-provider = { git = "https://github.com/paradigmxyz/reth", rev = "27a8c0f5a6 reth-chainspec = { git = "https://github.com/paradigmxyz/reth", rev = "27a8c0f5a6dfb27dea84c5751776ecabdd069646" } reth-evm = { git = "https://github.com/paradigmxyz/reth", rev = "27a8c0f5a6dfb27dea84c5751776ecabdd069646" } reth-evm-ethereum = { git = "https://github.com/paradigmxyz/reth", rev = "27a8c0f5a6dfb27dea84c5751776ecabdd069646" } +reth-revm = { git = "https://github.com/paradigmxyz/reth", rev = "27a8c0f5a6dfb27dea84c5751776ecabdd069646" } reth-execution-errors = { git = "https://github.com/paradigmxyz/reth", rev = "27a8c0f5a6dfb27dea84c5751776ecabdd069646" } reth-trie-db = { git = "https://github.com/paradigmxyz/reth", rev = "27a8c0f5a6dfb27dea84c5751776ecabdd069646" } reth-transaction-pool = { git = "https://github.com/paradigmxyz/reth", rev = "27a8c0f5a6dfb27dea84c5751776ecabdd069646" } diff --git a/crates/eth-sparse-mpt/Cargo.toml b/crates/eth-sparse-mpt/Cargo.toml index 2111bd4ca..e918ac349 100644 --- a/crates/eth-sparse-mpt/Cargo.toml +++ b/crates/eth-sparse-mpt/Cargo.toml @@ -22,11 +22,19 @@ nybbles = { version = "0.3.3", features = ["serde"] } tracing.workspace = true # reth +reth-chainspec = { workspace = true, optional = true } +reth-db = { workspace = true, optional = true } +reth-evm = { workspace = true, optional = true } +reth-evm-ethereum = { workspace = true, optional = true } reth-errors.workspace = true reth-execution-errors.workspace = true +reth-node-api = { workspace = true, optional = true } +reth-node-ethereum = { workspace = true, optional = true } reth-trie.workspace = true +reth-trie-parallel = { workspace = true, optional = true } reth-trie-db.workspace = true reth-provider.workspace = true +reth-revm = { workspace = true, optional = true } # revm revm.workspace = true @@ -48,6 +56,16 @@ flate2 = { workspace = true, optional = true } [features] benchmark-utils = ["dep:hash-db", "dep:triehash", "dep:flate2"] +dev-tools = [ + "dep:reth-chainspec", + "dep:reth-db", + "dep:reth-evm", + "dep:reth-evm-ethereum", + "dep:reth-node-api", + "dep:reth-node-ethereum", + "dep:reth-trie-parallel", + "dep:reth-revm", +] [dev-dependencies] criterion = { workspace = true, features = ["html_reports"] } @@ -71,3 +89,21 @@ harness = false name = "trie_do_bench" harness = false +[[bench]] +name = "trie_bench_v2_v_experimental" +harness = false + +[[bin]] +name = "correctness-harness" +path = "src/bin/correctness-harness.rs" +required-features = ["dev-tools"] + +[[bin]] +name = "cache-warm-compare" +path = "src/bin/cache-warm-compare.rs" +required-features = ["dev-tools"] + +[[bin]] +name = "correct_and_bench" +path = "src/bin/correct_and_bench.rs" +required-features = ["dev-tools"] diff --git a/crates/eth-sparse-mpt/benches/trie_bench_v2_v_experimental.rs b/crates/eth-sparse-mpt/benches/trie_bench_v2_v_experimental.rs new file mode 100644 index 000000000..b857a85dd --- /dev/null +++ b/crates/eth-sparse-mpt/benches/trie_bench_v2_v_experimental.rs @@ -0,0 +1,60 @@ +use alloy_primitives::{keccak256, Bytes, B256, U256}; +use criterion::{criterion_group, criterion_main, Criterion}; +use eth_sparse_mpt::{ + v2::trie::{proof_store::ProofStore as ProofStoreV2, Trie as TrieV2}, + v_experimental::trie::{ + proof_store::ProofStore as ProofStoreVExperimental, Trie as TrieVExperimental, + }, +}; + +fn prepare_key_value_data(n: usize) -> (Vec, Vec) { + let mut keys = Vec::with_capacity(n); + let mut values = Vec::with_capacity(n); + for i in 0u64..(n as u64) { + let b: B256 = U256::from(i).into(); + let data = keccak256(b).to_vec(); + let value = keccak256(&data).to_vec(); + keys.push(Bytes::copy_from_slice(data.as_slice())); + values.push(Bytes::copy_from_slice(value.as_slice())); + } + (keys, values) +} + +fn insert_nodes_v2_v_experimental(c: &mut Criterion) { + let (keys, values) = prepare_key_value_data(10000); + + let empty_proof_store_v2 = ProofStoreV2::default(); + let mut v2_hash = B256::ZERO; + let mut trie_v2 = TrieV2::new_empty(); + c.bench_function("insert_nodes_v2", |b| { + b.iter(|| { + trie_v2.clear_empty(); + for (key, value) in keys.iter().zip(values.iter()) { + trie_v2.insert(key, value).unwrap(); + } + v2_hash = trie_v2.root_hash(true, &empty_proof_store_v2).unwrap(); + }) + }); + + let empty_proof_store_v_experimental = ProofStoreVExperimental::default(); + let mut v_experimental_hash = B256::ZERO; + let mut trie_v_experimental = TrieVExperimental::new_empty(); + c.bench_function("insert_nodes_v_experimental", |b| { + b.iter(|| { + trie_v_experimental.clear_empty(); + for (key, value) in keys.iter().zip(values.iter()) { + trie_v_experimental.insert(key, value).unwrap(); + } + v_experimental_hash = trie_v_experimental + .root_hash(true, &empty_proof_store_v_experimental) + .unwrap(); + }) + }); + + if !v2_hash.is_zero() && !v_experimental_hash.is_zero() { + assert_eq!(v2_hash, v_experimental_hash); + } +} + +criterion_group!(benches, insert_nodes_v2_v_experimental); +criterion_main!(benches); diff --git a/crates/eth-sparse-mpt/scripts/plot_cache_warm.py b/crates/eth-sparse-mpt/scripts/plot_cache_warm.py new file mode 100755 index 000000000..9f7b604b7 --- /dev/null +++ b/crates/eth-sparse-mpt/scripts/plot_cache_warm.py @@ -0,0 +1,381 @@ +#!/usr/bin/env python3 +"""Create plots comparing V2 vs v_experimental cache-warm benchmark results.""" + +from __future__ import annotations + +import argparse +import csv +from collections import defaultdict +from pathlib import Path +from statistics import mean + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Generate V2 vs v_experimental plots from cache-warm-combined.csv." + ) + parser.add_argument( + "--input", + type=Path, + default=Path("cache-warm-combined.csv"), + help="Input CSV path (default: cache-warm-combined.csv)", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("plots"), + help="Output directory for generated PNGs (default: plots)", + ) + parser.add_argument( + "--source", + default="", + help=( + "Optional source_file filter, e.g. cache-warm-1001blocks.csv. " + "Leave empty to use all rows." + ), + ) + parser.add_argument( + "--exclude-100", + action="store_true", + help="Exclude rows where warm_pct is 100.", + ) + parser.add_argument( + "--moving-window", + type=int, + default=50, + help="Window size for moving-average block plots (default: 50).", + ) + return parser.parse_args() + + +def read_rows( + path: Path, + source_filter: str, + exclude_100: bool, +) -> list[dict[str, str]]: + with path.open(newline="", encoding="utf-8") as f: + rows = list(csv.DictReader(f)) + if source_filter: + rows = [r for r in rows if r.get("source_file", "") == source_filter] + if exclude_100: + rows = [r for r in rows if to_int(r, "warm_pct") != 100] + if not rows: + raise ValueError("No rows found after applying filters.") + return rows + + +def load_plt(): + try: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + except ModuleNotFoundError as exc: + raise SystemExit( + "matplotlib is required. Install with: python3 -m pip install matplotlib" + ) from exc + return plt + + +def to_float(row: dict[str, str], key: str) -> float: + return float(row[key]) + +def to_float_any(row: dict[str, str], keys: tuple[str, ...]) -> float: + for key in keys: + value = row.get(key) + if value not in (None, ""): + return float(value) + raise KeyError(f"missing expected columns: {', '.join(keys)}") + + +def to_float_v_experimental(row: dict[str, str], metric: str) -> float: + return to_float_any( + row, + ( + f"v_experimental_{metric}_ms", + f"v3_{metric}_ms", + ), + ) + + +def to_int(row: dict[str, str], key: str) -> int: + return int(float(row[key])) + + +def aggregate_by_warm(rows: list[dict[str, str]]) -> dict[int, dict[str, float]]: + bucket: dict[int, dict[str, list[float]]] = defaultdict( + lambda: defaultdict(list) + ) + for row in rows: + warm = to_int(row, "warm_pct") + bucket[warm]["v2_p50"].append(to_float(row, "v2_p50_ms")) + bucket[warm]["v_experimental_p50"].append(to_float_v_experimental(row, "p50")) + bucket[warm]["v2_p99"].append(to_float(row, "v2_p99_ms")) + bucket[warm]["v_experimental_p99"].append(to_float_v_experimental(row, "p99")) + bucket[warm]["v2_mean"].append(to_float(row, "v2_mean_ms")) + bucket[warm]["v_experimental_mean"].append(to_float_v_experimental(row, "mean")) + + out: dict[int, dict[str, float]] = {} + for warm, values in bucket.items(): + out[warm] = { + "v2_p50": mean(values["v2_p50"]), + "v_experimental_p50": mean(values["v_experimental_p50"]), + "v2_p99": mean(values["v2_p99"]), + "v_experimental_p99": mean(values["v_experimental_p99"]), + "v2_mean": mean(values["v2_mean"]), + "v_experimental_mean": mean(values["v_experimental_mean"]), + } + return out + + +def save_speedup_by_warm_plot( + plt, by_warm: dict[int, dict[str, float]], output_dir: Path +) -> None: + warms = sorted(by_warm.keys()) + p50 = [by_warm[w]["v2_p50"] / by_warm[w]["v_experimental_p50"] for w in warms] + p99 = [by_warm[w]["v2_p99"] / by_warm[w]["v_experimental_p99"] for w in warms] + means = [by_warm[w]["v2_mean"] / by_warm[w]["v_experimental_mean"] for w in warms] + + x = range(len(warms)) + width = 0.25 + + plt.figure(figsize=(11, 6)) + plt.bar([i - width for i in x], p50, width=width, label="p50 speedup") + plt.bar(x, p99, width=width, label="p99 speedup") + plt.bar([i + width for i in x], means, width=width, label="mean speedup") + plt.xticks(list(x), [str(w) for w in warms]) + plt.xlabel("Warm %") + plt.ylabel("Speedup (V2 / v_experimental)") + plt.title("v_experimental Speedup by Warm %") + plt.grid(axis="y", alpha=0.3) + plt.legend() + plt.tight_layout() + plt.savefig(output_dir / "speedup_by_warm_pct.png", dpi=180) + plt.close() + + +def save_latency_by_warm_plot( + plt, by_warm: dict[int, dict[str, float]], output_dir: Path +) -> None: + warms = sorted(by_warm.keys()) + v2_p50 = [by_warm[w]["v2_p50"] for w in warms] + v_experimental_p50 = [by_warm[w]["v_experimental_p50"] for w in warms] + v2_p99 = [by_warm[w]["v2_p99"] for w in warms] + v_experimental_p99 = [by_warm[w]["v_experimental_p99"] for w in warms] + + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5)) + width = 0.35 + x = range(len(warms)) + + ax1.bar([i - width / 2 for i in x], v2_p50, width=width, label="V2") + ax1.bar( + [i + width / 2 for i in x], + v_experimental_p50, + width=width, + label="v_experimental", + ) + ax1.set_title("p50 Latency by Warm %") + ax1.set_xticks(list(x)) + ax1.set_xticklabels([str(w) for w in warms]) + ax1.set_xlabel("Warm %") + ax1.set_ylabel("Milliseconds") + ax1.grid(axis="y", alpha=0.3) + ax1.legend() + + ax2.bar([i - width / 2 for i in x], v2_p99, width=width, label="V2") + ax2.bar( + [i + width / 2 for i in x], + v_experimental_p99, + width=width, + label="v_experimental", + ) + ax2.set_title("p99 Latency by Warm %") + ax2.set_xticks(list(x)) + ax2.set_xticklabels([str(w) for w in warms]) + ax2.set_xlabel("Warm %") + ax2.set_ylabel("Milliseconds") + ax2.grid(axis="y", alpha=0.3) + ax2.legend() + + fig.suptitle("V2 vs v_experimental Latency by Warm %") + fig.tight_layout() + fig.savefig(output_dir / "latency_by_warm_pct.png", dpi=180) + plt.close(fig) + + +def save_speedup_over_blocks_plot( + plt, rows: list[dict[str, str]], output_dir: Path +) -> None: + points: dict[int, list[tuple[int, float]]] = defaultdict(list) + for row in rows: + if row.get("source_type") != "interleaved": + continue + warm = to_int(row, "warm_pct") + block = to_int(row, "iteration") + speedup = to_float(row, "v2_p50_ms") / to_float_v_experimental(row, "p50") + points[warm].append((block, speedup)) + + if not points: + return + + plt.figure(figsize=(12, 6)) + for warm in sorted(points): + series = sorted(points[warm], key=lambda x: x[0]) + x = [p[0] for p in series] + y = [p[1] for p in series] + plt.plot(x, y, linewidth=1.2, label=f"warm={warm}%") + + plt.xlabel("Block") + plt.ylabel("p50 Speedup (V2 / v_experimental)") + plt.title("p50 Speedup Over Blocks (Interleaved Runs)") + plt.grid(alpha=0.3) + plt.legend() + plt.tight_layout() + plt.savefig(output_dir / "p50_speedup_over_blocks.png", dpi=180) + plt.close() + + +def save_p99_regression_plot( + plt, rows: list[dict[str, str]], output_dir: Path +) -> None: + x: list[int] = [] + y: list[float] = [] + c: list[int] = [] + + for row in rows: + if row.get("source_type") != "interleaved": + continue + v2_p99 = to_float(row, "v2_p99_ms") + v_experimental_p99 = to_float_v_experimental(row, "p99") + if v_experimental_p99 < v2_p99: + continue + x.append(to_int(row, "iteration")) + y.append(v_experimental_p99 - v2_p99) + c.append(to_int(row, "warm_pct")) + + if not x: + return + + plt.figure(figsize=(12, 6)) + sc = plt.scatter(x, y, c=c, cmap="viridis", s=30, alpha=0.9) + cb = plt.colorbar(sc) + cb.set_label("Warm %") + plt.xlabel("Block") + plt.ylabel("p99 Regression (v_experimental - V2, ms)") + plt.title("p99 Regressions Where v_experimental >= V2 (Interleaved Runs)") + plt.grid(alpha=0.3) + plt.tight_layout() + plt.savefig(output_dir / "p99_regressions.png", dpi=180) + plt.close() + + +def moving_average(values: list[float], window: int) -> list[float]: + if window <= 1: + return values[:] + out: list[float] = [] + acc = 0.0 + for i, value in enumerate(values): + acc += value + if i >= window: + acc -= values[i - window] + count = min(i + 1, window) + out.append(acc / count) + return out + + +def save_latency_moving_average_plot( + plt, + rows: list[dict[str, str]], + output_dir: Path, + window: int, +) -> None: + by_iter: dict[int, dict[str, list[float]]] = defaultdict( + lambda: defaultdict(list) + ) + for row in rows: + if row.get("source_type") != "interleaved": + continue + block = to_int(row, "iteration") + by_iter[block]["v2_p50"].append(to_float(row, "v2_p50_ms")) + by_iter[block]["v_experimental_p50"].append(to_float_v_experimental(row, "p50")) + by_iter[block]["v2_p99"].append(to_float(row, "v2_p99_ms")) + by_iter[block]["v_experimental_p99"].append(to_float_v_experimental(row, "p99")) + + if not by_iter: + return + + blocks = sorted(by_iter.keys()) + v2_p50 = [mean(by_iter[i]["v2_p50"]) for i in blocks] + v_experimental_p50 = [mean(by_iter[i]["v_experimental_p50"]) for i in blocks] + v2_p99 = [mean(by_iter[i]["v2_p99"]) for i in blocks] + v_experimental_p99 = [mean(by_iter[i]["v_experimental_p99"]) for i in blocks] + + v2_p50_ma = moving_average(v2_p50, window) + v_experimental_p50_ma = moving_average(v_experimental_p50, window) + v2_p99_ma = moving_average(v2_p99, window) + v_experimental_p99_ma = moving_average(v_experimental_p99, window) + + fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 9), sharex=True) + ax1.plot(blocks, v2_p50_ma, linewidth=1.8, label="V2 p50 (moving avg)") + ax1.plot( + blocks, + v_experimental_p50_ma, + linewidth=1.8, + label="v_experimental p50 (moving avg)", + ) + ax1.set_ylabel("Milliseconds") + ax1.set_title(f"Moving Average p50 (window={window} blocks)") + ax1.grid(alpha=0.3) + ax1.legend() + + ax2.plot(blocks, v2_p99_ma, linewidth=1.8, label="V2 p99 (moving avg)") + ax2.plot( + blocks, + v_experimental_p99_ma, + linewidth=1.8, + label="v_experimental p99 (moving avg)", + ) + ax2.set_xlabel("Block") + ax2.set_ylabel("Milliseconds") + ax2.set_title(f"Moving Average p99 (window={window} blocks)") + ax2.grid(alpha=0.3) + ax2.legend() + + fig.suptitle("V2 vs v_experimental Latency Moving Averages") + fig.tight_layout() + fig.savefig(output_dir / "latency_moving_average_over_blocks.png", dpi=180) + plt.close(fig) + + +def main() -> None: + args = parse_args() + plt = load_plt() + if args.moving_window < 1: + raise SystemExit("--moving-window must be >= 1") + + rows = read_rows(args.input, args.source, args.exclude_100) + args.output_dir.mkdir(parents=True, exist_ok=True) + + by_warm = aggregate_by_warm(rows) + save_speedup_by_warm_plot(plt, by_warm, args.output_dir) + save_latency_by_warm_plot(plt, by_warm, args.output_dir) + save_speedup_over_blocks_plot(plt, rows, args.output_dir) + save_p99_regression_plot(plt, rows, args.output_dir) + save_latency_moving_average_plot(plt, rows, args.output_dir, args.moving_window) + + print(f"Generated plots in: {args.output_dir}") + if args.exclude_100: + print("Filter: excluding warm_pct=100 rows") + for name in [ + "speedup_by_warm_pct.png", + "latency_by_warm_pct.png", + "p50_speedup_over_blocks.png", + "p99_regressions.png", + "latency_moving_average_over_blocks.png", + ]: + path = args.output_dir / name + if path.exists(): + print(f" - {path}") + + +if __name__ == "__main__": + main() diff --git a/crates/eth-sparse-mpt/scripts/run_correct_and_bench_thread_sweep.sh b/crates/eth-sparse-mpt/scripts/run_correct_and_bench_thread_sweep.sh new file mode 100755 index 000000000..b9fe505fa --- /dev/null +++ b/crates/eth-sparse-mpt/scripts/run_correct_and_bench_thread_sweep.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Usage: +# scripts/run_correct_and_bench_thread_sweep.sh [OUTDIR] +# +# Optional environment overrides: +# THREADS="1 2 4 8" +# BIN=/home/robert/rbuilder/target/release/correct_and_bench +# RETH_BIN=/usr/local/bin/reth +# DATADIR=/mnt/reth-full +# CHAIN=mainnet +# BLOCKS=25 +# CACHE_WARM_ITERS=25 +# CACHE_WARM_KEYS=10000 +# CACHE_WARM_PERCENTAGES=70,80,90,100 +# STATIC_FILE_BLOCKS_PER_FILE=10000 + +THREADS="${THREADS:-1 2 4 8}" +BIN="${BIN:-/home/robert/rbuilder/target/release/correct_and_bench}" +RETH_BIN="${RETH_BIN:-/usr/local/bin/reth}" +DATADIR="${DATADIR:-/mnt/reth-full}" +CHAIN="${CHAIN:-mainnet}" +BLOCKS="${BLOCKS:-25}" +CACHE_WARM_ITERS="${CACHE_WARM_ITERS:-25}" +CACHE_WARM_KEYS="${CACHE_WARM_KEYS:-10000}" +CACHE_WARM_PERCENTAGES="${CACHE_WARM_PERCENTAGES:-70,80,90,100}" +# /mnt/reth-full commonly uses 10k static-file windows near the tip. +STATIC_FILE_BLOCKS_PER_FILE="${STATIC_FILE_BLOCKS_PER_FILE:-10000}" + +OUTDIR="${1:-${OUTDIR:-bench-results/correct-and-bench-mainnet-${BLOCKS}blocks-$(date +%F-%H%M%S)}}" +mkdir -p "$OUTDIR" +echo "output dir: $OUTDIR" + +for t in $THREADS; do + echo "running threads=$t ..." + args=( + --cache-warm-iters "$CACHE_WARM_ITERS" + --cache-warm-keys "$CACHE_WARM_KEYS" + --cache-warm-percentages "$CACHE_WARM_PERCENTAGES" + --cache-warm-out-csv "$OUTDIR/correct-and-bench-threads-$t.csv" + --full + --blocks "$BLOCKS" + --reth-bin "$RETH_BIN" + --datadir "$DATADIR" + --chain "$CHAIN" + ) + if [[ -n "$STATIC_FILE_BLOCKS_PER_FILE" ]]; then + args+=(--static-file-blocks-per-file "$STATIC_FILE_BLOCKS_PER_FILE") + fi + + if ! RAYON_NUM_THREADS="$t" "$BIN" "${args[@]}" \ + > "$OUTDIR/correct-and-bench-threads-$t.log" 2>&1; then + echo "failed for threads=$t" + echo "log: $OUTDIR/correct-and-bench-threads-$t.log" + tail -n 80 "$OUTDIR/correct-and-bench-threads-$t.log" || true + exit 1 + fi + echo "completed threads=$t" +done + +{ + echo "threads,iteration,v_experimental_root_correct,warm_pct,v2_p50_ms,v_experimental_p50_ms,v2_p99_ms,v_experimental_p99_ms,v2_mean_ms,v_experimental_mean_ms,speedup" + for t in $THREADS; do + tail -n +2 "$OUTDIR/correct-and-bench-threads-$t.csv" | sed "s/^/$t,/" + done +} > "$OUTDIR/correct-and-bench-thread-sweep-combined.csv" + +echo "done: $OUTDIR" diff --git a/crates/eth-sparse-mpt/src/bin/cache-warm-compare-v-experimental.rs b/crates/eth-sparse-mpt/src/bin/cache-warm-compare-v-experimental.rs new file mode 100644 index 000000000..e138c969a --- /dev/null +++ b/crates/eth-sparse-mpt/src/bin/cache-warm-compare-v-experimental.rs @@ -0,0 +1,495 @@ +use alloy_primitives::{keccak256, B256, U256}; +use eth_sparse_mpt::{ + v2::trie::{proof_store::ProofStore as ProofStoreV2, Trie as TrieV2}, + v_experimental::trie::{ + proof_store::ProofStore as ProofStoreVExperimental, Trie as TrieVExperimental, + }, +}; +use eyre::{bail, Context, Result}; +use std::{ + env, + fs::{create_dir_all, File, OpenOptions}, + io::{BufRead, BufReader, BufWriter, Write}, + path::{Path, PathBuf}, + time::{Duration, Instant}, +}; + +#[derive(Debug, Clone)] +struct Cli { + iters: usize, + keys: usize, + percentages: Vec, + out_csv: Option, + append_csv: bool, + csv_iteration: Option, + v_experimental_root_correct: Option, +} + +#[derive(Debug, Clone, Copy)] +struct Stats { + median_ms: f64, + p99_ms: f64, + mean_ms: f64, +} + +fn main() -> Result<()> { + let cli = parse_cli()?; + let (keys, values) = prepare_key_value_data(cli.keys); + let empty_proof_store_v2 = ProofStoreV2::default(); + let empty_proof_store_v_experimental = ProofStoreVExperimental::default(); + let include_correctness_columns = + cli.csv_iteration.is_some() || cli.v_experimental_root_correct.is_some(); + let mut csv_writer = open_csv_writer( + cli.out_csv.as_deref(), + cli.append_csv, + include_correctness_columns, + )?; + + println!( + "Running cache-warm comparison with iters={}, keys={}, percentages={:?}", + cli.iters, cli.keys, cli.percentages + ); + println!( + "baseline=v2::Trie(root_hash parallel=true), optimized=v_experimental::Trie(root_hash parallel=true)" + ); + println!(); + println!( + "{:<8} {:>12} {:>12} {:>12} {:>12} {:>12} {:>12} {:>12}", + "warm_%", + "v2_p50", + "v_experimental_p50", + "v2_p99", + "v_experimental_p99", + "v2_mean", + "v_experimental_mean", + "speedup" + ); + if let Some((writer, write_header, include_correctness_columns)) = csv_writer.as_mut() { + if *write_header { + if *include_correctness_columns { + writeln!( + writer, + "iteration,v_experimental_root_correct,warm_pct,v2_p50_ms,v_experimental_p50_ms,v2_p99_ms,v_experimental_p99_ms,v2_mean_ms,v_experimental_mean_ms,speedup" + ) + .context("failed to write CSV header")?; + } else { + writeln!( + writer, + "warm_pct,v2_p50_ms,v_experimental_p50_ms,v2_p99_ms,v_experimental_p99_ms,v2_mean_ms,v_experimental_mean_ms,speedup" + ) + .context("failed to write CSV header")?; + } + writer.flush().context("failed to flush CSV header")?; + } + } + + for pct in &cli.percentages { + let prewarm = keys.len() * *pct / 100; + + let v2_template = make_v2_template(&keys, &values, prewarm, &empty_proof_store_v2)?; + let v_experimental_template = make_v_experimental_template( + &keys, + &values, + prewarm, + &empty_proof_store_v_experimental, + )?; + + let mut v2_times = Vec::with_capacity(cli.iters); + let mut v_experimental_times = Vec::with_capacity(cli.iters); + + for _ in 0..cli.iters { + let (v2_hash, v2_elapsed) = run_v2_once( + v2_template.clone(), + &keys, + &values, + prewarm, + &empty_proof_store_v2, + )?; + let (v_experimental_hash, v_experimental_elapsed) = run_v_experimental_once( + v_experimental_template.clone(), + &keys, + &values, + prewarm, + &empty_proof_store_v_experimental, + )?; + + if v2_hash != v_experimental_hash { + bail!( + "root hash mismatch at warm {}%: v2={v2_hash:?} v_experimental={v_experimental_hash:?}", + pct + ); + } + + v2_times.push(v2_elapsed); + v_experimental_times.push(v_experimental_elapsed); + } + + let v2_stats = calc_stats(&v2_times)?; + let v_experimental_stats = calc_stats(&v_experimental_times)?; + let speedup = if v_experimental_stats.mean_ms > 0.0 { + v2_stats.mean_ms / v_experimental_stats.mean_ms + } else { + f64::INFINITY + }; + + println!( + "{:<8} {:>10.3}ms {:>10.3}ms {:>10.3}ms {:>10.3}ms {:>10.3}ms {:>10.3}ms {:>10.2}x", + pct, + v2_stats.median_ms, + v_experimental_stats.median_ms, + v2_stats.p99_ms, + v_experimental_stats.p99_ms, + v2_stats.mean_ms, + v_experimental_stats.mean_ms, + speedup + ); + + if let Some((writer, _, include_correctness_columns)) = csv_writer.as_mut() { + if *include_correctness_columns { + let iteration = cli.csv_iteration.map(|v| v.to_string()).unwrap_or_default(); + let v_experimental_root_correct = cli + .v_experimental_root_correct + .map(|v| if v { "true" } else { "false" }) + .unwrap_or(""); + writeln!( + writer, + "{},{},{},{:.6},{:.6},{:.6},{:.6},{:.6},{:.6},{:.6}", + iteration, + v_experimental_root_correct, + pct, + v2_stats.median_ms, + v_experimental_stats.median_ms, + v2_stats.p99_ms, + v_experimental_stats.p99_ms, + v2_stats.mean_ms, + v_experimental_stats.mean_ms, + speedup + ) + .context("failed to write CSV row")?; + } else { + writeln!( + writer, + "{},{:.6},{:.6},{:.6},{:.6},{:.6},{:.6},{:.6}", + pct, + v2_stats.median_ms, + v_experimental_stats.median_ms, + v2_stats.p99_ms, + v_experimental_stats.p99_ms, + v2_stats.mean_ms, + v_experimental_stats.mean_ms, + speedup + ) + .context("failed to write CSV row")?; + } + writer.flush().context("failed to flush CSV row")?; + } + } + + if let Some((writer, _, _)) = csv_writer.as_mut() { + writer.flush().context("failed to flush CSV writer")?; + } + if let Some(path) = &cli.out_csv { + println!("wrote cache-warm CSV: {}", path.display()); + } + + Ok(()) +} + +fn parse_cli() -> Result { + let mut iters = 20usize; + let mut keys = 10_000usize; + let mut percentages = vec![70usize, 80, 90, 100]; + let mut out_csv = None; + let mut append_csv = false; + let mut csv_iteration = None; + let mut v_experimental_root_correct = None; + + let mut args = env::args().skip(1); + while let Some(arg) = args.next() { + match arg.as_str() { + "--iters" => { + let v = args + .next() + .ok_or_else(|| eyre::eyre!("--iters requires a value"))?; + iters = v.parse().with_context(|| format!("invalid --iters: {v}"))?; + } + "--keys" => { + let v = args + .next() + .ok_or_else(|| eyre::eyre!("--keys requires a value"))?; + keys = v.parse().with_context(|| format!("invalid --keys: {v}"))?; + } + "--percentages" => { + let v = args + .next() + .ok_or_else(|| eyre::eyre!("--percentages requires a value"))?; + percentages = parse_percentages(&v)?; + } + "--out-csv" => { + let v = args + .next() + .ok_or_else(|| eyre::eyre!("--out-csv requires a value"))?; + out_csv = Some(expand_tilde(&v)); + } + "--append-csv" => { + append_csv = true; + } + "--csv-iteration" => { + let v = args + .next() + .ok_or_else(|| eyre::eyre!("--csv-iteration requires a value"))?; + let parsed: usize = v + .parse() + .with_context(|| format!("invalid --csv-iteration: {v}"))?; + csv_iteration = Some(parsed); + } + "--v-experimental-root-correct" => { + let v = args + .next() + .ok_or_else(|| eyre::eyre!("--v-experimental-root-correct requires a value"))?; + let parsed: bool = v + .parse() + .with_context(|| format!("invalid correctness flag value: {v}"))?; + v_experimental_root_correct = Some(parsed); + } + "--help" | "-h" => { + print_help(); + std::process::exit(0); + } + _ => bail!("unknown arg: {arg}"), + } + } + + if iters == 0 { + bail!("--iters must be > 0"); + } + if keys == 0 { + bail!("--keys must be > 0"); + } + if percentages.is_empty() { + bail!("--percentages must not be empty"); + } + for p in &percentages { + if *p > 100 { + bail!("percentage out of range [0,100]: {p}"); + } + } + + Ok(Cli { + iters, + keys, + percentages, + out_csv, + append_csv, + csv_iteration, + v_experimental_root_correct, + }) +} + +fn parse_percentages(input: &str) -> Result> { + let mut out = Vec::new(); + for p in input.split(',') { + let p = p.trim(); + if p.is_empty() { + continue; + } + out.push( + p.parse() + .with_context(|| format!("invalid percentage: {p}"))?, + ); + } + Ok(out) +} + +fn print_help() { + println!("cache-warm-compare-v-experimental"); + println!(" --iters iterations per percentage (default: 20)"); + println!(" --keys number of key/value pairs (default: 10000)"); + println!(" --percentages pre-warm percentages (default: 70,80,90,100)"); + println!(" --out-csv write results to CSV"); + println!(" --append-csv append rows to CSV instead of truncating"); + println!(" --csv-iteration iteration id to include in CSV rows"); + println!(" --v-experimental-root-correct include v_experimental correctness (true/false) in CSV rows"); +} + +fn open_csv_writer( + path: Option<&Path>, + append: bool, + include_correctness_columns: bool, +) -> Result, bool, bool)>> { + let Some(path) = path else { + return Ok(None); + }; + + if let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) { + create_dir_all(parent).with_context(|| { + format!( + "failed to create parent directory for CSV: {}", + parent.display() + ) + })?; + } + + if append { + let existed = path.exists(); + let existing_len = if existed { + Some( + path.metadata() + .with_context(|| format!("failed to stat CSV file {}", path.display()))? + .len(), + ) + } else { + None + }; + let mut effective_include_correctness_columns = include_correctness_columns; + if let Some(len) = existing_len { + if len > 0 { + let mut first_line = String::new(); + BufReader::new( + File::open(path) + .with_context(|| format!("failed to read CSV file {}", path.display()))?, + ) + .read_line(&mut first_line) + .with_context(|| format!("failed to read CSV header {}", path.display()))?; + let header_has_correctness_column = first_line + .trim_end() + .split(',') + .any(|h| h == "v_experimental_root_correct"); + if include_correctness_columns && !header_has_correctness_column { + bail!( + "existing CSV {} does not include v_experimental_root_correct column; use a new --out-csv path", + path.display() + ); + } + effective_include_correctness_columns = header_has_correctness_column; + } + } + + let file = OpenOptions::new() + .create(true) + .append(true) + .open(path) + .with_context(|| format!("failed to open CSV file {}", path.display()))?; + let write_header = !existed || existing_len.is_some_and(|len| len == 0); + Ok(Some(( + BufWriter::new(file), + write_header, + effective_include_correctness_columns, + ))) + } else { + let file = File::create(path) + .with_context(|| format!("failed to create CSV file {}", path.display()))?; + Ok(Some(( + BufWriter::new(file), + true, + include_correctness_columns, + ))) + } +} + +fn expand_tilde(path: &str) -> PathBuf { + if let Some(rest) = path.strip_prefix("~/") { + if let Ok(home) = env::var("HOME") { + return PathBuf::from(home).join(rest); + } + } + PathBuf::from(path) +} + +fn prepare_key_value_data(n: usize) -> (Vec>, Vec>) { + let mut keys = Vec::with_capacity(n); + let mut values = Vec::with_capacity(n); + for i in 0u64..(n as u64) { + let b: B256 = U256::from(i).into(); + let key = keccak256(b).to_vec(); + let value = keccak256(&key).to_vec(); + keys.push(key); + values.push(value); + } + (keys, values) +} + +fn make_v2_template( + keys: &[Vec], + values: &[Vec], + prewarm: usize, + proof_store: &ProofStoreV2, +) -> Result { + let mut trie = TrieV2::new_empty(); + for i in 0..prewarm { + trie.insert(&keys[i], &values[i])?; + } + if prewarm > 0 { + let _ = trie.root_hash(true, proof_store)?; + } + Ok(trie) +} + +fn make_v_experimental_template( + keys: &[Vec], + values: &[Vec], + prewarm: usize, + proof_store: &ProofStoreVExperimental, +) -> Result { + let mut trie = TrieVExperimental::new_empty(); + for i in 0..prewarm { + trie.insert(&keys[i], &values[i])?; + } + if prewarm > 0 { + let _ = trie.root_hash(true, proof_store)?; + } + Ok(trie) +} + +fn run_v2_once( + mut trie: TrieV2, + keys: &[Vec], + values: &[Vec], + start: usize, + proof_store: &ProofStoreV2, +) -> Result<(B256, Duration)> { + let start_time = Instant::now(); + for i in start..keys.len() { + trie.insert(&keys[i], &values[i])?; + } + let hash = trie.root_hash(true, proof_store)?; + Ok((hash, start_time.elapsed())) +} + +fn run_v_experimental_once( + mut trie: TrieVExperimental, + keys: &[Vec], + values: &[Vec], + start: usize, + proof_store: &ProofStoreVExperimental, +) -> Result<(B256, Duration)> { + let start_time = Instant::now(); + for i in start..keys.len() { + trie.insert(&keys[i], &values[i])?; + } + let hash = trie.root_hash(true, proof_store)?; + Ok((hash, start_time.elapsed())) +} + +fn calc_stats(data: &[Duration]) -> Result { + if data.is_empty() { + bail!("no data points"); + } + let mut ms: Vec = data.iter().map(|d| d.as_secs_f64() * 1000.0).collect(); + ms.sort_by(|a, b| a.total_cmp(b)); + + let mean_ms = ms.iter().sum::() / ms.len() as f64; + let median_ms = if ms.len() % 2 == 0 { + (ms[ms.len() / 2 - 1] + ms[ms.len() / 2]) / 2.0 + } else { + ms[ms.len() / 2] + }; + let p99_idx = (ms.len() * 99 / 100).min(ms.len() - 1); + let p99_ms = ms[p99_idx]; + + Ok(Stats { + median_ms, + p99_ms, + mean_ms, + }) +} diff --git a/crates/eth-sparse-mpt/src/bin/cache-warm-compare.rs b/crates/eth-sparse-mpt/src/bin/cache-warm-compare.rs new file mode 100644 index 000000000..d7f869102 --- /dev/null +++ b/crates/eth-sparse-mpt/src/bin/cache-warm-compare.rs @@ -0,0 +1,502 @@ +use alloy_primitives::{keccak256, B256, U256}; +use eth_sparse_mpt::{ + v2::trie::{proof_store::ProofStore as ProofStoreV2, Trie as TrieV2}, + v_experimental::trie::{ + proof_store::ProofStore as ProofStoreVExperimental, Trie as TrieVExperimental, + }, +}; +use eyre::{bail, Context, Result}; +use std::{ + env, + fs::{create_dir_all, File, OpenOptions}, + io::{BufRead, BufReader, BufWriter, Write}, + path::{Path, PathBuf}, + time::{Duration, Instant}, +}; + +#[derive(Debug, Clone)] +struct Cli { + iters: usize, + keys: usize, + percentages: Vec, + out_csv: Option, + append_csv: bool, + csv_iteration: Option, + v_experimental_root_correct: Option, +} + +#[derive(Debug, Clone, Copy)] +struct Stats { + median_ms: f64, + p99_ms: f64, + mean_ms: f64, +} + +fn main() -> Result<()> { + let cli = parse_cli()?; + let (keys, values) = prepare_key_value_data(cli.keys); + let empty_proof_store_v2 = ProofStoreV2::default(); + let empty_proof_store_v_experimental = ProofStoreVExperimental::default(); + let include_correctness_columns = + cli.csv_iteration.is_some() || cli.v_experimental_root_correct.is_some(); + let mut csv_writer = open_csv_writer( + cli.out_csv.as_deref(), + cli.append_csv, + include_correctness_columns, + )?; + + println!( + "Running cache-warm comparison with iters={}, keys={}, percentages={:?}", + cli.iters, cli.keys, cli.percentages + ); + println!( + "baseline=v2::Trie(root_hash parallel=true), optimized=v_experimental::Trie(root_hash parallel=true)" + ); + println!(); + println!( + "{:<8} {:>12} {:>12} {:>12} {:>12} {:>12} {:>12} {:>12}", + "warm_%", + "v2_p50", + "v_experimental_p50", + "v2_p99", + "v_experimental_p99", + "v2_mean", + "v_experimental_mean", + "speedup" + ); + if let Some((writer, write_header, include_correctness_columns)) = csv_writer.as_mut() { + if *write_header { + if *include_correctness_columns { + writeln!( + writer, + "iteration,v_experimental_root_correct,warm_pct,v2_p50_ms,v_experimental_p50_ms,v2_p99_ms,v_experimental_p99_ms,v2_mean_ms,v_experimental_mean_ms,speedup" + ) + .context("failed to write CSV header")?; + } else { + writeln!( + writer, + "warm_pct,v2_p50_ms,v_experimental_p50_ms,v2_p99_ms,v_experimental_p99_ms,v2_mean_ms,v_experimental_mean_ms,speedup" + ) + .context("failed to write CSV header")?; + } + writer.flush().context("failed to flush CSV header")?; + } + } + + for pct in &cli.percentages { + let prewarm = keys.len() * *pct / 100; + + let v2_template = make_v2_template(&keys, &values, prewarm, &empty_proof_store_v2)?; + let v_experimental_template = make_v_experimental_template( + &keys, + &values, + prewarm, + &empty_proof_store_v_experimental, + )?; + + let mut v2_times = Vec::with_capacity(cli.iters); + let mut v_experimental_times = Vec::with_capacity(cli.iters); + + for _ in 0..cli.iters { + let (v2_hash, v2_elapsed) = run_v2_once( + v2_template.clone(), + &keys, + &values, + prewarm, + &empty_proof_store_v2, + )?; + let (v_experimental_hash, v_experimental_elapsed) = run_v_experimental_once( + v_experimental_template.clone(), + &keys, + &values, + prewarm, + &empty_proof_store_v_experimental, + )?; + + if v2_hash != v_experimental_hash { + bail!( + "root hash mismatch at warm {}%: v2={v2_hash:?} v_experimental={v_experimental_hash:?}", + pct + ); + } + + v2_times.push(v2_elapsed); + v_experimental_times.push(v_experimental_elapsed); + } + + let v2_stats = calc_stats(&v2_times)?; + let v_experimental_stats = calc_stats(&v_experimental_times)?; + let speedup = if v_experimental_stats.mean_ms > 0.0 { + v2_stats.mean_ms / v_experimental_stats.mean_ms + } else { + f64::INFINITY + }; + + println!( + "{:<8} {:>10.3}ms {:>10.3}ms {:>10.3}ms {:>10.3}ms {:>10.3}ms {:>10.3}ms {:>10.2}x", + pct, + v2_stats.median_ms, + v_experimental_stats.median_ms, + v2_stats.p99_ms, + v_experimental_stats.p99_ms, + v2_stats.mean_ms, + v_experimental_stats.mean_ms, + speedup + ); + + if let Some((writer, _, include_correctness_columns)) = csv_writer.as_mut() { + if *include_correctness_columns { + let iteration = cli.csv_iteration.map(|v| v.to_string()).unwrap_or_default(); + let v_experimental_root_correct = cli + .v_experimental_root_correct + .map(|v| if v { "true" } else { "false" }) + .unwrap_or(""); + writeln!( + writer, + "{},{},{},{:.6},{:.6},{:.6},{:.6},{:.6},{:.6},{:.6}", + iteration, + v_experimental_root_correct, + pct, + v2_stats.median_ms, + v_experimental_stats.median_ms, + v2_stats.p99_ms, + v_experimental_stats.p99_ms, + v2_stats.mean_ms, + v_experimental_stats.mean_ms, + speedup + ) + .context("failed to write CSV row")?; + } else { + writeln!( + writer, + "{},{:.6},{:.6},{:.6},{:.6},{:.6},{:.6},{:.6}", + pct, + v2_stats.median_ms, + v_experimental_stats.median_ms, + v2_stats.p99_ms, + v_experimental_stats.p99_ms, + v2_stats.mean_ms, + v_experimental_stats.mean_ms, + speedup + ) + .context("failed to write CSV row")?; + } + writer.flush().context("failed to flush CSV row")?; + } + } + + if let Some((writer, _, _)) = csv_writer.as_mut() { + writer.flush().context("failed to flush CSV writer")?; + } + if let Some(path) = &cli.out_csv { + println!("wrote cache-warm CSV: {}", path.display()); + } + + Ok(()) +} + +fn parse_cli() -> Result { + let mut iters = 20usize; + let mut keys = 10_000usize; + let mut percentages = vec![70usize, 80, 90, 100]; + let mut out_csv = None; + let mut append_csv = false; + let mut csv_iteration = None; + let mut v_experimental_root_correct = None; + + let mut args = env::args().skip(1); + while let Some(arg) = args.next() { + match arg.as_str() { + "--iters" => { + let v = args + .next() + .ok_or_else(|| eyre::eyre!("--iters requires a value"))?; + iters = v.parse().with_context(|| format!("invalid --iters: {v}"))?; + } + "--keys" => { + let v = args + .next() + .ok_or_else(|| eyre::eyre!("--keys requires a value"))?; + keys = v.parse().with_context(|| format!("invalid --keys: {v}"))?; + } + "--percentages" => { + let v = args + .next() + .ok_or_else(|| eyre::eyre!("--percentages requires a value"))?; + percentages = parse_percentages(&v)?; + } + "--out-csv" => { + let v = args + .next() + .ok_or_else(|| eyre::eyre!("--out-csv requires a value"))?; + out_csv = Some(expand_tilde(&v)); + } + "--append-csv" => { + append_csv = true; + } + "--csv-iteration" => { + let v = args + .next() + .ok_or_else(|| eyre::eyre!("--csv-iteration requires a value"))?; + let parsed: usize = v + .parse() + .with_context(|| format!("invalid --csv-iteration: {v}"))?; + csv_iteration = Some(parsed); + } + "--v-experimental-root-correct" | "--v3-root-correct" => { + let v = args + .next() + .ok_or_else(|| { + eyre::eyre!( + "--v-experimental-root-correct requires a value (or use legacy --v3-root-correct)" + ) + })?; + let parsed: bool = v + .parse() + .with_context(|| format!("invalid correctness flag value: {v}"))?; + v_experimental_root_correct = Some(parsed); + } + "--help" | "-h" => { + print_help(); + std::process::exit(0); + } + _ => bail!("unknown arg: {arg}"), + } + } + + if iters == 0 { + bail!("--iters must be > 0"); + } + if keys == 0 { + bail!("--keys must be > 0"); + } + if percentages.is_empty() { + bail!("--percentages must not be empty"); + } + for p in &percentages { + if *p > 100 { + bail!("percentage out of range [0,100]: {p}"); + } + } + + Ok(Cli { + iters, + keys, + percentages, + out_csv, + append_csv, + csv_iteration, + v_experimental_root_correct, + }) +} + +fn parse_percentages(input: &str) -> Result> { + let mut out = Vec::new(); + for p in input.split(',') { + let p = p.trim(); + if p.is_empty() { + continue; + } + out.push( + p.parse() + .with_context(|| format!("invalid percentage: {p}"))?, + ); + } + Ok(out) +} + +fn print_help() { + println!("cache-warm-compare"); + println!(" --iters iterations per percentage (default: 20)"); + println!(" --keys number of key/value pairs (default: 10000)"); + println!(" --percentages pre-warm percentages (default: 70,80,90,100)"); + println!(" --out-csv write results to CSV"); + println!(" --append-csv append rows to CSV instead of truncating"); + println!(" --csv-iteration iteration id to include in CSV rows"); + println!( + " --v-experimental-root-correct include v_experimental correctness (true/false) in CSV rows" + ); + println!(" --v3-root-correct legacy alias for --v-experimental-root-correct"); +} + +fn open_csv_writer( + path: Option<&Path>, + append: bool, + include_correctness_columns: bool, +) -> Result, bool, bool)>> { + let Some(path) = path else { + return Ok(None); + }; + + if let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) { + create_dir_all(parent).with_context(|| { + format!( + "failed to create parent directory for CSV: {}", + parent.display() + ) + })?; + } + + if append { + let existed = path.exists(); + let existing_len = if existed { + Some( + path.metadata() + .with_context(|| format!("failed to stat CSV file {}", path.display()))? + .len(), + ) + } else { + None + }; + let mut effective_include_correctness_columns = include_correctness_columns; + if let Some(len) = existing_len { + if len > 0 { + let mut first_line = String::new(); + BufReader::new( + File::open(path) + .with_context(|| format!("failed to read CSV file {}", path.display()))?, + ) + .read_line(&mut first_line) + .with_context(|| format!("failed to read CSV header {}", path.display()))?; + let header_has_correctness_column = first_line + .trim_end() + .split(',') + .any(|h| h == "v_experimental_root_correct" || h == "v3_root_correct"); + if include_correctness_columns && !header_has_correctness_column { + bail!( + "existing CSV {} does not include v_experimental_root_correct (or legacy v3_root_correct) column; use a new --out-csv path", + path.display() + ); + } + effective_include_correctness_columns = header_has_correctness_column; + } + } + + let file = OpenOptions::new() + .create(true) + .append(true) + .open(path) + .with_context(|| format!("failed to open CSV file {}", path.display()))?; + let write_header = !existed || existing_len.is_some_and(|len| len == 0); + Ok(Some(( + BufWriter::new(file), + write_header, + effective_include_correctness_columns, + ))) + } else { + let file = File::create(path) + .with_context(|| format!("failed to create CSV file {}", path.display()))?; + Ok(Some(( + BufWriter::new(file), + true, + include_correctness_columns, + ))) + } +} + +fn expand_tilde(path: &str) -> PathBuf { + if let Some(rest) = path.strip_prefix("~/") { + if let Ok(home) = env::var("HOME") { + return PathBuf::from(home).join(rest); + } + } + PathBuf::from(path) +} + +fn prepare_key_value_data(n: usize) -> (Vec>, Vec>) { + let mut keys = Vec::with_capacity(n); + let mut values = Vec::with_capacity(n); + for i in 0u64..(n as u64) { + let b: B256 = U256::from(i).into(); + let key = keccak256(b).to_vec(); + let value = keccak256(&key).to_vec(); + keys.push(key); + values.push(value); + } + (keys, values) +} + +fn make_v2_template( + keys: &[Vec], + values: &[Vec], + prewarm: usize, + proof_store: &ProofStoreV2, +) -> Result { + let mut trie = TrieV2::new_empty(); + for i in 0..prewarm { + trie.insert(&keys[i], &values[i])?; + } + if prewarm > 0 { + let _ = trie.root_hash(true, proof_store)?; + } + Ok(trie) +} + +fn make_v_experimental_template( + keys: &[Vec], + values: &[Vec], + prewarm: usize, + proof_store: &ProofStoreVExperimental, +) -> Result { + let mut trie = TrieVExperimental::new_empty(); + for i in 0..prewarm { + trie.insert(&keys[i], &values[i])?; + } + if prewarm > 0 { + let _ = trie.root_hash(true, proof_store)?; + } + Ok(trie) +} + +fn run_v2_once( + mut trie: TrieV2, + keys: &[Vec], + values: &[Vec], + start: usize, + proof_store: &ProofStoreV2, +) -> Result<(B256, Duration)> { + let start_time = Instant::now(); + for i in start..keys.len() { + trie.insert(&keys[i], &values[i])?; + } + let hash = trie.root_hash(true, proof_store)?; + Ok((hash, start_time.elapsed())) +} + +fn run_v_experimental_once( + mut trie: TrieVExperimental, + keys: &[Vec], + values: &[Vec], + start: usize, + proof_store: &ProofStoreVExperimental, +) -> Result<(B256, Duration)> { + let start_time = Instant::now(); + for i in start..keys.len() { + trie.insert(&keys[i], &values[i])?; + } + let hash = trie.root_hash(true, proof_store)?; + Ok((hash, start_time.elapsed())) +} + +fn calc_stats(data: &[Duration]) -> Result { + if data.is_empty() { + bail!("no data points"); + } + let mut ms: Vec = data.iter().map(|d| d.as_secs_f64() * 1000.0).collect(); + ms.sort_by(|a, b| a.total_cmp(b)); + + let mean_ms = ms.iter().sum::() / ms.len() as f64; + let median_ms = if ms.len() % 2 == 0 { + (ms[ms.len() / 2 - 1] + ms[ms.len() / 2]) / 2.0 + } else { + ms[ms.len() / 2] + }; + let p99_idx = (ms.len() * 99 / 100).min(ms.len() - 1); + let p99_ms = ms[p99_idx]; + + Ok(Stats { + median_ms, + p99_ms, + mean_ms, + }) +} diff --git a/crates/eth-sparse-mpt/src/bin/correct_and_bench.rs b/crates/eth-sparse-mpt/src/bin/correct_and_bench.rs new file mode 100644 index 000000000..dbbd2f971 --- /dev/null +++ b/crates/eth-sparse-mpt/src/bin/correct_and_bench.rs @@ -0,0 +1,400 @@ +use eyre::{bail, Context, Result}; +use std::{ + env, + path::{Path, PathBuf}, + process::Command, +}; + +#[derive(Debug, Clone)] +struct Cli { + correctness_bin: PathBuf, + cache_warm_bin: PathBuf, + cache_warm_iters: usize, + cache_warm_keys: usize, + cache_warm_percentages: String, + cache_warm_out_csv: Option, + correctness_args: Vec, +} + +#[derive(Debug, Clone)] +struct FullMode { + blocks: usize, + reth_bin: PathBuf, + datadir: PathBuf, + chain: String, +} + +fn main() -> Result<()> { + let cli = parse_cli()?; + + if let Some(full_mode) = parse_full_mode(&cli.correctness_args)? { + run_full_interleaved(&cli, &full_mode)?; + } else { + run_correctness(&cli)?; + run_cache_warm(&cli, false, Some(1), true)?; + } + + Ok(()) +} + +fn run_correctness(cli: &Cli) -> Result<()> { + run_correctness_with_args(cli, &cli.correctness_args) +} + +fn run_correctness_with_args(cli: &Cli, args: &[String]) -> Result<()> { + println!( + "running correctness harness: {} {}", + cli.correctness_bin.display(), + args.join(" ") + ); + + let status = Command::new(&cli.correctness_bin) + .args(args) + .status() + .with_context(|| { + format!( + "failed to start correctness harness at {}; build binaries with `cargo build --release --bins` or pass --correctness-bin", + cli.correctness_bin.display() + ) + })?; + + if !status.success() { + bail!("correctness harness failed with status {status}"); + } + + Ok(()) +} + +fn run_cache_warm( + cli: &Cli, + append_csv: bool, + csv_iteration: Option, + v_experimental_root_correct: bool, +) -> Result<()> { + println!(); + let mut summary = format!( + "running cache-warm benchmark: {} --iters {} --keys {} --percentages {}", + cli.cache_warm_bin.display(), + cli.cache_warm_iters, + cli.cache_warm_keys, + cli.cache_warm_percentages + ); + if let Some(out_csv) = &cli.cache_warm_out_csv { + summary.push_str(&format!(" --out-csv {}", out_csv.display())); + if append_csv { + summary.push_str(" --append-csv"); + } + if let Some(iteration) = csv_iteration { + summary.push_str(&format!(" --csv-iteration {iteration}")); + summary.push_str(&format!( + " --v-experimental-root-correct {}", + v_experimental_root_correct + )); + } + } + println!("{summary}"); + + let mut cmd = Command::new(&cli.cache_warm_bin); + cmd.arg("--iters") + .arg(cli.cache_warm_iters.to_string()) + .arg("--keys") + .arg(cli.cache_warm_keys.to_string()) + .arg("--percentages") + .arg(&cli.cache_warm_percentages); + if let Some(out_csv) = &cli.cache_warm_out_csv { + cmd.arg("--out-csv").arg(out_csv); + if append_csv { + cmd.arg("--append-csv"); + } + if let Some(iteration) = csv_iteration { + cmd.arg("--csv-iteration").arg(iteration.to_string()); + cmd.arg("--v-experimental-root-correct") + .arg(if v_experimental_root_correct { + "true" + } else { + "false" + }); + } + } + + let status = cmd + .status() + .with_context(|| { + format!( + "failed to start cache-warm benchmark at {}; build binaries with `cargo build --release --bins` or pass --cache-warm-bin", + cli.cache_warm_bin.display() + ) + })?; + + if !status.success() { + bail!("cache-warm benchmark failed with status {status}"); + } + + Ok(()) +} + +fn run_full_interleaved(cli: &Cli, full_mode: &FullMode) -> Result<()> { + let single_iteration_args = rewrite_to_single_full_iteration(&cli.correctness_args)?; + let append_csv = full_mode.blocks > 1 && cli.cache_warm_out_csv.is_some(); + + println!( + "running full interleaved mode: blocks={} (correctness -> cache-warm -> unwind)", + full_mode.blocks + ); + + for i in 0..full_mode.blocks { + println!(); + println!("interleaved iteration {}/{}", i + 1, full_mode.blocks); + run_correctness_with_args(cli, &single_iteration_args)?; + run_cache_warm(cli, append_csv, Some(i + 1), true)?; + if i + 1 < full_mode.blocks { + run_unwind_once(full_mode, i + 1)?; + } + } + + println!( + "full interleaved run passed for {} iterations", + full_mode.blocks + ); + Ok(()) +} + +fn run_unwind_once(full_mode: &FullMode, iteration: usize) -> Result<()> { + let output = Command::new(&full_mode.reth_bin) + .arg("stage") + .arg("unwind") + .arg(format!("--datadir={}", full_mode.datadir.display())) + .arg(format!("--chain={}", full_mode.chain)) + .arg("num-blocks") + .arg("1") + .output() + .with_context(|| { + format!("failed to execute reth unwind command at iteration {iteration}") + })?; + + if !output.status.success() { + bail!( + "reth unwind failed at iteration {} (status: {}):\nstdout:\n{}\nstderr:\n{}", + iteration, + output.status, + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + } + Ok(()) +} + +fn parse_full_mode(correctness_args: &[String]) -> Result> { + let mut full = false; + let mut blocks = 1usize; + let mut reth_bin = expand_tilde("~/reth/target/maxperf/reth"); + let mut datadir = expand_tilde("/home/ubuntu/merkle/"); + let mut chain = "mainnet".to_string(); + + let mut i = 0usize; + while i < correctness_args.len() { + match correctness_args[i].as_str() { + "--full" => { + full = true; + i += 1; + } + "--blocks" => { + let value = correctness_args + .get(i + 1) + .ok_or_else(|| eyre::eyre!("--blocks requires a value"))?; + blocks = value + .parse() + .with_context(|| format!("invalid --blocks value: {value}"))?; + i += 2; + } + "--reth-bin" => { + let value = correctness_args + .get(i + 1) + .ok_or_else(|| eyre::eyre!("--reth-bin requires a value"))?; + reth_bin = expand_tilde(value); + i += 2; + } + "--datadir" => { + let value = correctness_args + .get(i + 1) + .ok_or_else(|| eyre::eyre!("--datadir requires a value"))?; + datadir = expand_tilde(value); + i += 2; + } + "--chain" => { + let value = correctness_args + .get(i + 1) + .ok_or_else(|| eyre::eyre!("--chain requires a value"))?; + chain = value.clone(); + i += 2; + } + _ => { + i += 1; + } + } + } + + if !full { + return Ok(None); + } + if blocks == 0 { + bail!("--blocks must be > 0"); + } + + Ok(Some(FullMode { + blocks, + reth_bin, + datadir, + chain, + })) +} + +fn rewrite_to_single_full_iteration(correctness_args: &[String]) -> Result> { + let mut out = Vec::with_capacity(correctness_args.len() + 2); + let mut saw_full = false; + let mut i = 0usize; + + while i < correctness_args.len() { + match correctness_args[i].as_str() { + "--full" => { + saw_full = true; + out.push(correctness_args[i].clone()); + i += 1; + } + "--blocks" => { + if i + 1 >= correctness_args.len() { + bail!("--blocks requires a value"); + } + i += 2; + } + _ => { + out.push(correctness_args[i].clone()); + i += 1; + } + } + } + + if !saw_full { + out.push("--full".to_string()); + } + out.push("--blocks".to_string()); + out.push("1".to_string()); + Ok(out) +} + +fn parse_cli() -> Result { + let mut correctness_bin = sibling_binary("correctness-harness"); + let mut cache_warm_bin = sibling_binary("cache-warm-compare"); + let mut cache_warm_iters = 20usize; + let mut cache_warm_keys = 10_000usize; + let mut cache_warm_percentages = "70,80,90,100".to_string(); + let mut cache_warm_out_csv = None; + let mut correctness_args = Vec::new(); + + let mut args = env::args().skip(1); + while let Some(arg) = args.next() { + match arg.as_str() { + "--correctness-bin" => { + let value = args + .next() + .ok_or_else(|| eyre::eyre!("--correctness-bin requires a value"))?; + correctness_bin = expand_tilde(&value); + } + "--cache-warm-bin" => { + let value = args + .next() + .ok_or_else(|| eyre::eyre!("--cache-warm-bin requires a value"))?; + cache_warm_bin = expand_tilde(&value); + } + "--cache-warm-iters" => { + let value = args + .next() + .ok_or_else(|| eyre::eyre!("--cache-warm-iters requires a value"))?; + cache_warm_iters = value + .parse() + .with_context(|| format!("invalid --cache-warm-iters value: {value}"))?; + } + "--cache-warm-keys" => { + let value = args + .next() + .ok_or_else(|| eyre::eyre!("--cache-warm-keys requires a value"))?; + cache_warm_keys = value + .parse() + .with_context(|| format!("invalid --cache-warm-keys value: {value}"))?; + } + "--cache-warm-percentages" => { + cache_warm_percentages = args + .next() + .ok_or_else(|| eyre::eyre!("--cache-warm-percentages requires a value"))?; + } + "--cache-warm-out-csv" => { + let value = args + .next() + .ok_or_else(|| eyre::eyre!("--cache-warm-out-csv requires a value"))?; + cache_warm_out_csv = Some(expand_tilde(&value)); + } + "--help" | "-h" => { + print_help(); + std::process::exit(0); + } + _ => correctness_args.push(arg), + } + } + + if cache_warm_iters == 0 { + bail!("--cache-warm-iters must be > 0"); + } + if cache_warm_keys == 0 { + bail!("--cache-warm-keys must be > 0"); + } + + Ok(Cli { + correctness_bin, + cache_warm_bin, + cache_warm_iters, + cache_warm_keys, + cache_warm_percentages, + cache_warm_out_csv, + correctness_args, + }) +} + +fn print_help() { + println!("correct_and_bench"); + println!(" Runs correctness-harness and cache-warm-compare."); + println!(" In --full mode, runs cache-warm after each correctness iteration."); + println!(); + println!(" wrapper flags:"); + println!(" --correctness-bin correctness-harness binary path"); + println!(" --cache-warm-bin cache-warm-compare binary path"); + println!(" --cache-warm-iters benchmark iterations (default: 20)"); + println!(" --cache-warm-keys benchmark keys (default: 10000)"); + println!(" --cache-warm-percentages "); + println!(" benchmark warm percentages (default: 70,80,90,100)"); + println!(" --cache-warm-out-csv write cache-warm result rows to CSV"); + println!(); + println!(" all other flags are passed through to correctness-harness."); + println!(); + println!(" example:"); + println!( + " correct_and_bench --datadir /mnt/reth-archive --chain mainnet --tip-fallback --cache-warm-iters 10" + ); +} + +fn sibling_binary(name: &str) -> PathBuf { + let filename = format!("{name}{}", env::consts::EXE_SUFFIX); + env::current_exe() + .ok() + .and_then(|exe| exe.parent().map(Path::to_path_buf)) + .map(|dir| dir.join(&filename)) + .unwrap_or_else(|| PathBuf::from(filename)) +} + +fn expand_tilde(path: &str) -> PathBuf { + if let Some(rest) = path.strip_prefix("~/") { + if let Ok(home) = env::var("HOME") { + return PathBuf::from(home).join(rest); + } + } + PathBuf::from(path) +} diff --git a/crates/eth-sparse-mpt/src/bin/correctness-harness.rs b/crates/eth-sparse-mpt/src/bin/correctness-harness.rs new file mode 100644 index 000000000..203211b76 --- /dev/null +++ b/crates/eth-sparse-mpt/src/bin/correctness-harness.rs @@ -0,0 +1,647 @@ +use alloy_primitives::B256; +use eth_sparse_mpt::{ + calculate_root_hash_with_sparse_trie, ETHSpareMPTVersion, SparseTrieLocalCache, + SparseTrieSharedCache, +}; +use eyre::{bail, Context, Result}; +use reth_chainspec::{ChainSpec, DEV, HOLESKY, HOODI, MAINNET, SEPOLIA}; +use reth_db::{ + mdbx::{DatabaseArguments, MaxReadTransactionDuration}, + open_db_read_only, ClientVersion, DatabaseEnv, +}; +use reth_evm::{execute::Executor, ConfigureEvm}; +use reth_evm_ethereum::EthEvmConfig; +use reth_node_api::NodeTypesWithDBAdapter; +use reth_node_ethereum::EthereumNode; +use reth_provider::{ + providers::{ConsistentDbView, OverlayStateProviderFactory, StaticFileProvider}, + BlockHashReader, BlockNumReader, BlockReader, ChainSpecProvider, HeaderProvider, + ProviderFactory, TransactionVariant, +}; +use reth_revm::database::StateProviderDatabase; +use reth_trie::TrieInput; +use reth_trie_parallel::root::ParallelStateRoot; +use revm::database::BundleState; +use std::{ + env, + path::{Path, PathBuf}, + process::Command, + sync::Arc, + time::Instant, +}; + +type NodeTypes = NodeTypesWithDBAdapter>; +type Factory = ProviderFactory; + +#[derive(Debug, Clone)] +struct Cli { + full: bool, + blocks: usize, + reth_bin: PathBuf, + datadir: PathBuf, + chain: String, + tip_fallback: bool, + static_file_blocks_per_file: Option, + skip_v2: bool, + skip_v1: bool, +} + +#[derive(Debug, Clone, Copy)] +enum CompareMode { + TipPlusOne, + TipFallback, + CurrentTip, +} + +#[derive(Debug)] +struct CompareResult { + mode: CompareMode, + parent_block: u64, + target_block: u64, + canonical_root: B256, + reth_root: B256, + v2_root: Option, + v_experimental_root: B256, + v1_root: Option, + execute_ms: f64, + reth_ms: f64, + v2_ms: Option, + v_experimental_ms: f64, + v1_ms: Option, + v2_fetched_nodes: Option, + v_experimental_fetched_nodes: usize, + v1_fetched_nodes: Option, +} + +fn main() -> Result<()> { + let cli = parse_cli()?; + + if cli.full { + run_full_mode(&cli)?; + } else { + if cli.blocks != 1 { + bail!("--blocks is only supported with --full"); + } + let factory = + open_provider_factory(&cli.datadir, &cli.chain, cli.static_file_blocks_per_file)?; + let result = compare_tip_plus_one(&factory, &cli)?; + print_compare_result(None, &result); + println!("correctness check passed"); + } + Ok(()) +} + +fn run_full_mode(cli: &Cli) -> Result<()> { + println!( + "running full mode: blocks={} (compare current tip, then unwind-by-1 each iteration)", + cli.blocks + ); + + for i in 0..cli.blocks { + let factory = + open_provider_factory(&cli.datadir, &cli.chain, cli.static_file_blocks_per_file)?; + let result = compare_current_tip(&factory, cli)?; + print_compare_result(Some(i + 1), &result); + if i + 1 < cli.blocks { + run_unwind_once(cli, i + 1)?; + } + } + + println!("full correctness loop passed for {} iterations", cli.blocks); + Ok(()) +} + +fn compare_tip_plus_one(factory: &Factory, cli: &Cli) -> Result { + let parent_block = factory + .best_block_number() + .context("failed to read best (executed) block number")?; + let target_block = parent_block + .checked_add(1) + .ok_or_else(|| eyre::eyre!("tip overflow"))?; + let tip_plus_one_exists = factory + .header_by_number(target_block) + .with_context(|| format!("failed to fetch target header for block {target_block}"))? + .is_some(); + if tip_plus_one_exists { + compare_block( + factory, + parent_block, + target_block, + CompareMode::TipPlusOne, + cli.skip_v2, + cli.skip_v1, + ) + } else if cli.tip_fallback { + let fallback_parent = parent_block + .checked_sub(1) + .ok_or_else(|| eyre::eyre!("cannot fallback to tip mode at block 0"))?; + compare_block( + factory, + fallback_parent, + parent_block, + CompareMode::TipFallback, + cli.skip_v2, + cli.skip_v1, + ) + } else { + bail!( + "target header block {} not found in local DB; this harness expects tip+1 to exist after unwind", + target_block + ); + } +} + +fn compare_current_tip(factory: &Factory, cli: &Cli) -> Result { + let target_block = factory + .best_block_number() + .context("failed to read best (executed) block number")?; + let parent_block = target_block + .checked_sub(1) + .ok_or_else(|| eyre::eyre!("tip is block 0; need at least one executed block in DB"))?; + compare_block( + factory, + parent_block, + target_block, + CompareMode::CurrentTip, + cli.skip_v2, + cli.skip_v1, + ) +} + +fn compare_block( + factory: &Factory, + parent_block: u64, + target_block: u64, + mode: CompareMode, + skip_v2: bool, + skip_v1: bool, +) -> Result { + if parent_block == 0 { + bail!("parent block is 0; need at least one executed block in DB"); + } + + let parent_hash = factory + .block_hash(parent_block) + .with_context(|| format!("failed to fetch parent hash for block {parent_block}"))? + .ok_or_else(|| eyre::eyre!("missing parent hash for block {parent_block}"))?; + + let parent_header = factory + .header_by_number(parent_block) + .with_context(|| format!("failed to fetch parent header for block {parent_block}"))? + .ok_or_else(|| eyre::eyre!("missing parent header for block {parent_block}"))?; + + let target_header = factory + .header_by_number(target_block) + .with_context(|| format!("failed to fetch target header for block {target_block}"))? + .ok_or_else(|| { + eyre::eyre!( + "target header block {} not found in local DB; \ + this harness expects tip+1 to exist after unwind", + target_block + ) + })?; + let canonical_root = target_header.state_root; + + let recovered_block = factory + .recovered_block(target_block.into(), TransactionVariant::NoHash) + .with_context(|| format!("failed to fetch recovered block {target_block}"))? + .ok_or_else(|| { + eyre::eyre!( + "target block {} not found in local DB; \ + ensure block bodies are present for tip+1", + target_block + ) + })?; + + let evm_config = EthEvmConfig::new(factory.chain_spec()); + let executor = evm_config.batch_executor(StateProviderDatabase( + factory + .history_by_block_hash(parent_hash) + .with_context(|| { + format!("failed to open state provider at parent hash {parent_hash:?}") + })?, + )); + + let execute_start = Instant::now(); + let execution = executor + .execute(&recovered_block) + .with_context(|| format!("failed to execute block {target_block}"))?; + let execute_ms = elapsed_ms(execute_start); + let outcome = execution.state; + + let reth_start = Instant::now(); + let reth_root = calculate_reth_root(factory, parent_hash, &outcome) + .with_context(|| format!("failed to compute reth root for block {target_block}"))?; + let reth_ms = elapsed_ms(reth_start); + + // v2/v_experimental fetchers validate an expected chain-tip hash to detect DB movement while reading. + // For current-tip comparisons, this must be the actual DB tip hash (not parent hash). + let expected_tip_block = factory + .last_block_number() + .context("failed to read last block number for sparse-trie consistency check")?; + let expected_tip_hash = factory + .block_hash(expected_tip_block) + .with_context(|| format!("failed to fetch tip hash for block {expected_tip_block}"))? + .ok_or_else(|| eyre::eyre!("missing tip hash for block {expected_tip_block}"))?; + + let shared_cache = SparseTrieSharedCache::new_with_parent_block_data( + expected_tip_hash, + parent_header.state_root, + ); + + let (v2_root, v2_ms, v2_fetched_nodes) = if skip_v2 { + (None, None, None) + } else { + let mut local_v2 = SparseTrieLocalCache::default(); + let v2_start = Instant::now(); + let (v2_root_result, v2_metrics) = calculate_root_hash_with_sparse_trie( + ConsistentDbView::new(factory.clone(), Some((parent_hash, parent_block))), + &outcome, + &[], + &shared_cache, + &mut local_v2, + &None, + ETHSpareMPTVersion::V2, + ); + let v2_ms = elapsed_ms(v2_start); + let v2_root = v2_root_result.with_context(|| "v2 root hash calculation failed")?; + (Some(v2_root), Some(v2_ms), Some(v2_metrics.fetched_nodes)) + }; + + let mut local_v_experimental = SparseTrieLocalCache::default(); + let v_experimental_start = Instant::now(); + let (v_experimental_root_result, v_experimental_metrics) = calculate_root_hash_with_sparse_trie( + ConsistentDbView::new(factory.clone(), Some((parent_hash, parent_block))), + &outcome, + &[], + &shared_cache, + &mut local_v_experimental, + &None, + ETHSpareMPTVersion::VExperimental, + ); + let v_experimental_ms = elapsed_ms(v_experimental_start); + let v_experimental_root = v_experimental_root_result + .with_context(|| "v_experimental root hash calculation failed")?; + + let (v1_root, v1_ms, v1_fetched_nodes) = if skip_v1 { + (None, None, None) + } else { + let mut local_v1 = SparseTrieLocalCache::default(); + let v1_start = Instant::now(); + let (v1_root_result, v1_metrics) = calculate_root_hash_with_sparse_trie( + ConsistentDbView::new(factory.clone(), Some((parent_hash, parent_block))), + &outcome, + &[], + &shared_cache, + &mut local_v1, + &None, + ETHSpareMPTVersion::V1, + ); + let v1_ms = elapsed_ms(v1_start); + let v1_root = v1_root_result.with_context(|| "v1 root hash calculation failed")?; + (Some(v1_root), Some(v1_ms), Some(v1_metrics.fetched_nodes)) + }; + + let v2_mismatch = v2_root.is_some_and(|root| root != canonical_root); + let v1_mismatch = v1_root.is_some_and(|root| root != canonical_root); + if v2_mismatch + || v_experimental_root != canonical_root + || reth_root != canonical_root + || v1_mismatch + { + let v2_root_display = format_optional_root(v2_root.as_ref()); + let v1_root_display = format_optional_root(v1_root.as_ref()); + bail!( + "root mismatch for block {}\n canonical: {:?}\n reth: {:?}\n v2: {}\n v_experimental: {:?}\n v1: {}", + target_block, + canonical_root, + reth_root, + v2_root_display, + v_experimental_root, + v1_root_display + ); + } + + Ok(CompareResult { + mode, + parent_block, + target_block, + canonical_root, + reth_root, + v2_root, + v_experimental_root, + v1_root, + execute_ms, + reth_ms, + v2_ms, + v_experimental_ms, + v1_ms, + v2_fetched_nodes, + v_experimental_fetched_nodes: v_experimental_metrics.fetched_nodes, + v1_fetched_nodes, + }) +} + +fn calculate_reth_root( + factory: &Factory, + parent_hash: B256, + outcome: &BundleState, +) -> Result { + let overlay = OverlayStateProviderFactory::new(factory.clone()); + let hasher = factory + .history_by_block_hash(parent_hash) + .with_context(|| format!("failed to open state provider at parent hash {parent_hash:?}"))?; + let hashed_post_state = hasher.hashed_post_state(outcome); + let trie_input = TrieInput::from_state(hashed_post_state); + ParallelStateRoot::new(overlay, trie_input.prefix_sets.freeze()) + .incremental_root() + .with_context(|| "parallel state root failed") +} + +fn run_unwind_once(cli: &Cli, iteration: usize) -> Result<()> { + let output = Command::new(&cli.reth_bin) + .arg("stage") + .arg("unwind") + .arg(format!("--datadir={}", cli.datadir.display())) + .arg(format!("--chain={}", cli.chain)) + .arg("num-blocks") + .arg("1") + .output() + .with_context(|| { + format!("failed to execute reth unwind command at iteration {iteration}") + })?; + + if !output.status.success() { + bail!( + "reth unwind failed at iteration {} (status: {}):\nstdout:\n{}\nstderr:\n{}", + iteration, + output.status, + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + } + Ok(()) +} + +fn open_provider_factory( + datadir: &Path, + chain: &str, + static_file_blocks_per_file: Option, +) -> Result { + let db_path = datadir.join("db"); + let static_files_path = datadir.join("static_files"); + let chain_spec = parse_chain_spec(chain)?; + + let db = Arc::new( + // Root computation can hold RO transactions long enough to hit MDBX timeout on large blocks. + // Disable timeout for this offline harness path. + open_db_read_only( + &db_path, + DatabaseArguments::new(ClientVersion::default()) + .with_max_read_transaction_duration(Some(MaxReadTransactionDuration::Unbounded)), + ) + .with_context(|| format!("failed to open reth db at {}", db_path.display()))?, + ); + + let static_files = + StaticFileProvider::read_only(&static_files_path, false).with_context(|| { + format!( + "failed to open static files at {}", + static_files_path.display() + ) + })?; + let static_files = if let Some(blocks_per_file) = static_file_blocks_per_file { + static_files.with_custom_blocks_per_file(blocks_per_file) + } else { + static_files + }; + + Ok(ProviderFactory::::new( + db, + chain_spec, + static_files, + )) +} + +fn parse_chain_spec(chain: &str) -> Result> { + match chain.to_ascii_lowercase().as_str() { + "mainnet" => Ok(MAINNET.clone()), + "sepolia" => Ok(SEPOLIA.clone()), + "holesky" => Ok(HOLESKY.clone()), + "hoodi" => Ok(HOODI.clone()), + "dev" => Ok(DEV.clone()), + _ => bail!("unsupported chain: {chain} (supported: mainnet,sepolia,holesky,hoodi,dev)"), + } +} + +fn parse_cli() -> Result { + let mut full = false; + let mut blocks: usize = 1; + let mut reth_bin = expand_tilde("~/reth/target/maxperf/reth"); + let mut datadir = expand_tilde("/home/ubuntu/merkle/"); + let mut chain = "mainnet".to_string(); + let mut tip_fallback = false; + let mut static_file_blocks_per_file = None; + let mut skip_v2 = false; + // v1 is disabled by default on modern DB layouts; opt in with --v1. + let mut skip_v1 = true; + + let mut args = env::args().skip(1); + while let Some(arg) = args.next() { + match arg.as_str() { + "--full" => { + full = true; + } + "--blocks" => { + let value = args + .next() + .ok_or_else(|| eyre::eyre!("--blocks requires a value"))?; + blocks = value + .parse() + .with_context(|| format!("invalid --blocks value: {value}"))?; + } + "--reth-bin" => { + let value = args + .next() + .ok_or_else(|| eyre::eyre!("--reth-bin requires a value"))?; + reth_bin = expand_tilde(&value); + } + "--datadir" => { + let value = args + .next() + .ok_or_else(|| eyre::eyre!("--datadir requires a value"))?; + datadir = expand_tilde(&value); + } + "--chain" => { + chain = args + .next() + .ok_or_else(|| eyre::eyre!("--chain requires a value"))?; + } + "--tip-fallback" => { + tip_fallback = true; + } + "--static-file-blocks-per-file" => { + let value = args + .next() + .ok_or_else(|| eyre::eyre!("--static-file-blocks-per-file requires a value"))?; + let parsed: u64 = value.parse().with_context(|| { + format!("invalid --static-file-blocks-per-file value: {value}") + })?; + if parsed == 0 { + bail!("--static-file-blocks-per-file must be > 0"); + } + static_file_blocks_per_file = Some(parsed); + } + "--skip-v2" => { + skip_v2 = true; + } + "--v1" => { + skip_v1 = false; + } + "--skip-v1" => { + skip_v1 = true; + } + "--help" | "-h" => { + print_help(); + std::process::exit(0); + } + _ => bail!("unknown arg: {arg}"), + } + } + + if blocks == 0 { + bail!("--blocks must be > 0"); + } + + Ok(Cli { + full, + blocks, + reth_bin, + datadir, + chain, + tip_fallback, + static_file_blocks_per_file, + skip_v2, + skip_v1, + }) +} + +fn print_help() { + println!("correctness-harness"); + println!(" default mode:"); + println!(" compares v1, v2, v_experimental and reth for block tip+1"); + println!(" expected setup: reth already unwound by 1 block"); + println!(); + println!(" --full"); + println!(" compares current tip, then unwinds by 1 per iteration"); + println!(); + println!(" flags:"); + println!(" --full run unwind+compare loop"); + println!(" --blocks iterations in --full mode (default: 1)"); + println!(" --reth-bin reth binary (default: ~/reth/target/maxperf/reth)"); + println!(" --datadir reth datadir (default: /home/ubuntu/merkle/)"); + println!(" --chain chain (default: mainnet)"); + println!(" --tip-fallback compare current tip if tip+1 is unavailable"); + println!(" --static-file-blocks-per-file "); + println!(" override static file blocks-per-file lookup"); + println!(" --skip-v2 skip v2 sparse trie comparison"); + println!(" --v1 enable v1 sparse trie comparison (default: off)"); + println!(" --skip-v1 skip v1 sparse trie comparison (default: on)"); +} + +fn print_compare_result(iteration: Option, result: &CompareResult) { + if let Some(iteration) = iteration { + println!("iteration {iteration}:"); + } + + println!( + " mode={} parent={} target={} status=ok", + match result.mode { + CompareMode::TipPlusOne => "tip+1", + CompareMode::TipFallback => "tip-fallback", + CompareMode::CurrentTip => "tip", + }, + result.parent_block, + result.target_block + ); + println!(" roots:"); + println!(" canonical: {:?}", result.canonical_root); + println!( + " reth: {:?} {}", + result.reth_root, + root_match_marker(result.reth_root == result.canonical_root) + ); + println!( + " v2: {} {}", + format_optional_root(result.v2_root.as_ref()), + optional_root_match_marker(result.v2_root.as_ref(), &result.canonical_root) + ); + println!( + " v_experimental: {:?} {}", + result.v_experimental_root, + root_match_marker(result.v_experimental_root == result.canonical_root) + ); + println!( + " v1: {} {}", + format_optional_root(result.v1_root.as_ref()), + optional_root_match_marker(result.v1_root.as_ref(), &result.canonical_root) + ); + let v2_ms = result + .v2_ms + .map(|v| format!("{v:.2}")) + .unwrap_or_else(|| "skipped".to_string()); + let v1_ms = result + .v1_ms + .map(|v| format!("{v:.2}")) + .unwrap_or_else(|| "skipped".to_string()); + println!( + " time_ms: execute={:.2} reth={:.2} v2={} v_experimental={:.2} v1={}", + result.execute_ms, result.reth_ms, v2_ms, result.v_experimental_ms, v1_ms + ); + let v2_nodes = result + .v2_fetched_nodes + .map(|v| v.to_string()) + .unwrap_or_else(|| "skipped".to_string()); + let v1_nodes = result + .v1_fetched_nodes + .map(|v| v.to_string()) + .unwrap_or_else(|| "skipped".to_string()); + println!( + " fetched_nodes: v2={} v_experimental={} v1={}", + v2_nodes, result.v_experimental_fetched_nodes, v1_nodes + ); +} + +fn expand_tilde(path: &str) -> PathBuf { + if let Some(rest) = path.strip_prefix("~/") { + if let Ok(home) = env::var("HOME") { + return PathBuf::from(home).join(rest); + } + } + PathBuf::from(path) +} + +fn elapsed_ms(start: Instant) -> f64 { + start.elapsed().as_secs_f64() * 1000.0 +} + +fn format_optional_root(root: Option<&B256>) -> String { + root.map(|value| format!("{value:?}")) + .unwrap_or_else(|| "None".to_string()) +} + +fn root_match_marker(matches: bool) -> &'static str { + if matches { + "(= canonical)" + } else { + "(!= canonical)" + } +} + +fn optional_root_match_marker(root: Option<&B256>, canonical_root: &B256) -> &'static str { + match root { + Some(value) if value == canonical_root => "(= canonical)", + Some(_) => "(!= canonical)", + None => "(skipped)", + } +} diff --git a/crates/eth-sparse-mpt/src/lib.rs b/crates/eth-sparse-mpt/src/lib.rs index f1782f99a..306fd51e5 100644 --- a/crates/eth-sparse-mpt/src/lib.rs +++ b/crates/eth-sparse-mpt/src/lib.rs @@ -18,6 +18,7 @@ pub mod utils; pub mod v1; pub mod v2; +pub mod v_experimental; #[derive(Debug)] pub struct ChangedAccountData { @@ -65,6 +66,7 @@ impl Default for RootHashThreadPool { pub struct SparseTrieSharedCache { cache_v1: v1::reth_sparse_trie::SparseTrieSharedCache, cache_v2: v2::SharedCacheV2, + cache_v_experimental: v_experimental::SharedCacheVExperimental, } impl SparseTrieSharedCache { @@ -74,19 +76,27 @@ impl SparseTrieSharedCache { ); let mut cache_v2 = v2::SharedCacheV2::default(); cache_v2.last_block_hash = parent_block_hash; - Self { cache_v1, cache_v2 } + let mut cache_v_experimental = v_experimental::SharedCacheVExperimental::default(); + cache_v_experimental.last_block_hash = parent_block_hash; + Self { + cache_v1, + cache_v2, + cache_v_experimental, + } } } #[derive(Debug, Default, Clone)] pub struct SparseTrieLocalCache { - calc: v2::RootHashCalculator, + calc_v2: v2::RootHashCalculator, + calc_v_experimental: v_experimental::RootHashCalculatorExperimental, } #[derive(Debug, Clone, Copy)] pub enum ETHSpareMPTVersion { V1, V2, + VExperimental, } pub fn prefetch_tries_for_accounts<'a, Provider>( @@ -115,6 +125,11 @@ where ETHSpareMPTVersion::V2 => { v2::prefetch_proofs(consistent_db_view, &shared_cache.cache_v2, changed_data) } + ETHSpareMPTVersion::VExperimental => v_experimental::prefetch_proofs( + consistent_db_view, + &shared_cache.cache_v_experimental, + changed_data, + ), } } @@ -160,7 +175,7 @@ where Default::default(), ), ETHSpareMPTVersion::V2 => { - let result = local_cache.calc.calculate_root_hash_with_sparse_trie( + let result = local_cache.calc_v2.calculate_root_hash_with_sparse_trie( consistent_db_view, shared_cache.cache_v2.clone(), outcome, @@ -172,6 +187,21 @@ where Err(err) => (Err(err), Default::default()), } } + ETHSpareMPTVersion::VExperimental => { + let result = local_cache + .calc_v_experimental + .calculate_root_hash_with_sparse_trie( + consistent_db_view, + shared_cache.cache_v_experimental.clone(), + outcome, + &[], + proof_targets, + ); + match result { + Ok((_, proofs, metrics)) => (Ok(proofs), metrics), + Err(err) => (Err(err), Default::default()), + } + } }; if let Some(thread_pool) = thread_pool { thread_pool.rayon_pool.install(calculate) @@ -239,7 +269,7 @@ where (result, metrics) } ETHSpareMPTVersion::V2 => { - let result = local_cache.calc.calculate_root_hash_with_sparse_trie( + let result = local_cache.calc_v2.calculate_root_hash_with_sparse_trie( consistent_db_view, shared_cache.cache_v2.clone(), outcome, @@ -251,5 +281,20 @@ where Err(err) => (Err(err), Default::default()), } } + ETHSpareMPTVersion::VExperimental => { + let result = local_cache + .calc_v_experimental + .calculate_root_hash_with_sparse_trie( + consistent_db_view, + shared_cache.cache_v_experimental.clone(), + outcome, + incremental_change, + &Default::default(), + ); + match result { + Ok((res, _, metrics)) => (Ok(res), metrics), + Err(err) => (Err(err), Default::default()), + } + } } } diff --git a/crates/eth-sparse-mpt/src/v2/trie/mod.rs b/crates/eth-sparse-mpt/src/v2/trie/mod.rs index 459da833a..241202fab 100644 --- a/crates/eth-sparse-mpt/src/v2/trie/mod.rs +++ b/crates/eth-sparse-mpt/src/v2/trie/mod.rs @@ -789,7 +789,7 @@ impl Trie { if self.nodes.is_empty() { return Err(NodeNotFound(Nibbles::new())); } - self.root_hash_node(0, &mut rlp, parallel, proof_store); + self.root_hash_node(0, &mut rlp, parallel, 0, proof_store); self.rlp_encode_node(0, &mut rlp, proof_store); Ok(keccak256(&rlp)) } @@ -799,6 +799,7 @@ impl Trie { node_idx: usize, rlp: &mut Vec, parallel: bool, + depth: usize, proof_store: &ProofStore, ) { if self.hashed_nodes[node_idx] { @@ -810,20 +811,18 @@ impl Trie { .expect("root_hash_node: node not found"); match node { DiffTrieNode::Branch { children } => { - let compute_children_this_thread = if !parallel { - true - } else { + let use_rayon = parallel && depth == 0 && { let local_children = self.branch_node_children[*children] .iter() .filter(|c| matches!(c, Some(NodePtr::Local(_)))) .count(); - local_children <= 1 + local_children > 1 }; - if compute_children_this_thread { + if !use_rayon { for child in self.branch_node_children[*children].into_iter().flatten() { if let NodePtr::Local(child) = child { - self.root_hash_node(child, rlp, parallel, proof_store); + self.root_hash_node(child, rlp, false, depth + 1, proof_store); } } } else { @@ -832,7 +831,13 @@ impl Trie { if let NodePtr::Local(child) = child { scope.spawn(move |_| { let mut rlp = Vec::new(); - self.root_hash_node(child, &mut rlp, parallel, proof_store); + self.root_hash_node( + child, + &mut rlp, + false, + depth + 1, + proof_store, + ); }) } } @@ -842,7 +847,7 @@ impl Trie { } DiffTrieNode::Extension { next_node, .. } => { if let NodePtr::Local(child) = next_node { - self.root_hash_node(*child, rlp, parallel, proof_store); + self.root_hash_node(*child, rlp, parallel, depth + 1, proof_store); } self.calculate_rlp_pointer_node(node_idx, rlp, proof_store); } diff --git a/crates/eth-sparse-mpt/src/v_experimental/fetch.rs b/crates/eth-sparse-mpt/src/v_experimental/fetch.rs new file mode 100644 index 000000000..43f27fcaf --- /dev/null +++ b/crates/eth-sparse-mpt/src/v_experimental/fetch.rs @@ -0,0 +1,157 @@ +use std::sync::Arc; + +use crate::{ + utils::{convert_nibbles_to_reth_nybbles, convert_reth_nybbles_to_nibbles, HashMap}, + SparseTrieError, +}; +use alloy_primitives::map::B256Set; +use parking_lot::Mutex; +use rayon::prelude::*; + +use alloy_primitives::B256; +use nybbles::Nibbles; +use reth_provider::{ + providers::ConsistentDbView, BlockHashReader, BlockNumReader, BlockReader, DBProvider, + DatabaseProviderFactory, +}; +use reth_trie::{ + proof::{Proof, StorageProof}, + MultiProofTargets, +}; +use reth_trie_db::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory}; + +use super::SharedCacheV2; + +#[derive(Debug, Default)] +pub struct MissingNodesFetcher { + storage_proof_targets: HashMap)>, + account_proof_targets: Vec, + account_proof_requested_nodes: Vec, +} + +impl MissingNodesFetcher { + pub fn is_empty(&self) -> bool { + self.storage_proof_targets.is_empty() && self.account_proof_targets.is_empty() + } + + pub fn add_missing_storage_node(&mut self, hashed_address: &B256, node: Nibbles) { + let entry = self + .storage_proof_targets + .entry(*hashed_address) + .or_default(); + entry.0.insert(pad_path(node.clone())); + entry.1.push(node); + } + + pub fn add_missing_account_node(&mut self, node: Nibbles) { + self.account_proof_targets.push(pad_path(node.clone())); + self.account_proof_requested_nodes.push(node); + } + + // fetch currently accumulated nodes into shared cache + pub fn fetch_nodes( + &mut self, + shared_cache: &SharedCacheV2, + consistent_db_view: &ConsistentDbView, + ) -> Result + where + Provider: DatabaseProviderFactory + Send + Sync, + { + let fetched_nodes: Arc> = Default::default(); + + let last_block_hash = shared_cache.last_block_hash; + std::mem::take(&mut self.storage_proof_targets) + .into_par_iter() + .map( + |(hashed_address, (targets, requested_proofs))| -> Result<(), SparseTrieError> { + let provider = consistent_db_view + .provider_ro() + .map_err(SparseTrieError::other)?; + if !last_block_hash.is_zero() { + let block_number = provider + .last_block_number() + .map_err(SparseTrieError::other)?; + let block_hash = provider + .block_hash(block_number) + .map_err(SparseTrieError::other)?; + if block_hash != Some(shared_cache.last_block_hash) { + return Err(SparseTrieError::WrongDatabaseTrieError); + } + } + + let proof = StorageProof::new_hashed( + DatabaseTrieCursorFactory::new(provider.tx_ref()), + DatabaseHashedCursorFactory::new(provider.tx_ref()), + hashed_address, + ); + let storge_multiproof = proof + .storage_multiproof(targets) + .map_err(SparseTrieError::other)?; + *fetched_nodes.lock() += requested_proofs.len(); + for requested_proof in requested_proofs { + let proof_for_node = storge_multiproof.subtree.matching_nodes_sorted( + &convert_nibbles_to_reth_nybbles(requested_proof.clone()), + ); + let reth_proof_for_node = proof_for_node + .into_iter() + .map(|(k, v)| (convert_reth_nybbles_to_nibbles(k), v)) + .collect(); + let proof_store = + shared_cache.account_proof_store_hashed_address(&hashed_address); + proof_store + .add_proof(requested_proof, reth_proof_for_node) + .map_err(SparseTrieError::other)?; + } + Ok(()) + }, + ) + .collect::>()?; + + let provider = consistent_db_view + .provider_ro() + .map_err(SparseTrieError::other)?; + if !last_block_hash.is_zero() { + let block_number = provider + .last_block_number() + .map_err(SparseTrieError::other)?; + let block_hash = provider + .block_hash(block_number) + .map_err(SparseTrieError::other)?; + if block_hash != Some(shared_cache.last_block_hash) { + return Err(SparseTrieError::WrongDatabaseTrieError); + } + } + + let proof = Proof::new( + DatabaseTrieCursorFactory::new(provider.tx_ref()), + DatabaseHashedCursorFactory::new(provider.tx_ref()), + ); + let targets = MultiProofTargets::accounts(std::mem::take(&mut self.account_proof_targets)); + let multiproof = proof.multiproof(targets).map_err(SparseTrieError::other)?; + + *fetched_nodes.lock() += self.account_proof_requested_nodes.len(); + for requested_node in self.account_proof_requested_nodes.drain(..) { + let proof_for_node = multiproof + .account_subtree + .matching_nodes_sorted(&convert_nibbles_to_reth_nybbles(requested_node.clone())); + + let reth_proof_for_node = proof_for_node + .into_iter() + .map(|(k, v)| (convert_reth_nybbles_to_nibbles(k), v)) + .collect(); + shared_cache + .account_trie + .add_proof(requested_node, reth_proof_for_node) + .map_err(SparseTrieError::other)?; + } + let fetched_nodes = *fetched_nodes.lock(); + Ok(fetched_nodes) + } +} + +fn pad_path(mut path: Nibbles) -> B256 { + path.as_mut_vec_unchecked().resize(64, 0); + let mut res = B256::default(); + path.pack_to(res.as_mut_slice()); + res +} diff --git a/crates/eth-sparse-mpt/src/v_experimental/mod.rs b/crates/eth-sparse-mpt/src/v_experimental/mod.rs new file mode 100644 index 000000000..612c0233f --- /dev/null +++ b/crates/eth-sparse-mpt/src/v_experimental/mod.rs @@ -0,0 +1,1207 @@ +use alloy_primitives::{keccak256, Address, Bytes, B256, U256}; +use dashmap::DashMap; +use fetch::MissingNodesFetcher; +use nybbles::Nibbles; +use parking_lot::{Mutex, RwLock}; +use rayon::prelude::*; +use reth_provider::{providers::ConsistentDbView, BlockReader, DatabaseProviderFactory}; +use reth_trie::TrieAccount; +use revm::{ + database::{BundleAccount, BundleState}, + state::AccountInfo, +}; +use rustc_hash::FxBuildHasher; +use std::{ops::Range, sync::Arc, time::Instant}; +use trie::{ + proof_store::ProofStore, DeletionError, InsertValue, NodeNotFound, ProofError, ProofWithValue, + Trie, +}; + +use crate::{ + utils::{HashMap, HashSet}, + ChangedAccountData, SparseTrieError, SparseTrieMetrics, +}; + +pub mod fetch; +pub mod trie; + +const PARALLEL_HASHING_STORAGE_NODES: bool = true; + +pub type SharedCacheVExperimental = SharedCacheV2; +pub type RootHashCalculatorExperimental = RootHashCalculator; + +#[derive(Debug, Default, Clone)] +pub struct SharedCacheV2 { + pub account_trie: ProofStore, + pub storage_tries: Arc>, + pub last_block_hash: B256, +} + +impl SharedCacheV2 { + pub fn account_proof_store_hashed_address(&self, hashed_address: &B256) -> ProofStore { + if let Some(store) = self.storage_tries.get(hashed_address) { + store.value().clone() + } else { + self.storage_tries + .entry(*hashed_address) + .or_default() + .clone() + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum StorageTrieStatus { + InsertsNotProcessed, + InsertsProcessed, + Hashed, +} + +impl StorageTrieStatus { + fn needs_processing(&self) -> bool { + match self { + StorageTrieStatus::InsertsNotProcessed => true, + StorageTrieStatus::InsertsProcessed => false, + StorageTrieStatus::Hashed => false, + } + } +} + +#[derive(Debug, Default)] +pub struct RootHashCalculator { + storage: DashMap>, FxBuildHasher>, + changed_account: Arc>>, + + account_trie: AccountTrieCalculator, + + shared_cache: SharedCacheV2, + + // if set, only changes for these accounts will be incrementally applied + incremental_account_change: HashSet
, +} + +impl Clone for RootHashCalculator { + fn clone(&self) -> Self { + Self { + storage: self + .storage + .iter() + .map(|entry| { + ( + *entry.key(), + Arc::new(Mutex::new(entry.value().lock().clone())), + ) + }) + .collect(), + changed_account: Arc::new(RwLock::new(self.changed_account.read().clone())), + account_trie: self.account_trie.clone(), + shared_cache: self.shared_cache.clone(), + incremental_account_change: self.incremental_account_change.clone(), + } + } +} + +#[derive(Debug, Clone)] +struct AppliedAccountOp { + inserted_value: Option, + revert_key: Nibbles, + revert_value: Option>, +} + +#[derive(Debug, Clone, Default)] +struct AccountTrieCalculator { + trie: Trie, + + applied_account_ops_current_iteration: HashMap, + + revert_account_ops: Vec, + revert_account_ops_done: Vec, + + insert_keys: Vec, + insert_values: Vec>, + insert_account_keys: Vec
, + insert_account_values: Vec, + insert_ok: Vec, + + delete_keys: Vec, + delete_account_keys: Vec
, + delete_ok: Vec, + + proof_keys: Vec, + proof_account_keys: Vec
, + proof_ok: Vec, + proof_result: Vec<(Address, Vec<(Nibbles, Vec)>)>, + + missing_nodes: Vec, + missing_nodes_requested: Vec, + + applied_account_ops_previous_iteration: HashMap, +} + +impl AccountTrieCalculator { + fn clear(&mut self) { + self.applied_account_ops_current_iteration.clear(); + self.revert_account_ops.clear(); + self.revert_account_ops_done.clear(); + + self.insert_keys.clear(); + self.insert_account_keys.clear(); + self.insert_account_values.clear(); + self.insert_values.clear(); + self.insert_ok.clear(); + self.delete_keys.clear(); + self.delete_account_keys.clear(); + self.delete_ok.clear(); + self.missing_nodes.clear(); + + self.proof_keys.clear(); + self.proof_account_keys.clear(); + self.proof_ok.clear(); + self.proof_result.clear(); + // self.trie.clear(); + } +} + +#[derive(Debug, Clone)] +struct AppliedStorageOp { + inserted_value: U256, + revert_key: Nibbles, + revert_value: Option>, +} + +#[derive(Debug, Clone, Default)] +struct StorageCalculator { + hashed_address: B256, + unpacked_hashed_address: Nibbles, + + account_info: Option, + + applied_storage_ops_current_iteration: HashMap, + + revert_storage_ops: Vec, + revert_storage_ops_done: Vec, + + insert_keys: Vec, + insert_values: Vec>, + insert_storage_key: Vec, + insert_storage_value: Vec, + insert_ok: Vec, + + delete_keys: Vec, + delete_storage_key: Vec, + delete_ok: Vec, + + missing_nodes: Vec, + missing_nodes_requested: Vec, + + trie: Trie, + applied_storage_ops_previous_iteration: HashMap, + + proof_store: ProofStore, + hash: B256, +} + +impl StorageCalculator { + fn new(addres: Address, shared_cache: &SharedCacheV2) -> Self { + let hashed_address = keccak256(addres); + let unpacked_hashed_address = Nibbles::unpack(hashed_address.as_slice()); + let proof_store = shared_cache.account_proof_store_hashed_address(&hashed_address); + + StorageCalculator { + hashed_address, + unpacked_hashed_address, + account_info: None, + revert_storage_ops: Vec::new(), + revert_storage_ops_done: Vec::new(), + insert_keys: Vec::new(), + insert_storage_key: Vec::new(), + insert_storage_value: Vec::new(), + insert_values: Vec::new(), + insert_ok: Vec::new(), + delete_keys: Vec::new(), + delete_storage_key: Vec::new(), + delete_ok: Vec::new(), + missing_nodes: Vec::new(), + missing_nodes_requested: Vec::new(), + trie: Trie::default(), + applied_storage_ops_previous_iteration: HashMap::default(), + applied_storage_ops_current_iteration: HashMap::default(), + hash: B256::ZERO, + proof_store, + } + } + // it does not clear the trie with currently applied op! + fn clear(&mut self) { + self.account_info = None; + self.applied_storage_ops_current_iteration.clear(); + self.revert_storage_ops.clear(); + self.revert_storage_ops_done.clear(); + + self.insert_keys.clear(); + self.insert_values.clear(); + self.insert_storage_key.clear(); + self.insert_storage_value.clear(); + self.insert_ok.clear(); + + self.delete_keys.clear(); + self.delete_storage_key.clear(); + self.delete_ok.clear(); + + self.missing_nodes.clear(); + self.missing_nodes_requested.clear(); + } +} + +pub fn prefetch_proofs<'a, Provider>( + consistent_db_view: ConsistentDbView, + shared_cache: &SharedCacheV2, + changed_data: impl Iterator, +) -> Result +where + Provider: DatabaseProviderFactory + Send + Sync, +{ + let mut metrics = SparseTrieMetrics::default(); + let mut fetcher = MissingNodesFetcher::default(); + + for data in changed_data { + let hashed_address = keccak256(data.address.as_slice()); + let account_node = Nibbles::unpack(hashed_address); + fetcher.add_missing_account_node(account_node); + for (slot, _) in &data.slots { + let storage_node = Nibbles::unpack(keccak256(B256::from(*slot)).as_slice()); + fetcher.add_missing_storage_node(&hashed_address, storage_node); + } + } + metrics.fetched_nodes += fetcher.fetch_nodes(shared_cache, &consistent_db_view)?; + + Ok(metrics) +} + +impl RootHashCalculator { + pub fn new(shared_cache: SharedCacheV2) -> Self { + Self { + account_trie: AccountTrieCalculator::default(), + storage: DashMap::default(), + changed_account: Default::default(), + shared_cache, + incremental_account_change: Default::default(), + } + } + + fn get_account_storage(&self, address: &Address) -> Arc> { + if let Some(v) = self.storage.get(address) { + v.value().clone() + } else { + self.storage + .entry(*address) + .or_insert_with(|| { + Arc::new(Mutex::new(StorageCalculator::new( + *address, + &self.shared_cache, + ))) + }) + .value() + .clone() + } + } + + fn prepare_changes_for_one_storage_trie( + &self, + address: Address, + bundle_account: &BundleAccount, + ) { + let storage_calc = self.get_account_storage(&address); + let mut storage_calc = storage_calc.lock(); + let storage_calc = &mut *storage_calc; + + storage_calc.clear(); + storage_calc.account_info = bundle_account.account_info().map(|a| a.without_code()); + + if storage_calc.account_info.is_none() { + // account processed, no need to compute storage hash + storage_calc.hash = B256::ZERO; + self.changed_account + .write() + .push((address, StorageTrieStatus::Hashed)); + return; + } + + for (storage_key, storage_value) in &bundle_account.storage { + if !storage_value.is_changed() { + continue; + } + + let storage_value = storage_value.present_value(); + + if let Some(applied_op) = storage_calc + .applied_storage_ops_previous_iteration + .remove(storage_key) + { + if applied_op.inserted_value == storage_value { + storage_calc + .applied_storage_ops_current_iteration + .insert(*storage_key, applied_op); + continue; + } else { + storage_calc.revert_storage_ops.push(applied_op); + storage_calc.revert_storage_ops_done.push(false); + } + } + + let hashed_key = Nibbles::unpack(keccak256(B256::from(*storage_key)).as_slice()); + if !storage_value.is_zero() { + let value = alloy_rlp::encode(storage_value); + storage_calc.insert_keys.push(hashed_key); + storage_calc.insert_values.push(value); + storage_calc.insert_storage_key.push(*storage_key); + storage_calc.insert_storage_value.push(storage_value); + storage_calc.insert_ok.push(false); + } else { + storage_calc.delete_keys.push(hashed_key); + storage_calc.delete_storage_key.push(*storage_key); + storage_calc.delete_ok.push(false); + } + } + + // revert all applied ops from previous iteration + let mut applied_storage_ops_previous_iteration = + std::mem::take(&mut storage_calc.applied_storage_ops_previous_iteration); + for (_, applied_op) in applied_storage_ops_previous_iteration.drain() { + storage_calc.revert_storage_ops.push(applied_op); + storage_calc.revert_storage_ops_done.push(false); + } + storage_calc.applied_storage_ops_previous_iteration = + applied_storage_ops_previous_iteration; + + if storage_calc.delete_keys.is_empty() + && storage_calc.insert_keys.is_empty() + && storage_calc.revert_storage_ops.is_empty() + && !storage_calc.hash.is_zero() + && !storage_calc.trie.is_uninit() + { + std::mem::swap( + &mut storage_calc.applied_storage_ops_previous_iteration, + &mut storage_calc.applied_storage_ops_current_iteration, + ); + self.changed_account + .write() + .push((address, StorageTrieStatus::Hashed)); + } else { + storage_calc.hash = B256::ZERO; + self.changed_account + .write() + .push((address, StorageTrieStatus::InsertsNotProcessed)); + } + } + + fn prepare_changes_for_storage_trie(&mut self, outcome: &BundleState) -> eyre::Result<()> { + self.changed_account.write().clear(); + + let incremental_change = !self.incremental_account_change.is_empty(); + + if incremental_change { + self.incremental_account_change.iter().for_each(|address| { + let bundle_account = outcome + .account(address) + .expect("account with incremental change is not in the BundleState"); + self.prepare_changes_for_one_storage_trie(*address, bundle_account) + }); + } else { + outcome + .state() + .iter() + .map(|(a, acc)| (*a, acc)) + .par_bridge() + .for_each(|(address, bundle_account)| { + if bundle_account.status.is_not_modified() { + return; + } + self.prepare_changes_for_one_storage_trie(address, bundle_account) + }); + } + + Ok(()) + } + + fn do_first_fetch( + &mut self, + consistent_db_view: &ConsistentDbView, + stats: &mut Stats, + ) -> Result<(), SparseTrieError> + where + Provider: DatabaseProviderFactory + Send + Sync, + { + stats.start(); + + let fetcher = Arc::new(Mutex::new(MissingNodesFetcher::default())); + self.changed_account + .read() + .par_iter() + .for_each(|(address, _)| { + let fetcher = fetcher.clone(); + let storage_calc = self.get_account_storage(address); + let storage_calc = storage_calc.lock(); + + if !self + .shared_cache + .account_trie + .has_proof(&storage_calc.unpacked_hashed_address) + { + fetcher + .lock() + .add_missing_account_node(storage_calc.unpacked_hashed_address.clone()); + } + + if storage_calc.insert_keys.is_empty() && storage_calc.delete_keys.is_empty() { + let node = Nibbles::new(); + if !storage_calc.proof_store.has_proof(&node) { + fetcher + .lock() + .add_missing_storage_node(&storage_calc.hashed_address, node); + } + } + + for node in storage_calc + .insert_keys + .iter() + .chain(storage_calc.delete_keys.iter()) + { + if !storage_calc.proof_store.has_proof(node) { + fetcher + .lock() + .add_missing_storage_node(&storage_calc.hashed_address, node.clone()); + } + } + }); + + stats.measure_other(); + + let mut fetcher = fetcher.lock(); + if !fetcher.is_empty() { + stats.start_proof_fetch_db(); + let nodes_fetched = fetcher.fetch_nodes(&self.shared_cache, consistent_db_view)?; + stats.fetched_nodes += nodes_fetched; + stats.measure_proof_fetch_db_part(); + } + + Ok(()) + } + + // true if no missing nodes + fn process_storage_tries_update(&self, address: Address) -> eyre::Result { + let storage_calc = self.get_account_storage(&address); + let mut storage_calc = storage_calc.lock(); + let storage_calc = &mut *storage_calc; + storage_calc.missing_nodes.clear(); + + for i in 0..storage_calc.revert_storage_ops_done.len() { + if storage_calc.revert_storage_ops_done[i] { + continue; + } + + let applied_op = &storage_calc.revert_storage_ops[i]; + + let missing_node = match applied_op.revert_value.clone() { + Some(old_value) => { + match storage_calc.trie.insert_nibble_key( + &applied_op.revert_key, + InsertValue::StoredValue(old_value), + ) { + Ok(_) => None, + Err(node_not_found) => Some(node_not_found.0), + } + } + None => { + match storage_calc.trie.delete_nibbles_key(&applied_op.revert_key) { + Ok(_) => None, + Err(DeletionError::KeyNotFound) => { + eyre::bail!("reverting nodes, can't delete key that is not in the trie (storage)"); + } + Err(DeletionError::NodeNotFound(node_not_found)) => Some(node_not_found.0), + } + } + }; + match missing_node { + Some(node) => { + storage_calc.missing_nodes.push(node); + } + None => { + storage_calc.revert_storage_ops_done[i] = true; + } + } + } + // this way we alway process reverts first + if !storage_calc.missing_nodes.is_empty() { + return Ok(false); + } + + for i in 0..storage_calc.insert_ok.len() { + if storage_calc.insert_ok[i] { + continue; + } + + let insertion_result = storage_calc.trie.insert_nibble_key( + &storage_calc.insert_keys[i], + InsertValue::Value(&storage_calc.insert_values[i]), + ); + + match insertion_result { + Ok(revert_value) => { + storage_calc.insert_ok[i] = true; + storage_calc.applied_storage_ops_current_iteration.insert( + storage_calc.insert_storage_key[i], + AppliedStorageOp { + inserted_value: storage_calc.insert_storage_value[i], + revert_key: storage_calc.insert_keys[i].clone(), + revert_value, + }, + ); + } + Err(NodeNotFound(missing_node)) => { + storage_calc.missing_nodes.push(missing_node); + } + } + } + + for i in 0..storage_calc.delete_ok.len() { + if storage_calc.delete_ok[i] { + continue; + } + let deletion_result = storage_calc + .trie + .delete_nibbles_key(&storage_calc.delete_keys[i]); + match deletion_result { + Ok(revert_value) => { + storage_calc.delete_ok[i] = true; + storage_calc.applied_storage_ops_current_iteration.insert( + storage_calc.delete_storage_key[i], + AppliedStorageOp { + inserted_value: U256::ZERO, + revert_key: storage_calc.delete_keys[i].clone(), + revert_value: Some(revert_value), + }, + ); + } + Err(DeletionError::NodeNotFound(NodeNotFound(missing_node))) => { + storage_calc.missing_nodes.push(missing_node); + } + Err(DeletionError::KeyNotFound) => { + eyre::bail!("Deleting key that is not in the trie"); + } + } + } + + if storage_calc.trie.is_uninit() { + storage_calc.missing_nodes.push(Nibbles::new()); + } + Ok(storage_calc.missing_nodes.is_empty()) + } + + fn process_storage_tries_updates(&mut self) -> eyre::Result { + let all_changed_processed = Arc::new(Mutex::new(true)); + self.changed_account + .write() + .par_iter_mut() + .map(|(address, status)| -> eyre::Result<()> { + // if account is done, just return + if !status.needs_processing() { + return Ok(()); + } + if self.process_storage_tries_update(*address)? { + *status = StorageTrieStatus::InsertsProcessed; + } else { + *all_changed_processed.lock() = false; + } + Ok(()) + }) + .collect::>()?; + let res = *all_changed_processed.lock(); + Ok(res) + } + + fn fetch_missing_storage_nodes( + &mut self, + consistent_db_view: &ConsistentDbView, + stats: &mut Stats, + ) -> Result<(), SparseTrieError> + where + Provider: DatabaseProviderFactory + Send + Sync, + { + let fetcher = Arc::new(Mutex::new(MissingNodesFetcher::default())); + + self.changed_account.read().par_iter().for_each(|(address, status)| { + if !status.needs_processing() { + return; + } + let storage_calc = self.get_account_storage(address); + let mut storage_calc = storage_calc.lock(); + let storage_calc = &mut *storage_calc; + storage_calc.missing_nodes_requested.clear(); + for missing_node in storage_calc.missing_nodes.drain(..) { + if storage_calc.proof_store.has_proof(&missing_node) { + let ok = storage_calc.trie.try_add_proof_from_proof_store(&missing_node, &storage_calc.proof_store).expect("should be able to insert proofs from proof store when they are found (storage trie)"); + assert!(ok, "proof is not added (storage trie)"); + } else { + storage_calc.missing_nodes_requested.push(missing_node.clone()); + fetcher.lock().add_missing_storage_node(&storage_calc.hashed_address, missing_node); + } + } + }); + + let mut fetcher = fetcher.lock(); + if !fetcher.is_empty() { + stats.start_proof_fetch_db(); + let nodes_fetched = fetcher.fetch_nodes(&self.shared_cache, consistent_db_view)?; + stats.fetched_nodes += nodes_fetched; + stats.measure_proof_fetch_db_part(); + } + + self.changed_account.read().par_iter().for_each(|(address, status)| { + if !status.needs_processing() { + return; + } + let storage_calc = self.get_account_storage(address); + let mut storage_calc = storage_calc.lock(); + let storage_calc = &mut *storage_calc; + for missing_node in storage_calc.missing_nodes_requested.drain(..) { + if storage_calc.proof_store.has_proof(&missing_node) { + let ok = storage_calc.trie.try_add_proof_from_proof_store(&missing_node, &storage_calc.proof_store).expect("should be able to insert proofs from proof store when they are found (storage trie)"); + assert!(ok, "proof is not added (storage trie)"); + } else { + panic!("Missing node that was just fetched is not there (storage trie)"); + } + } + }); + Ok(()) + } + + fn hash_storage_tries(&mut self) { + self.changed_account + .read() + .par_iter() + .for_each(|(address, status)| { + if status == &StorageTrieStatus::Hashed { + return; + } + assert_eq!( + status, + &StorageTrieStatus::InsertsProcessed, + "storage trie is not updated properly" + ); + let storage_calc = self.get_account_storage(address); + let mut storage_calc = storage_calc.lock(); + let storage_calc = &mut *storage_calc; + let hash = storage_calc + .trie + .root_hash(PARALLEL_HASHING_STORAGE_NODES, &storage_calc.proof_store) + .expect("missing node while hahsing storage trie"); + storage_calc.hash = hash; + std::mem::swap( + &mut storage_calc.applied_storage_ops_previous_iteration, + &mut storage_calc.applied_storage_ops_current_iteration, + ); + }); + } + + fn prepare_changes_account_trie(&mut self, proof_targets: &HashSet
) { + self.account_trie.clear(); + + for (address, _) in &*self.changed_account.read() { + let storage_calc = self.get_account_storage(address); + let storage_calc = storage_calc.lock(); + + let trie_account = storage_calc.account_info.as_ref().map(|account_info| { + assert!(!storage_calc.hash.is_zero()); + TrieAccount { + nonce: account_info.nonce, + balance: account_info.balance, + storage_root: storage_calc.hash, + code_hash: account_info.code_hash, + } + }); + + if let Some(applied_op) = self + .account_trie + .applied_account_ops_previous_iteration + .remove(address) + { + if applied_op.inserted_value == trie_account { + self.account_trie + .applied_account_ops_current_iteration + .insert(*address, applied_op); + continue; + } else { + self.account_trie.revert_account_ops.push(applied_op); + self.account_trie.revert_account_ops_done.push(false); + } + } + + let key = storage_calc.unpacked_hashed_address.clone(); + if let Some(trie_account) = trie_account { + let value = alloy_rlp::encode(trie_account); + self.account_trie.insert_keys.push(key); + self.account_trie.insert_values.push(value); + self.account_trie.insert_account_keys.push(*address); + self.account_trie.insert_account_values.push(trie_account); + self.account_trie.insert_ok.push(false); + } else { + self.account_trie.delete_keys.push(key); + self.account_trie.delete_account_keys.push(*address); + self.account_trie.delete_ok.push(false); + } + } + + let incremental_change = !self.incremental_account_change.is_empty(); + + for (address, applied_op) in self + .account_trie + .applied_account_ops_previous_iteration + .drain() + { + if incremental_change && !self.incremental_account_change.contains(&address) { + self.account_trie + .applied_account_ops_current_iteration + .insert(address, applied_op); + } else { + self.account_trie.revert_account_ops.push(applied_op); + self.account_trie.revert_account_ops_done.push(false); + } + } + + for address in proof_targets { + let storage_calc = self.get_account_storage(address); + let storage_calc = storage_calc.lock(); + let key = storage_calc.unpacked_hashed_address.clone(); + self.account_trie.proof_keys.push(key); + self.account_trie.proof_account_keys.push(*address); + self.account_trie.proof_ok.push(false); + } + } + + fn process_account_tries_update(&mut self) -> eyre::Result { + let account_trie = &mut self.account_trie; + account_trie.missing_nodes.clear(); + + for i in 0..account_trie.revert_account_ops_done.len() { + if account_trie.revert_account_ops_done[i] { + continue; + } + let applied_op = &account_trie.revert_account_ops[i]; + let missing_node = match applied_op.revert_value.clone() { + Some(old_value) => { + match account_trie.trie.insert_nibble_key( + &applied_op.revert_key, + InsertValue::StoredValue(old_value), + ) { + Ok(_) => None, + Err(node_not_found) => Some(node_not_found.0), + } + } + None => { + match account_trie.trie.delete_nibbles_key(&applied_op.revert_key) { + Ok(_) => None, + Err(DeletionError::KeyNotFound) => { + eyre::bail!("reverting nodes, can't delete key that is not in the trie (accounts)"); + } + Err(DeletionError::NodeNotFound(node_not_found)) => Some(node_not_found.0), + } + } + }; + + match missing_node { + Some(node) => { + account_trie.missing_nodes.push(node); + } + None => { + account_trie.revert_account_ops_done[i] = true; + } + } + } + // this way we alway process reverts first + if !account_trie.missing_nodes.is_empty() { + return Ok(false); + } + + for i in 0..account_trie.insert_ok.len() { + if account_trie.insert_ok[i] { + continue; + } + let insertion_result = account_trie.trie.insert_nibble_key( + &account_trie.insert_keys[i], + InsertValue::Value(&account_trie.insert_values[i]), + ); + match insertion_result { + Ok(revert_value) => { + account_trie.insert_ok[i] = true; + account_trie.applied_account_ops_current_iteration.insert( + account_trie.insert_account_keys[i], + AppliedAccountOp { + inserted_value: Some(account_trie.insert_account_values[i]), + revert_key: account_trie.insert_keys[i].clone(), + revert_value, + }, + ); + } + Err(NodeNotFound(missing_node)) => { + account_trie.missing_nodes.push(missing_node); + } + } + } + + for i in 0..account_trie.delete_ok.len() { + if account_trie.delete_ok[i] { + continue; + } + let deletion_result = account_trie + .trie + .delete_nibbles_key(&account_trie.delete_keys[i]); + match deletion_result { + Ok(revert_value) => { + account_trie.delete_ok[i] = true; + account_trie.applied_account_ops_current_iteration.insert( + account_trie.delete_account_keys[i], + AppliedAccountOp { + inserted_value: None, + revert_key: account_trie.delete_keys[i].clone(), + revert_value: Some(revert_value), + }, + ); + } + Err(DeletionError::NodeNotFound(NodeNotFound(missing_node))) => { + account_trie.missing_nodes.push(missing_node); + } + Err(DeletionError::KeyNotFound) => { + eyre::bail!("Deleting key that is not in the trie (account trie)"); + } + } + } + Ok(account_trie.missing_nodes.is_empty()) + } + + fn process_account_trie_proofs(&mut self, shared_cache: &SharedCacheV2) -> eyre::Result { + let account_trie = &mut self.account_trie; + account_trie.missing_nodes.clear(); + for i in 0..account_trie.proof_ok.len() { + if account_trie.proof_ok[i] { + continue; + } + let proof_result = account_trie + .trie + .get_proof_nibbles_key(&account_trie.proof_keys[i], &shared_cache.account_trie); + match proof_result { + Ok(ProofWithValue { proof, .. }) => { + account_trie.proof_ok[i] = true; + account_trie + .proof_result + .push((account_trie.proof_account_keys[i], proof)); + } + Err(ProofError::TrieIsDirty) => { + eyre::bail!("Trie is not hashed before fetching proofs") + } + Err(ProofError::NodeNotFound(NodeNotFound(missing_node))) => { + account_trie.missing_nodes.push(missing_node); + } + } + } + Ok(account_trie.missing_nodes.is_empty()) + } + + fn fetch_missing_account_trie_nodes( + &mut self, + consistent_db_view: &ConsistentDbView, + stats: &mut Stats, + ) -> Result<(), SparseTrieError> + where + Provider: DatabaseProviderFactory + Send + Sync, + { + let mut fetcher = MissingNodesFetcher::default(); + + let proof_store = &self.shared_cache.account_trie; + let account_trie = &mut self.account_trie; + account_trie.missing_nodes_requested.clear(); + for missing_node in account_trie.missing_nodes.drain(..) { + if proof_store.has_proof(&missing_node) { + let ok = account_trie.trie.try_add_proof_from_proof_store(&missing_node, proof_store).expect("should be able to insert proofs from proof store when they are found (storage trie)"); + assert!(ok, "proof is not added (storage trie)"); + } else { + account_trie + .missing_nodes_requested + .push(missing_node.clone()); + fetcher.add_missing_account_node(missing_node); + } + } + + if !fetcher.is_empty() { + stats.start_proof_fetch_db(); + let nodes_fetched = fetcher.fetch_nodes(&self.shared_cache, consistent_db_view)?; + stats.fetched_nodes += nodes_fetched; + stats.measure_proof_fetch_db_part(); + } + + for missing_node in account_trie.missing_nodes_requested.drain(..) { + if proof_store.has_proof(&missing_node) { + let ok = account_trie.trie.try_add_proof_from_proof_store(&missing_node, proof_store).expect("should be able to insert proofs from proof store when they are found (account trie)"); + assert!(ok, "proof is not added (account trie)"); + } else { + panic!("Missing node that was just fetched is not there (account trie)"); + } + } + + Ok(()) + } + + pub fn hash_account_trie(&mut self, shared_cache: &SharedCacheV2) -> B256 { + let hash = self + .account_trie + .trie + .root_hash(true, &shared_cache.account_trie) + .expect("failed to hash account trie"); + std::mem::swap( + &mut self.account_trie.applied_account_ops_current_iteration, + &mut self.account_trie.applied_account_ops_previous_iteration, + ); + hash + } + + pub fn calculate_root_hash_with_sparse_trie( + &mut self, + consistent_db_view: ConsistentDbView, + shared_cache: SharedCacheV2, + outcome: &BundleState, + incremental_change: &[Address], + proof_targets: &HashSet
, + ) -> Result<(B256, HashMap>, SparseTrieMetrics), SparseTrieError> + where + Provider: DatabaseProviderFactory + Send + Sync, + { + if !incremental_change.is_empty() { + self.incremental_account_change.extend(incremental_change); + } else { + self.incremental_account_change.clear(); + } + + let mut stats = Stats::default(); + stats.start_global(); + + self.shared_cache = shared_cache.clone(); + + stats.start(); + self.prepare_changes_for_storage_trie(outcome)?; + stats.measure_prepare(true); + if self.incremental_account_change.is_empty() { + self.do_first_fetch(&consistent_db_view, &mut stats)?; + } + + let mut loop_break = false; + for _ in 0..10 { + stats.start(); + let ok = self.process_storage_tries_updates()?; + stats.measure_insert(true); + if !ok { + stats.start(); + self.fetch_missing_storage_nodes(&consistent_db_view, &mut stats)?; + stats.measure_proof_fetch(true); + continue; + } + stats.start(); + self.hash_storage_tries(); + stats.measure_hash(true); + loop_break = true; + break; + } + assert!(loop_break, "storage trie are not processed after 10 iters"); + + stats.start(); + self.prepare_changes_account_trie(proof_targets); + stats.measure_prepare(false); + + let mut loop_break = false; + let mut root_hash = B256::ZERO; + for _ in 0..10 { + stats.start(); + let ok = self.process_account_tries_update()?; + stats.measure_insert(false); + if !ok { + stats.start(); + self.fetch_missing_account_trie_nodes(&consistent_db_view, &mut stats)?; + stats.measure_proof_fetch(false); + continue; + } + stats.start(); + root_hash = self.hash_account_trie(&shared_cache); + stats.measure_hash(false); + loop_break = true; + break; + } + assert!(loop_break, "account trie are not processed after 10 iters"); + + let mut loop_break = false; + for _ in 0..10 { + stats.start(); + let ok = self.process_account_trie_proofs(&shared_cache)?; + stats.measure_other(); + if !ok { + stats.start(); + self.fetch_missing_account_trie_nodes(&consistent_db_view, &mut stats)?; + stats.measure_proof_fetch(false); + // if we fetched proofs we need to rehash account trie + stats.start(); + self.account_trie + .trie + .root_hash(true, &shared_cache.account_trie) + .map_err(|err| { + eyre::eyre!("failed to hash account trie (account proofs) {err:?}") + })?; + stats.measure_other(); + continue; + } + loop_break = true; + break; + } + assert!( + loop_break, + "account trie proofs are not processed after 10 iters" + ); + + let mut proofs = HashMap::default(); + for (address, proof) in self.account_trie.proof_result.drain(..) { + proofs.insert( + address, + proof.into_iter().map(|(_, node)| node.into()).collect(), + ); + } + for proof_target in proof_targets { + if !proofs.contains_key(proof_target) { + return Err(SparseTrieError::Other(eyre::eyre!( + "Proof was not fethed correctly" + ))); + } + } + + let mut metrics = SparseTrieMetrics::default(); + metrics.fetched_nodes = stats.fetched_nodes; + stats.finalize_and_print(); + Ok((root_hash, proofs, metrics)) + } +} + +#[derive(Debug, Clone, Default)] +struct Stats { + global_start: Option, + current_start: Option, + proof_fetch_db_start: Option, + + prepare_storage: f64, + prepare_account: f64, + + insert_storage: f64, + insert_account: f64, + + proof_fetch_storage: f64, + proof_fetch_account: f64, + + proof_fetch_db_part: f64, + + hash_storage: f64, + hash_account: f64, + + other: f64, + + fetched_nodes: usize, +} + +impl Stats { + fn start_global(&mut self) { + self.global_start = Some(Instant::now()); + } + + fn finalize_and_print(self) { + let elapsed = elapsed_ms(self.global_start.unwrap()); + let elapsed_no_db = elapsed - self.proof_fetch_db_part; + + let Stats { + prepare_storage, + prepare_account, + insert_storage, + insert_account, + proof_fetch_storage, + proof_fetch_account, + proof_fetch_db_part, + hash_storage, + hash_account, + other, + fetched_nodes, + .. + } = self; + + tracing::trace!( + elapsed, + elapsed_no_db, + prepare_storage, + prepare_account, + insert_storage, + insert_account, + proof_fetch_storage, + proof_fetch_account, + proof_fetch_db_part, + hash_storage, + hash_account, + other, + fetched_nodes, + "Root hash" + ); + } + + fn start(&mut self) { + self.current_start = Some(Instant::now()); + } + + fn start_proof_fetch_db(&mut self) { + self.proof_fetch_db_start = Some(Instant::now()); + } + + fn measure_prepare(&mut self, storage: bool) { + let elapsed = elapsed_ms(self.current_start.unwrap()); + if storage { + self.prepare_storage += elapsed; + } else { + self.prepare_account += elapsed; + } + } + fn measure_insert(&mut self, storage: bool) { + let elapsed = elapsed_ms(self.current_start.unwrap()); + if storage { + self.insert_storage += elapsed; + } else { + self.insert_account += elapsed; + } + } + fn measure_proof_fetch(&mut self, storage: bool) { + let elapsed = elapsed_ms(self.current_start.unwrap()); + if storage { + self.proof_fetch_storage += elapsed; + } else { + self.proof_fetch_account += elapsed; + } + } + fn measure_proof_fetch_db_part(&mut self) { + let elapsed = elapsed_ms(self.proof_fetch_db_start.unwrap()); + self.proof_fetch_db_part += elapsed; + } + fn measure_hash(&mut self, storage: bool) { + let elapsed = elapsed_ms(self.current_start.unwrap()); + if storage { + self.hash_storage += elapsed; + } else { + self.hash_account += elapsed; + } + } + fn measure_other(&mut self) { + let elapsed = elapsed_ms(self.current_start.unwrap()); + self.other += elapsed; + } +} + +fn elapsed_ms(start: Instant) -> f64 { + start.elapsed().as_nanos() as f64 / 1_000_000.0 +} diff --git a/crates/eth-sparse-mpt/src/v_experimental/trie/mod.rs b/crates/eth-sparse-mpt/src/v_experimental/trie/mod.rs new file mode 100644 index 000000000..0a25fcd18 --- /dev/null +++ b/crates/eth-sparse-mpt/src/v_experimental/trie/mod.rs @@ -0,0 +1,1193 @@ +use std::ops::Range; + +use alloy_primitives::{keccak256, B256}; +use alloy_rlp::EMPTY_STRING_CODE; +use arrayvec::ArrayVec; + +use nybbles::Nibbles; +use proof_store::{ProofNode, ProofStore}; + +pub mod proof_store; + +#[cfg(test)] +mod tests; + +use crate::utils::{encode_branch_node, encode_extension, encode_leaf, mismatch}; + +#[derive(Debug, Clone, Copy)] +enum NodePtr { + Local(usize), + Remote(usize), +} + +impl NodePtr { + #[inline] + fn is_remote(&self) -> bool { + matches!(self, Self::Remote(_)) + } + + #[inline] + fn as_local(&self) -> Option { + match self { + Self::Local(idx) => Some(*idx), + Self::Remote(_) => None, + } + } + + #[inline] + fn expect_local(&self, msg: &str) -> usize { + self.as_local() + .unwrap_or_else(|| panic!("eth-sparse-mpt expect local node: {msg}")) + } +} + +#[derive(Debug, Default)] +pub struct Trie { + // 3 arrays below are of the same length + hashed_nodes: Vec, + rlp_ptrs_local: Vec>, + nodes: Vec, + + values: Vec, + keys: Vec, + branch_node_children: Vec<[Option; 16]>, + + // scratchpad + walk_path: Vec<(usize, u8)>, // node index, nibble +} + +impl Clone for Trie { + fn clone(&self) -> Self { + fn clone_vec_with_capacity(src: &Vec) -> Vec { + let mut out = Vec::with_capacity(src.capacity()); + out.extend(src.iter().cloned()); + out + } + + Self { + hashed_nodes: clone_vec_with_capacity(&self.hashed_nodes), + rlp_ptrs_local: clone_vec_with_capacity(&self.rlp_ptrs_local), + nodes: clone_vec_with_capacity(&self.nodes), + values: clone_vec_with_capacity(&self.values), + keys: clone_vec_with_capacity(&self.keys), + branch_node_children: clone_vec_with_capacity(&self.branch_node_children), + walk_path: clone_vec_with_capacity(&self.walk_path), + } + } +} + +#[derive(Debug, Clone)] +enum DiffTrieNode { + Leaf { + key: Range, + value: Range, + }, + Extension { + key: Range, + next_node: NodePtr, + }, + Branch { + children: usize, + }, + Null, +} + +#[derive(Debug, thiserror::Error)] +pub enum DeletionError { + #[error("Deletion error: {0:?}")] + NodeNotFound(#[from] NodeNotFound), + #[error("Key node not found in the trie")] + KeyNotFound, +} + +#[derive(Debug, thiserror::Error)] +pub enum ProofError { + #[error("Proof node not found: {0:?}")] + NodeNotFound(#[from] NodeNotFound), + #[error("Trie is dirty")] + TrieIsDirty, +} + +#[derive(Debug, thiserror::Error)] +#[error("Node not found")] +pub struct NodeNotFound(pub Nibbles); + +impl NodeNotFound { + fn new(path: &[u8]) -> Self { + Self(Nibbles::from_nibbles_unchecked(path)) + } +} + +#[derive(Debug)] +pub enum InsertValue<'a> { + Value(&'a [u8]), + StoredValue(Range), +} + +#[derive(Debug, Clone)] +pub struct ProofWithValue { + pub proof: Vec<(Nibbles, Vec)>, + pub value: Option>, +} + +impl Trie { + pub fn is_uninit(&self) -> bool { + self.nodes.is_empty() + } + + pub fn new_empty() -> Self { + let mut def = Self::default(); + def.clear_empty(); + def + } + + pub fn clear_empty(&mut self) { + self.clear(); + self.push_node(DiffTrieNode::Null); + } + + fn push_node(&mut self, node: DiffTrieNode) -> NodePtr { + let idx = self.nodes.len(); + self.nodes.push(node); + self.hashed_nodes.push(false); + self.rlp_ptrs_local.push(Default::default()); + NodePtr::Local(idx) + } + + fn insert_value(&mut self, insert_value: InsertValue<'_>) -> Range { + match insert_value { + InsertValue::Value(slice) => self.copy_value(slice), + InsertValue::StoredValue(range) => range.clone(), + } + } + + fn copy_value(&mut self, value: &[u8]) -> Range { + let offset = self.values.len(); + self.values.extend_from_slice(value); + offset..offset + value.len() + } + + fn insert_key(&mut self, key: &[u8]) -> Range { + let offset = self.keys.len(); + self.keys.extend_from_slice(key); + offset..offset + key.len() + } + + fn create_branch_children(&mut self) -> usize { + let idx = self.branch_node_children.len(); + self.branch_node_children.push(Default::default()); + idx + } + + pub fn clear(&mut self) { + self.hashed_nodes.clear(); + self.rlp_ptrs_local.clear(); + self.nodes.clear(); + + self.values.clear(); + self.keys.clear(); + self.branch_node_children.clear(); + } + + // return prefix (as part of path_left, stripped nibble of suffix 1, suffix 1 as part of path_left, stripped nibble from suffix2, suffix2 as part of key stored) + fn extract_prefix_and_suffix<'a>( + &self, + path_left: &'a [u8], + key: Range, + ) -> (&'a [u8], u8, &'a [u8], u8, Range) { + let p = mismatch(path_left, &self.keys[key.clone()]); + let prefix = &path_left[..p]; + let n1 = path_left[p]; + let suff1 = &path_left[(p + 1)..]; + let n2 = self.keys[key.start + p]; + let suff2 = (key.start + p + 1)..key.end; + (prefix, n1, suff1, n2, suff2) + } + + pub fn insert( + &mut self, + key: &[u8], + insert_value: &[u8], + ) -> Result>, NodeNotFound> { + let n = Nibbles::unpack(key); + self.insert_nibble_key(&n, InsertValue::Value(insert_value)) + } + + // returns old value + pub fn insert_nibble_key( + &mut self, + nibbles_key: &Nibbles, + insert_value: InsertValue<'_>, + ) -> Result>, NodeNotFound> { + let ins_key = nibbles_key.as_slice(); + + let mut current_node = 0; + let mut path_walked = 0; + + let mut old_value = None; + + loop { + let node = self + .nodes + .get(current_node) + .ok_or_else(|| NodeNotFound(nibbles_key.clone()))?; + self.hashed_nodes[current_node] = false; + match node { + DiffTrieNode::Branch { children } => { + let children = *children; + + let n = ins_key[path_walked] as usize; + path_walked += 1; + if let Some(child_ptr) = self.branch_node_children[children][n] { + current_node = child_ptr + .as_local() + .ok_or_else(|| NodeNotFound(nibbles_key.clone()))?; + continue; + } else { + let new_leaf_key = self.insert_key(&ins_key[path_walked..]); + let leaf_value = self.insert_value(insert_value); + let leaf_ptr = self.push_node(DiffTrieNode::Leaf { + key: new_leaf_key, + value: leaf_value, + }); + self.branch_node_children[children][n] = Some(leaf_ptr); + } + } + DiffTrieNode::Extension { key, next_node } => { + let key = key.clone(); + let next_node = *next_node; + + if ins_key[path_walked..].starts_with(&self.keys[key.clone()]) { + path_walked += key.len(); + current_node = next_node + .as_local() + .ok_or_else(|| NodeNotFound(nibbles_key.clone()))?; + continue; + } + + let (prefix, n1, suff1, n2, suff2) = + self.extract_prefix_and_suffix(&ins_key[path_walked..], key); + + let has_extension_node = !prefix.is_empty(); + if has_extension_node { + let new_ext_key = self.insert_key(prefix); + // next node will be a branch node that we will push below + let ext_next_node = NodePtr::Local(self.nodes.len()); + self.nodes[current_node] = DiffTrieNode::Extension { + key: new_ext_key, + next_node: ext_next_node, + }; + }; + let branch_children = self.create_branch_children(); + let branch_node = DiffTrieNode::Branch { + children: branch_children, + }; + if has_extension_node { + self.push_node(branch_node); + } else { + self.nodes[current_node] = branch_node; + } + + let new_leaf_key = self.insert_key(suff1); + let new_leaf_value = self.insert_value(insert_value); + + let new_leaf_ptr = self.push_node(DiffTrieNode::Leaf { + key: new_leaf_key, + value: new_leaf_value, + }); + + let branch_child = if !suff2.is_empty() { + self.push_node(DiffTrieNode::Extension { + key: suff2, + next_node, + }) + } else { + next_node + }; + + self.branch_node_children[branch_children][n1 as usize] = Some(new_leaf_ptr); + self.branch_node_children[branch_children][n2 as usize] = Some(branch_child); + } + DiffTrieNode::Leaf { key, value } => { + let key = key.clone(); + let value = value.clone(); + + if self.keys[key.clone()] == ins_key[path_walked..] { + // update leaf in place + old_value = Some(value.clone()); + let new_value = self.insert_value(insert_value); + self.nodes[current_node] = DiffTrieNode::Leaf { + key: key.clone(), + value: new_value, + }; + break; + } + + let (prefix, n1, suff1, n2, suff2) = + self.extract_prefix_and_suffix(&ins_key[path_walked..], key); + + let has_extension_node = !prefix.is_empty(); + if has_extension_node { + let new_ext_key = self.insert_key(prefix); + // next node will branch node that we will push below + let ext_next_node = NodePtr::Local(self.nodes.len()); + self.nodes[current_node] = DiffTrieNode::Extension { + key: new_ext_key, + next_node: ext_next_node, + }; + }; + let branch_children = self.create_branch_children(); + let branch_node = DiffTrieNode::Branch { + children: branch_children, + }; + if has_extension_node { + self.push_node(branch_node); + } else { + self.nodes[current_node] = branch_node; + } + + let first_leaf_key = self.insert_key(suff1); + let first_leaf_value = self.insert_value(insert_value); + let first_leaf_ptr = self.push_node(DiffTrieNode::Leaf { + key: first_leaf_key, + value: first_leaf_value, + }); + + let second_leaf_key = suff2; + let second_leaf_value = value; + let second_leaf_ptr = self.push_node(DiffTrieNode::Leaf { + key: second_leaf_key, + value: second_leaf_value, + }); + + self.branch_node_children[branch_children][n1 as usize] = Some(first_leaf_ptr); + self.branch_node_children[branch_children][n2 as usize] = Some(second_leaf_ptr); + } + DiffTrieNode::Null => { + let new_leaf_key = self.insert_key(&ins_key[path_walked..]); + let new_leaf_value = self.insert_value(insert_value); + self.nodes[current_node] = DiffTrieNode::Leaf { + key: new_leaf_key, + value: new_leaf_value, + }; + } + } + break; + } + Ok(old_value) + } + + fn merge_keys(&mut self, key1: Range, nibble: u8, key2: Range) -> Range { + let new_start = self.keys.len(); + let new_len = self.keys.len() + key1.len() + key2.len() + 1; + self.keys.resize(new_len, 0); + self.keys.copy_within(key1.clone(), new_start); + self.keys[new_start + key1.len()] = nibble; + self.keys.copy_within(key2, new_start + key1.len() + 1); + new_start..new_len + } + + // returns old value + pub fn delete(&mut self, key: &[u8]) -> Result, DeletionError> { + let n = Nibbles::unpack(key); + self.delete_nibbles_key(&n) + } + + pub fn delete_nibbles_key( + &mut self, + nibbles_key: &Nibbles, + ) -> Result, DeletionError> { + let del_key = nibbles_key.as_slice(); + + let mut current_node = 0; + let mut path_walked = 0; + + self.walk_path.clear(); + + let old_value; + + loop { + let node = self + .nodes + .get(current_node) + .ok_or_else(|| NodeNotFound(nibbles_key.clone()))?; + self.hashed_nodes[current_node] = false; + match node { + DiffTrieNode::Branch { children } => { + // deleting from branch, key not found + if del_key.len() == path_walked { + return Err(DeletionError::KeyNotFound); + } + + let children = *children; + + let n = del_key[path_walked]; + self.walk_path.push((current_node, n)); + path_walked += 1; + + if let Some(child_ptr) = self.branch_node_children[children][n as usize] { + current_node = child_ptr + .as_local() + .ok_or_else(|| NodeNotFound(nibbles_key.clone()))?; + if self.branch_node_children[children] + .iter() + .filter(|c| c.is_some()) + .count() + == 2 + { + // we will have an orphan, make sure that we we have it in the trie + let (orphan_nibble, orphan_ptr) = self.branch_node_children[children] + .iter() + .enumerate() + .find(|(nib, c)| c.is_some() && *nib as u8 != n) + .unwrap(); + let orphan_ptr = orphan_ptr.unwrap(); + if orphan_ptr.is_remote() { + let mut orphan_path = Nibbles::with_capacity(path_walked); + orphan_path + .extend_from_slice_unchecked(&del_key[..(path_walked - 1)]); + orphan_path.push_unchecked(orphan_nibble as u8); + return Err(NodeNotFound(orphan_path).into()); + } + } + continue; + } else { + return Err(DeletionError::KeyNotFound); + } + } + DiffTrieNode::Extension { key, next_node } => { + let key = key.clone(); + let next_node = *next_node; + + if del_key[path_walked..].starts_with(&self.keys[key.clone()]) { + self.walk_path.push((current_node, 0)); + path_walked += key.len(); + current_node = next_node + .as_local() + .ok_or_else(|| NodeNotFound(nibbles_key.clone()))?; + continue; + } + + return Err(DeletionError::KeyNotFound); + } + DiffTrieNode::Leaf { key, value } => { + if self.keys[key.clone()] == del_key[path_walked..] { + old_value = value.clone(); + self.walk_path.push((current_node, 0)); + break; + } + return Err(DeletionError::KeyNotFound); + } + DiffTrieNode::Null => { + return Err(DeletionError::KeyNotFound); + } + } + } + + #[derive(Debug)] + enum NodeDeletionResult { + NodeDeleted, + NodeUpdated, + BranchBelowRemovedWithOneChild { + child_nibble: u8, + child_ptr: NodePtr, + }, + } + + let mut deletion_result = NodeDeletionResult::NodeDeleted; + + for (current_node, current_node_child) in self.walk_path.iter().rev() { + let current_node = *current_node; + let current_node_child = *current_node_child; + match deletion_result { + NodeDeletionResult::NodeDeleted => { + match &self.nodes[current_node] { + DiffTrieNode::Leaf { .. } => { + deletion_result = NodeDeletionResult::NodeDeleted; + } + DiffTrieNode::Branch { children } => { + let children = &mut self.branch_node_children[*children]; + let children_count = children.iter().filter(|c| c.is_some()).count(); + match children_count { + 3.. => { + children[current_node_child as usize] = None; + deletion_result = NodeDeletionResult::NodeUpdated; + } + 2 => { + children[current_node_child as usize] = None; + let (orphan_nibble, orphan_ptr) = children + .iter() + .enumerate() + .find(|(_, c)| c.is_some()) + .unwrap(); + let orphan_ptr = orphan_ptr.unwrap(); + + if orphan_ptr.is_remote() { + unreachable!("trie delete: orphan must be fetched before walking back"); + } + + deletion_result = + NodeDeletionResult::BranchBelowRemovedWithOneChild { + child_nibble: orphan_nibble as u8, + child_ptr: orphan_ptr, + }; + } + _ => unreachable!(), + } + } + _ => unreachable!(), + } + } + NodeDeletionResult::BranchBelowRemovedWithOneChild { + child_nibble: orphan_nibble, + child_ptr: orphan_ptr, + } => { + // we need to merge orphaned node and its nibble into the new parent + let orphan_ptr_idx = + orphan_ptr.expect_local("deletion walking back: orphan is not in the trie"); + let new_parent = self.nodes[current_node].clone(); + let orphaned_node = self.nodes[orphan_ptr_idx].clone(); + match (new_parent, orphaned_node) { + ( + DiffTrieNode::Extension { + key: parent_key, .. + }, + DiffTrieNode::Leaf { + key: orphan_key, + value, + }, + ) => { + // replace extension node by merging its path into leaf with child_nibble + let new_leaf_key = + self.merge_keys(parent_key, orphan_nibble, orphan_key); + self.nodes[current_node] = DiffTrieNode::Leaf { + key: new_leaf_key, + value, + } + } + ( + DiffTrieNode::Extension { + key: parent_key, .. + }, + DiffTrieNode::Extension { + key: orphan_key, + next_node: orphan_next, + }, + ) => { + // parent extension absorbs orphan extenion into itself + let new_ext_key = + self.merge_keys(parent_key, orphan_nibble, orphan_key); + self.nodes[current_node] = DiffTrieNode::Extension { + key: new_ext_key, + next_node: orphan_next, + } + } + ( + DiffTrieNode::Extension { + key: parent_key, .. + }, + DiffTrieNode::Branch { .. }, + ) => { + // parent extension eats the orphan nibble and start to point to orphan branch + // orphan branch is not changed + let new_ext_key = self.merge_keys(parent_key, orphan_nibble, 0..0); + self.nodes[current_node] = DiffTrieNode::Extension { + key: new_ext_key, + next_node: orphan_ptr, + } + } + ( + DiffTrieNode::Branch { + children: parent_children, + }, + DiffTrieNode::Leaf { + key: orphan_key, + value, + }, + ) => { + // parent branch starts to point to orphan leaf + // orphan leaf eats nibble + let new_leaf_key = self.merge_keys(0..0, orphan_nibble, orphan_key); + self.nodes[orphan_ptr_idx] = DiffTrieNode::Leaf { + key: new_leaf_key, + value, + }; + self.hashed_nodes[orphan_ptr_idx] = false; + self.branch_node_children[parent_children] + [current_node_child as usize] = Some(orphan_ptr); + } + ( + DiffTrieNode::Branch { + children: parent_children, + }, + DiffTrieNode::Extension { + key: orphan_key, + next_node, + }, + ) => { + // parent branch points to the orphan extension + // orphan extension eats nibble + let new_ext_key = self.merge_keys(0..0, orphan_nibble, orphan_key); + self.nodes[orphan_ptr_idx] = DiffTrieNode::Extension { + key: new_ext_key, + next_node, + }; + self.hashed_nodes[orphan_ptr_idx] = false; + self.branch_node_children[parent_children] + [current_node_child as usize] = Some(orphan_ptr); + } + ( + DiffTrieNode::Branch { + children: parent_children, + }, + DiffTrieNode::Branch { .. }, + ) => { + // parent branch points to new extension + // new extension node created with nibble inside pointing to orphan branch + // orphan branch is not changed + let new_ext_key = self.insert_key(&[orphan_nibble]); + let next_ext_ptr = self.push_node(DiffTrieNode::Extension { + key: new_ext_key, + next_node: orphan_ptr, + }); + self.branch_node_children[parent_children] + [current_node_child as usize] = Some(next_ext_ptr); + } + _ => unreachable!(), + } + deletion_result = NodeDeletionResult::NodeUpdated; + break; + } + NodeDeletionResult::NodeUpdated => break, + } + } + + // here we handle the case when deletion reaches the head + match deletion_result { + // updates terminated before reaching the top + NodeDeletionResult::NodeUpdated => {} + NodeDeletionResult::NodeDeleted => { + // trie is empty, insert the null node on top + self.nodes[0] = DiffTrieNode::Null; + self.hashed_nodes[0] = false; + } + // orphan becomes head + NodeDeletionResult::BranchBelowRemovedWithOneChild { + child_nibble: orphan_nibble, + child_ptr: orphan_ptr, + } => { + let orphan_ptr_idx = + orphan_ptr.expect_local("deletion reached trie top: orphan is not in the trie"); + self.hashed_nodes[0] = false; + match &self.nodes[orphan_ptr_idx] { + DiffTrieNode::Leaf { + key: orphan_key, + value, + } => { + let value = value.clone(); + // orphan leaf eats nibble + let new_leaf_key = self.merge_keys(0..0, orphan_nibble, orphan_key.clone()); + self.nodes[0] = DiffTrieNode::Leaf { + key: new_leaf_key, + value, + }; + } + DiffTrieNode::Extension { + key: orphan_key, + next_node, + } => { + // orphan extension eats nibble + let next_node = *next_node; + let new_ext_key = self.merge_keys(0..0, orphan_nibble, orphan_key.clone()); + self.nodes[0] = DiffTrieNode::Extension { + key: new_ext_key, + next_node, + } + } + DiffTrieNode::Branch { .. } => { + // new extension is created and it eats nibble + let new_ext_key = self.insert_key(&[orphan_nibble]); + self.nodes[0] = DiffTrieNode::Extension { + key: new_ext_key, + next_node: orphan_ptr, + } + } + DiffTrieNode::Null => unreachable!(), + } + } + } + + Ok(old_value) + } + + fn rlp_encode_node(&self, node_idx: usize, rlp: &mut Vec, proof_store: &ProofStore) { + rlp.clear(); + let node = self + .nodes + .get(node_idx) + .expect("rlp_encode_node: node not found") + .clone(); + match node { + DiffTrieNode::Branch { children } => { + let remote_nodes = proof_store.rlp_ptrs(); + let mut children_rlp_ptrs: [Option<&[u8]>; 16] = Default::default(); + for (nibble, child) in self.branch_node_children[children].iter().enumerate() { + match child { + Some(NodePtr::Local(idx)) => { + debug_assert!(self.hashed_nodes[*idx]); + children_rlp_ptrs[nibble] = Some(self.rlp_ptrs_local[*idx].as_slice()); + } + Some(NodePtr::Remote(idx)) => { + children_rlp_ptrs[nibble] = Some(remote_nodes[*idx].as_slice()); + } + None => {} + } + } + encode_branch_node(&children_rlp_ptrs, rlp); + } + DiffTrieNode::Extension { key, next_node } => { + let remote_nodes; + + let key = Nibbles::from_nibbles_unchecked(&self.keys[key]); + let child_rlp_ptr = match next_node { + NodePtr::Local(idx) => { + debug_assert!(self.hashed_nodes[idx]); + self.rlp_ptrs_local[idx].as_slice() + } + NodePtr::Remote(idx) => { + remote_nodes = proof_store.rlp_ptrs(); + remote_nodes[idx].as_slice() + } + }; + encode_extension(&key, child_rlp_ptr, rlp); + } + DiffTrieNode::Leaf { key, value } => { + let key = Nibbles::from_nibbles_unchecked(&self.keys[key]); + encode_leaf(&key, &self.values[value], rlp); + } + DiffTrieNode::Null => { + rlp.push(EMPTY_STRING_CODE); + } + } + } + + // children must be hashed + // NOT thread safe + // this function is unsafe to avoid borrow checker when hashing node as it updates rlp_ptrs_local and hashed_node + // 1. its not public + // 2. pulic caller is root_hash and its &mut + // 3. root_hash ensures that we don't have a race here + fn calculate_rlp_pointer_node( + &self, + node_idx: usize, + rlp: &mut Vec, + proof_store: &ProofStore, + ) { + self.rlp_encode_node(node_idx, rlp, proof_store); + let result = + unsafe { &mut *(self.rlp_ptrs_local.as_ptr().add(node_idx) as *mut ArrayVec) }; + result.clear(); + if rlp.len() < 32 { + result.try_extend_from_slice(rlp).unwrap(); + } else { + let hash = keccak256(rlp); + result.push(EMPTY_STRING_CODE + 32); + result.try_extend_from_slice(hash.as_slice()).unwrap(); + } + + let hashed_node = unsafe { &mut *(self.hashed_nodes.as_ptr().add(node_idx) as *mut bool) }; + *hashed_node = true; + } + + /// Calculates hash of the trie, when parallel is set rayon is used. + pub fn root_hash( + &self, + parallel: bool, + proof_store: &ProofStore, + ) -> Result { + let mut rlp = Vec::new(); + if self.nodes.is_empty() { + return Err(NodeNotFound(Nibbles::new())); + } + self.root_hash_node(0, &mut rlp, parallel, 0, proof_store); + self.rlp_encode_node(0, &mut rlp, proof_store); + Ok(keccak256(&rlp)) + } + + fn root_hash_node( + &self, + node_idx: usize, + rlp: &mut Vec, + parallel: bool, + depth: usize, + proof_store: &ProofStore, + ) { + if self.hashed_nodes[node_idx] { + return; + } + let node = self + .nodes + .get(node_idx) + .expect("root_hash_node: node not found"); + match node { + DiffTrieNode::Branch { children } => { + let use_rayon = parallel && depth == 0 && { + let local_children = self.branch_node_children[*children] + .iter() + .filter(|c| matches!(c, Some(NodePtr::Local(_)))) + .count(); + local_children > 1 + }; + + if !use_rayon { + for child in self.branch_node_children[*children].into_iter().flatten() { + if let NodePtr::Local(child) = child { + self.root_hash_node(child, rlp, false, depth + 1, proof_store); + } + } + } else { + rayon::scope(|scope| { + for child in self.branch_node_children[*children].into_iter().flatten() { + if let NodePtr::Local(child) = child { + scope.spawn(move |_| { + let mut rlp = Vec::new(); + self.root_hash_node( + child, + &mut rlp, + false, + depth + 1, + proof_store, + ); + }) + } + } + }); + } + self.calculate_rlp_pointer_node(node_idx, rlp, proof_store); + } + DiffTrieNode::Extension { next_node, .. } => { + if let NodePtr::Local(child) = next_node { + self.root_hash_node(*child, rlp, parallel, depth + 1, proof_store); + } + self.calculate_rlp_pointer_node(node_idx, rlp, proof_store); + } + DiffTrieNode::Null | DiffTrieNode::Leaf { .. } => { + self.calculate_rlp_pointer_node(node_idx, rlp, proof_store); + } + } + } + + pub fn get_proof( + &self, + key: &[u8], + proof_store: &ProofStore, + ) -> Result { + let n = Nibbles::unpack(key); + self.get_proof_nibbles_key(&n, proof_store) + } + + /// Generate proof for the target key. + pub fn get_proof_nibbles_key( + &self, + target_key: &Nibbles, + proof_store: &ProofStore, + ) -> Result { + let mut buf = Vec::new(); + let mut result = ProofWithValue { + proof: Vec::new(), + value: None, + }; + + let mut current_node = 0; + let mut path_walked = 0; + + loop { + let node = self + .nodes + .get(current_node) + .ok_or_else(|| NodeNotFound(target_key.clone()))?; + + if !self.hashed_nodes[current_node] { + return Err(ProofError::TrieIsDirty); + } + + self.rlp_encode_node(current_node, &mut buf, proof_store); + let current_node_path = + Nibbles::from_nibbles_unchecked(&target_key.as_slice()[..path_walked]); + result.proof.push((current_node_path, buf.clone())); + + match node { + DiffTrieNode::Branch { children } => { + if target_key.len() == path_walked { + break; + } + + let children = *children; + + let n = target_key[path_walked]; + path_walked += 1; + + if let Some(child_ptr) = self.branch_node_children[children][n as usize] { + current_node = child_ptr + .as_local() + .ok_or_else(|| NodeNotFound(target_key.clone()))?; + continue; + } + + break; + } + DiffTrieNode::Extension { key, next_node } => { + let key = key.clone(); + let next_node = *next_node; + + if target_key[path_walked..].starts_with(&self.keys[key.clone()]) { + path_walked += key.len(); + current_node = next_node + .as_local() + .ok_or_else(|| NodeNotFound(target_key.clone()))?; + continue; + } + + break; + } + DiffTrieNode::Leaf { key, value } => { + if self.keys[key.clone()] == target_key[path_walked..] { + result.value = Some(self.values[value.clone()].to_vec()); + } + break; + } + DiffTrieNode::Null => { + break; + } + } + } + + Ok(result) + } + + pub fn debug_print_node(&self, node_idx: usize) { + let node = self + .nodes + .get(node_idx) + .expect("print_node: node not found") + .clone(); + let h = alloy_primitives::hex::encode; + match node { + DiffTrieNode::Branch { children } => { + println!("{node_idx} Branch"); + println!("{}", h(self.rlp_ptrs_local[node_idx].as_slice())); + for (idx, child) in self.branch_node_children[children].into_iter().enumerate() { + if child.is_some() { + println!(" {idx} -> {child:?}"); + } + } + for child in self.branch_node_children[children].into_iter().flatten() { + if let NodePtr::Local(idx) = child { + self.debug_print_node(idx); + } + } + } + DiffTrieNode::Extension { next_node, key } => { + println!( + "{} Extension {:?} -> {:?}", + node_idx, + h(&self.keys[key]), + next_node + ); + println!("{}", h(self.rlp_ptrs_local[node_idx].as_slice())); + if let NodePtr::Local(idx) = next_node { + self.debug_print_node(idx); + } + } + DiffTrieNode::Leaf { key, value } => { + println!( + "{} Leaf {:?} : {:?}", + node_idx, + h(&self.keys[key]), + h(&self.values[value]) + ); + println!("{}", h(self.rlp_ptrs_local[node_idx].as_slice())); + } + DiffTrieNode::Null => { + println!("{node_idx} Null"); + println!("{}", h(self.rlp_ptrs_local[node_idx].as_slice())); + } + } + } + + pub fn try_add_proof_from_proof_store( + &mut self, + key: &Nibbles, + proof_store: &ProofStore, + ) -> Result { + let proof = if let Some(proof) = proof_store.proofs.get(key) { + proof + } else { + return Ok(false); + }; + + for (path, node) in proof.value() { + self.add_node_from_proof(path, node, proof_store)?; + } + + Ok(true) + } + + // node can be added only if all of its parents are actually in the trie + fn add_node_from_proof( + &mut self, + path: &Nibbles, + node: &ProofNode, + proof_store: &ProofStore, + ) -> Result<(), NodeNotFound> { + if path.is_empty() && !self.nodes.is_empty() { + return Ok(()); + } + + if self.nodes.is_empty() { + if !path.is_empty() { + return Err(NodeNotFound(Nibbles::new())); + } + match node { + ProofNode::Branch { children } => { + let branch_node_children = self.create_branch_children(); + let branch_children = &mut self.branch_node_children[branch_node_children]; + for b in 0..16 { + if let Some(child_rlp) = &children[b] { + let child_ptr = NodePtr::Remote(*child_rlp); + branch_children[b] = Some(child_ptr); + } + } + self.push_node(DiffTrieNode::Branch { + children: branch_node_children, + }); + } + ProofNode::Extension { key, child } => { + let key = &proof_store.keys_guard()[*key]; + let key = self.insert_key(key); + let next_node = NodePtr::Remote(*child); + self.push_node(DiffTrieNode::Extension { key, next_node }); + } + ProofNode::Leaf { key, value } => { + let key = &proof_store.keys_guard()[*key]; + let key = self.insert_key(key); + let value = &proof_store.values_guard()[*value]; + let value = self.copy_value(value); + self.push_node(DiffTrieNode::Leaf { key, value }); + } + ProofNode::Empty => { + self.push_node(DiffTrieNode::Null); + } + } + return Ok(()); + } + + let mut current_node = 0; + let mut path_walked = 0; + + let mut parent_ptr = None; + let mut parent_nibble = 0; + + loop { + let node = self + .nodes + .get(current_node) + .ok_or_else(|| NodeNotFound::new(&path[..path_walked]))?; + self.hashed_nodes[current_node] = false; + match node { + DiffTrieNode::Branch { children } => { + let children = *children; + + let n = path[path_walked] as usize; + path_walked += 1; + if path[path_walked..].is_empty() { + parent_ptr = self.branch_node_children[children][n]; + parent_nibble = n; + break; + } + if let Some(child_ptr) = self.branch_node_children[children][n] { + current_node = child_ptr + .as_local() + .ok_or_else(|| NodeNotFound::new(&path[..path_walked]))?; + continue; + } else { + return Err(NodeNotFound::new(&path[..path_walked])); + } + } + DiffTrieNode::Extension { key, next_node } => { + let key = key.clone(); + let next_node = *next_node; + + if path[path_walked..].starts_with(&self.keys[key.clone()]) { + path_walked += key.len(); + + if path[path_walked..].is_empty() { + parent_ptr = Some(next_node); + parent_nibble = 0; + break; + } + current_node = next_node.as_local().ok_or_else(|| { + NodeNotFound(Nibbles::from_nibbles_unchecked(&path[..path_walked])) + })?; + continue; + } + } + _ => { + // no proofs can be added here, + return Ok(()); + } + } + break; + } + + match parent_ptr { + Some(NodePtr::Remote(_)) => {} + _ => { + // node is not needed + return Ok(()); + } + }; + + let new_node = match node { + ProofNode::Leaf { key, value } => { + let key = &proof_store.keys_guard()[*key]; + let key = self.insert_key(key); + let value = &proof_store.values_guard()[*value]; + let value = self.copy_value(value); + self.push_node(DiffTrieNode::Leaf { key, value }) + } + ProofNode::Extension { key, child } => { + let key = &proof_store.keys_guard()[*key]; + let key = self.insert_key(key); + let next_node = NodePtr::Remote(*child); + self.push_node(DiffTrieNode::Extension { key, next_node }) + } + ProofNode::Branch { children } => { + let branch_node_children = self.create_branch_children(); + for b in 0..16 { + if let Some(child_rlp) = &children[b] { + let child_ptr = NodePtr::Remote(*child_rlp); + self.branch_node_children[branch_node_children][b] = Some(child_ptr); + } + } + self.push_node(DiffTrieNode::Branch { + children: branch_node_children, + }) + } + ProofNode::Empty => panic!("inserting empty to node to non empty trie"), + }; + + // give pointer to a parent + match &mut self.nodes[current_node] { + DiffTrieNode::Branch { children } => { + self.branch_node_children[*children][parent_nibble] = Some(new_node); + } + DiffTrieNode::Extension { next_node, .. } => { + *next_node = new_node; + } + _ => unreachable!(), + } + + Ok(()) + } +} diff --git a/crates/eth-sparse-mpt/src/v_experimental/trie/proof_store.rs b/crates/eth-sparse-mpt/src/v_experimental/trie/proof_store.rs new file mode 100644 index 000000000..5abca12aa --- /dev/null +++ b/crates/eth-sparse-mpt/src/v_experimental/trie/proof_store.rs @@ -0,0 +1,114 @@ +use alloy_rlp::Decodable; +use parking_lot::{lock_api::RwLockReadGuard, RawRwLock, RwLock}; +use std::sync::Arc; + +use alloy_trie::nodes::TrieNode as AlloyTrieNode; +use arrayvec::ArrayVec; +use dashmap::DashMap; +use nybbles::Nibbles; +use rustc_hash::FxBuildHasher; + +#[derive(Debug, Clone)] +pub enum ProofNode { + Leaf { key: usize, value: usize }, + Extension { key: usize, child: usize }, + Branch { children: [Option; 16] }, + Empty, +} + +#[derive(Debug, Clone, Default)] +pub struct ProofStore { + keys: Arc>>, + values: Arc>>>, + rlp_ptrs: Arc>>>, + + pub proofs: Arc, FxBuildHasher>>, +} + +impl ProofStore { + fn add_rlp_ptr(&self, data: ArrayVec) -> usize { + let mut arr = self.rlp_ptrs.write(); + let idx = arr.len(); + arr.push(data); + idx + } + + fn add_key(&self, data: Nibbles) -> usize { + let mut arr = self.keys.write(); + let idx = arr.len(); + arr.push(data); + idx + } + + fn add_value(&self, data: Vec) -> usize { + let mut arr = self.values.write(); + let idx = arr.len(); + arr.push(data.into()); + idx + } + + pub fn has_proof(&self, key: &Nibbles) -> bool { + self.proofs.contains_key(key) + } + + pub fn add_proof>( + &self, + key: Nibbles, + proof: Vec<(Nibbles, P)>, + ) -> Result<(), alloy_rlp::Error> { + if self.proofs.contains_key(&key) { + return Ok(()); + } + + let mut parsed_proof: Vec<(Nibbles, ProofNode)> = Vec::with_capacity(proof.len()); + + for (path, encoded_node) in proof { + let alloy_trie_node = AlloyTrieNode::decode(&mut encoded_node.as_ref())?; + let decoded_node = match alloy_trie_node { + AlloyTrieNode::Branch(alloy_node) => { + let mut children: [Option; 16] = Default::default(); + let mut stack_iter = alloy_node.stack.into_iter(); + for index in 0..16 { + if alloy_node.state_mask.is_bit_set(index) { + let rlp_ptr: ArrayVec = stack_iter + .next() + .expect("stack must be the same size as mask") + .as_slice() + .try_into() + .unwrap(); + children[index as usize] = Some(self.add_rlp_ptr(rlp_ptr)); + } + } + ProofNode::Branch { children } + } + AlloyTrieNode::Extension(node) => ProofNode::Extension { + key: self.add_key(node.key), + child: self.add_rlp_ptr(node.child.as_slice().try_into().unwrap()), + }, + AlloyTrieNode::Leaf(node) => ProofNode::Leaf { + key: self.add_key(node.key), + value: self.add_value(node.value), + }, + AlloyTrieNode::EmptyRoot => ProofNode::Empty, + }; + parsed_proof.push((path, decoded_node)); + } + + self.proofs.insert(key, parsed_proof); + + Ok(()) + } + + // panics if ptr is not stored in this proof store + pub fn rlp_ptrs(&self) -> RwLockReadGuard<'_, RawRwLock, Vec>> { + self.rlp_ptrs.read() + } + + pub fn keys_guard(&self) -> RwLockReadGuard<'_, RawRwLock, Vec> { + self.keys.read() + } + + pub fn values_guard(&self) -> RwLockReadGuard<'_, RawRwLock, Vec>> { + self.values.read() + } +} diff --git a/crates/eth-sparse-mpt/src/v_experimental/trie/tests.rs b/crates/eth-sparse-mpt/src/v_experimental/trie/tests.rs new file mode 100644 index 000000000..ca50155c7 --- /dev/null +++ b/crates/eth-sparse-mpt/src/v_experimental/trie/tests.rs @@ -0,0 +1,485 @@ +use std::{collections::HashMap, env}; + +use super::*; +use crate::{test_utils::reference_trie_hash_vec, utils::HashSet}; +use proptest::prelude::*; + +fn compare_impls(data: &[(Vec, Vec)]) { + let mut trie = Trie::new_empty(); + let empty_proof_store = ProofStore::default(); + for (key, value) in data { + trie.insert(key, value).unwrap(); + } + let got_hash = trie.root_hash(false, &empty_proof_store).unwrap(); + + if std::env::var("ETH_SPARSE_MPT_TEST_PRINT").is_ok() { + trie.debug_print_node(0); + } + let expected_hash = reference_trie_hash_vec(data); + assert_eq!(expected_hash, got_hash); + + for (key, expected_value) in calculate_final_trie_values(data, &[]) { + let ProofWithValue { proof, value } = trie + .get_proof(&key, &empty_proof_store) + .expect("failed to get proof"); + assert_eq!(expected_value, value, "proof value mismatch"); + let proof_hash = verify_proof(&key, proof); + assert_eq!(got_hash, proof_hash); + } +} + +fn compare_with_removals(data: &[(Vec, Vec)], remove: &[Vec]) -> eyre::Result<()> { + let empty_proof_store = ProofStore::default(); + let mut trie = Trie::new_empty(); + for (key, value) in data { + trie.insert(key, value)?; + } + + if env::var("ETH_SPARSE_MPT_TEST_PRINT").is_ok() { + println!("Trie before deletes"); + trie.debug_print_node(0); + println!(); + } + for key in remove { + trie.delete(key)?; + } + let got_hash = trie.root_hash(false, &empty_proof_store).unwrap(); + if env::var("ETH_SPARSE_MPT_TEST_PRINT").is_ok() { + println!("Trie after deletes"); + trie.debug_print_node(0); + println!(); + } + + let filtered_data: Vec<_> = data + .iter() + .filter(|(key, _)| !remove.contains(key)) + .cloned() + .collect(); + + if env::var("ETH_SPARSE_MPT_TEST_PRINT").is_ok() { + // for reference trie without any removals + println!("Trie from filtered data"); + let mut trie = Trie::new_empty(); + for (key, value) in &filtered_data { + trie.insert(key, value).unwrap(); + } + trie.root_hash(false, &empty_proof_store).unwrap(); + trie.debug_print_node(0); + println!(); + } + let expected_hash = reference_trie_hash_vec(&filtered_data); + assert_eq!(expected_hash, got_hash); + + for (key, expected_value) in calculate_final_trie_values(data, remove) { + let ProofWithValue { proof, value } = trie + .get_proof(&key, &empty_proof_store) + .expect("failed to get proof"); + assert_eq!(expected_value, value, "proof value mismatch"); + let proof_hash = verify_proof(&key, proof); + assert_eq!(got_hash, proof_hash); + } + + Ok(()) +} + +// resolves multiple inserts and removals +fn calculate_final_trie_values( + data: &[(Vec, Vec)], + remove: &[Vec], +) -> Vec<(Vec, Option>)> { + let mut result: HashMap, Option>> = HashMap::default(); + for (key, value) in data { + result.insert(key.clone(), Some(value.clone())); + } + for key in remove { + result.insert(key.clone(), None); + } + let mut result: Vec<_> = result.into_iter().collect(); + result.sort_by_key(|(k, _)| k.clone()); + result +} + +fn verify_proof(key: &[u8], proof: Vec<(Nibbles, Vec)>) -> B256 { + let nibble_key = Nibbles::unpack(key); + let proof_store = ProofStore::default(); + proof_store + .add_proof(nibble_key.clone(), proof) + .expect("failed to add proof to proof store"); + let mut trie = Trie::default(); + let found = trie + .try_add_proof_from_proof_store(&nibble_key, &proof_store) + .expect("failed to add proof to the trie"); + assert!(found, "proof was not found in proof store"); + trie.root_hash(false, &proof_store) + .expect("failed to calc root hash from proof") +} + +#[test] +fn do_empty_trie() { + compare_impls(&[]); +} + +#[test] +fn do_one_element_trie() { + let data = [(vec![1, 1], vec![0xa, 0xa])]; + compare_impls(&data); +} + +#[test] +fn do_update_leaf_node() { + let data = &[(vec![1], vec![2]), (vec![1], vec![3])]; + compare_impls(data); +} + +#[test] +fn do_insert_into_leaf_node_no_extension() { + let data = &[(vec![0x11], vec![0x0a]), (vec![0x22], vec![0x0b])]; + compare_impls(data); + + let data = &[(vec![0x22], vec![0x0b]), (vec![0x11], vec![0x0a])]; + compare_impls(data); +} +#[test] +fn do_insert_into_leaf_node_with_extension() { + let data = &[ + (vec![0x33, 0x22], vec![0x0a]), + (vec![0x33, 0x11], vec![0x0b]), + ]; + compare_impls(data); +} + +#[test] +fn do_insert_into_extension_node_no_extension_above() { + let data = &[ + (vec![0x33, 0x22], vec![0x0a]), + (vec![0x33, 0x11], vec![0x0b]), + (vec![0x44, 0x33], vec![0x0c]), + ]; + compare_impls(data); +} + +#[test] +fn do_insert_into_extension_node_with_extension_above() { + let data = &[ + (vec![0x33, 0x33, 0x22], vec![0x0a]), + (vec![0x33, 0x33, 0x11], vec![0x0b]), + (vec![0x33, 0x44, 0x33], vec![0x0c]), + ]; + compare_impls(data); +} + +#[test] +fn do_insert_into_extension_node_collapse_extension() { + let data = &[ + (vec![0x33, 0x22, 0x44], vec![0x0a]), + (vec![0x33, 0x11, 0x44], vec![0x0b]), + (vec![0x34, 0x33, 0x44], vec![0x0c]), + ]; + compare_impls(data); +} + +#[test] +fn do_insert_into_extension_node_collapse_extension_no_ext_above() { + let data = &[ + (vec![0x31, 0x11], vec![0x0a]), + (vec![0x32, 0x22], vec![0x0b]), + (vec![0x11, 0x33], vec![0x0c]), + ]; + compare_impls(data); +} + +#[test] +fn do_insert_into_branch_empty_child() { + let data = &[ + (vec![0x11], vec![0x0a]), + (vec![0x22], vec![0x0b]), + (vec![0x33], vec![0x0c]), + ]; + compare_impls(data); +} + +#[test] +fn do_insert_into_branch_leaf_child() { + let data = &[ + (vec![0x11], vec![0x0a]), + (vec![0x22], vec![0x0b]), + (vec![0x33], vec![0x0c]), + (vec![0x33], vec![0x0d]), + ]; + compare_impls(data); +} + +#[test] +fn do_remove_empty_trie_err() { + let add = &[]; + + let remove = &[vec![0x12]]; + + let _ = compare_with_removals(add, remove).unwrap_err(); +} + +#[test] +fn do_remove_leaf() { + let add = &[(vec![0x11], vec![0x0a])]; + + let remove = &[vec![0x11]]; + + compare_with_removals(add, remove).unwrap(); +} + +#[test] +fn do_remove_leaf_key_error() { + let add = &[(vec![0x11], vec![0x0a])]; + + let remove = &[vec![0x12]]; + + let _ = compare_with_removals(add, remove).unwrap_err(); +} + +#[test] +fn do_remove_extension_node_error() { + let add = &[(vec![0x11, 0x1], vec![0x0a]), (vec![0x11, 0x2], vec![0x0b])]; + + let remove = &[vec![0x12]]; + + let _ = compare_with_removals(add, remove).unwrap_err(); +} + +#[test] +fn do_remove_branch_err() { + let add = &[ + (vec![0x01, 0x10], vec![0x0a]), + (vec![0x01, 0x20], vec![0x0b]), + (vec![0x01, 0x30], vec![0x0c]), + ]; + + let remove = &[vec![0x01]]; + + let _ = compare_with_removals(add, remove).unwrap_err(); +} + +#[test] +fn do_remove_branch_leave_2_children() { + let add = &[ + (vec![0x01], vec![0x0a]), + (vec![0x02], vec![0x0b]), + (vec![0x03], vec![0x0c]), + ]; + + let remove = &[vec![0x01]]; + + compare_with_removals(add, remove).unwrap(); +} + +#[test] +fn do_remove_branch_leave_1_children_leaf_below_branch_above() { + let add = &[ + (vec![0x11], vec![0x0a]), + (vec![0x12], vec![0x0b]), + (vec![0x23], vec![0x0b]), + (vec![0x33], vec![0x0c]), + ]; + + let remove = &[vec![0x11]]; + + compare_with_removals(add, remove).unwrap(); +} + +#[test] +fn do_remove_branch_leave_1_children_branch_below_branch_above() { + let add = &[ + (vec![0x11, 0x00], vec![0x0a]), + (vec![0x12, 0x10], vec![0x0b]), + (vec![0x12, 0x20], vec![0x0b]), + (vec![0x23, 0x00], vec![0x0b]), + (vec![0x33, 0x00], vec![0x0c]), + ]; + + let remove = &[vec![0x11, 0x00]]; + + compare_with_removals(add, remove).unwrap(); +} + +#[test] +fn do_remove_branch_leave_1_children_ext_below_branch_above() { + let add = &[ + (vec![0x11, 0x00, 0x00], vec![0x0a]), + (vec![0x12, 0x10, 0x20], vec![0x0b]), + (vec![0x12, 0x10, 0x30], vec![0x0b]), + (vec![0x23, 0x00, 0x00], vec![0x0b]), + (vec![0x33, 0x00, 0x00], vec![0x0c]), + ]; + + let remove = &[vec![0x11, 0x00, 0x00]]; + + compare_with_removals(add, remove).unwrap(); +} + +#[test] +fn do_remove_branch_leave_1_children_leaf_below_ext_above() { + let add = &[(vec![0x11], vec![0x0a]), (vec![0x12], vec![0x0b])]; + + let remove = &[vec![0x11]]; + + compare_with_removals(add, remove).unwrap(); +} + +#[test] +fn do_remove_branch_leave_1_children_branch_below_ext_above() { + let add = &[ + (vec![0x11, 0x00], vec![0x0a]), + (vec![0x12, 0x10], vec![0x0b]), + (vec![0x12, 0x20], vec![0x0b]), + ]; + + let remove = &[vec![0x11, 0x00]]; + + compare_with_removals(add, remove).unwrap(); +} + +#[test] +fn do_remove_branch_leave_1_children_branch_below_null_above() { + let add = &[ + (vec![0x10], vec![0xa]), + (vec![0x23], vec![0xb]), + (vec![0x24], vec![0xc]), + ]; + + let remove = &[vec![0x10]]; + + compare_with_removals(add, remove).unwrap(); +} + +#[test] +fn do_remove_branch_leave_1_children_ext_below_null_above() { + let add = &[ + (vec![0x10, 0x00], vec![0xa]), + (vec![0x23, 0x01], vec![0xb]), + (vec![0x23, 0x02], vec![0xb]), + ]; + + let remove = &[vec![0x10, 0x00]]; + + compare_with_removals(add, remove).unwrap(); +} + +#[test] +fn do_remove_branch_leave_1_children_leaf_below_null_above() { + let add = &[(vec![0x10, 0x00], vec![0xa]), (vec![0x23, 0x01], vec![0xb])]; + + let remove = &[vec![0x10, 0x00]]; + + compare_with_removals(add, remove).unwrap(); +} + +#[test] +fn do_remove_branch_leave_1_children_ext_below_ext_above() { + let add = &[ + (vec![0x11, 0x00], vec![0x0a]), + (vec![0x12, 0x11], vec![0x0b]), + (vec![0x12, 0x12], vec![0x0b]), + ]; + + let remove = &[vec![0x11, 0x00]]; + + compare_with_removals(add, remove).unwrap(); +} + +fn compare_insert_key_hashed_map( + fill_values: &[(Vec, Vec)], + last_value: &(Vec, Vec), +) { + let mut trie = Trie::new_empty(); + let proof_store = ProofStore::default(); + + for (key, value) in fill_values { + trie.insert(key, value).expect("insert fill"); + } + trie.root_hash(true, &proof_store).expect("hash fill"); + + trie.insert(&last_value.0, &last_value.1) + .expect("insert last"); + let hash1 = trie.root_hash(true, &proof_store).expect("hash last"); + + let mut trie = Trie::new_empty(); + + for (key, value) in fill_values { + trie.insert(key, value).expect("insert fill 2"); + } + trie.insert(&last_value.0, &last_value.1) + .expect("insert last 2"); + let hash2 = trie.root_hash(true, &proof_store).expect("hash 2"); + + assert_eq!(hash1, hash2); +} + +fn compare_insert_key_reinsert( + // key, value, new value (or delete) + input: &[(Vec, Vec, Option>)], +) { + let proof_store = ProofStore::default(); + + let mut trie = Trie::new_empty(); + for (key, value, _) in input { + trie.insert(key, value).expect("insert fill"); + } + trie.root_hash(true, &proof_store).expect("hash fill"); + for (key, _, new_value) in input { + match new_value.as_ref() { + Some(value) => { + trie.insert(key, value).expect("second pass insert"); + } + None => { + trie.delete(key).expect("second pass delete"); + } + } + } + let hash1 = trie + .root_hash(true, &proof_store) + .expect("second pass hash"); + + let mut ref_trie = Trie::new_empty(); + for (key, _, new_value) in input { + if let Some(value) = new_value.as_ref() { + ref_trie.insert(key, value).expect("ref pass insert"); + } + } + let hash2 = ref_trie.root_hash(true, &proof_store).expect("ref hash"); + + assert_eq!(hash1, hash2); +} + +proptest! { + #[test] + fn proptest_random_insert_small_values(key_values in any::>()) { + let data: Vec<_> = key_values.into_iter().map(|(k, v)| (k.to_vec(), v.to_vec())).collect(); + compare_impls(&data); + } + + #[test] + fn proptest_random_insert_reinsert(key_values in any::)>>()) { + let data: Vec<_> = key_values.into_iter().map(|(k, v1, v2)| (k.to_vec(), v1.to_vec(), v2.map(|v| v.to_vec()))).collect(); + compare_insert_key_reinsert(&data); + } + + #[test] + fn proptest_random_insert_remove_any_values(key_values in any::)>>()) { + let mut keys_to_remove_set = HashSet::default(); + let mut keys_to_remove = Vec::new(); + let data: Vec<_> = key_values.into_iter().map(|((k, remove), v)| { + if remove && !keys_to_remove_set.contains(&k) { + keys_to_remove_set.insert(k); + keys_to_remove.push(k.to_vec()); + } + (k.to_vec(), v) + }).collect(); + compare_with_removals(&data, &keys_to_remove).unwrap(); + } + + #[test] + fn proptest_insert_key_hashed_map(key_values in any::)>>(), last_value in any::<([u8; 3], Vec)>()) { + let data: Vec<_> = key_values.into_iter().map(|(k, v)| (k.to_vec(), v)).collect(); + let last_value = (last_value.0.to_vec(), last_value.1); + compare_insert_key_hashed_map(&data, &last_value); + } +}