perf: compact merkle tree for transaction placeholder proof

crates/rbuilder-primitives/Cargo.toml

@@ -65,3 +65,7 @@ ring = "0.17"
 [[bench]]
 name = "sha_pair"
 harness = false
+
+[[bench]]
+name = "ssz_proof"
+harness = false

crates/rbuilder-primitives/benches/ssz_proof.rs

@@ -0,0 +1,186 @@
+use alloy_primitives::Bytes;
+use criterion::{
+    criterion_group, criterion_main, measurement::WallTime, BenchmarkGroup, Criterion,
+};
+use proptest::{prelude::*, strategy::ValueTree as _, test_runner::TestRunner};
+
+use impls::SszTransactionProof;
+
+criterion_main!(ssz_proof);
+criterion_group!(ssz_proof, ssz_proof_bench);
+
+fn ssz_proof_bench(c: &mut Criterion) {
+    let mut group = c.benchmark_group("ssz_proof");
+
+    // Start with asserting equivalence of all implementations.
+    impls::assert_equivalence();
+
+    for num_txs in [100, 500, 1_000] {
+        let target = num_txs - 1;
+
+        for tx_size in [128, 1_024] {
+            let mut runner = TestRunner::deterministic();
+            let txs = generate_test_data(&mut runner, num_txs, tx_size);
+
+            run_bench::<impls::VanillaSszTxProof>(&mut group, &txs, target);
+            run_bench::<impls::VanillaBufferedSszTxProof>(&mut group, &txs, target);
+            run_bench::<impls::CompactSszTxProof>(&mut group, &txs, target);
+        }
+    }
+}
+
+fn run_bench<T: SszTransactionProof>(
+    group: &mut BenchmarkGroup<'_, WallTime>,
+    txs: &[Bytes],
+    target: usize,
+) {
+    let tx_size = txs.first().unwrap().0.len();
+    let id = format!(
+        "{} | num txs {} | tx size {} bytes",
+        T::description(),
+        txs.len(),
+        tx_size
+    );
+    group.bench_function(id, |b| {
+        b.iter_with_setup(
+            || T::default(),
+            |mut gen| {
+                gen.generate(txs, target);
+            },
+        )
+    });
+}
+
+fn generate_test_data(runner: &mut TestRunner, num_txs: usize, tx_size: usize) -> Vec<Bytes> {
+    proptest::collection::vec(proptest::collection::vec(any::<u8>(), tx_size), num_txs)
+        .new_tree(runner)
+        .unwrap()
+        .current()
+        .into_iter()
+        .map(Bytes::from)
+        .collect::<Vec<_>>()
+}
+
+mod impls {
+    use super::*;
+    use alloy_primitives::{Bytes, B256};
+    use rbuilder_primitives::mev_boost::ssz_roots::{
+        sha_pair, tx_ssz_leaf_root, CompactSszTransactionTree,
+    };
+
+    const TREE_DEPTH: usize = 20; // log₂(MAX_TRANSACTIONS_PER_PAYLOAD)
+    const MAX_CHUNK_COUNT: usize = 1 << TREE_DEPTH;
+
+    pub fn assert_equivalence() {
+        let num_txs = 100;
+        let proof_target = num_txs - 1;
+        let tx_size = 1_024;
+        let mut runner = TestRunner::deterministic();
+
+        let mut vanilla = VanillaSszTxProof::default();
+        let mut vanilla_buf = VanillaBufferedSszTxProof::default();
+        let mut compact = CompactSszTxProof::default();
+        for _ in 0..100 {
+            let txs = generate_test_data(&mut runner, num_txs, tx_size);
+            let expected = vanilla.generate(&txs, proof_target);
+            assert_eq!(expected, vanilla_buf.generate(&txs, proof_target));
+            assert_eq!(expected, compact.generate(&txs, proof_target));
+        }
+    }
+
+    pub trait SszTransactionProof: Default {
+        fn description() -> &'static str;
+
+        fn generate(&mut self, txs: &[Bytes], target: usize) -> Vec<B256>;
+    }
+
+    /// === VanillaSszTransactionProof ===
+    #[derive(Default)]
+    pub struct VanillaSszTxProof;
+
+    impl SszTransactionProof for VanillaSszTxProof {
+        fn description() -> &'static str {
+            "vanilla"
+        }
+
+        fn generate(&mut self, txs: &[Bytes], target: usize) -> Vec<B256> {
+            vanilla_transaction_proof_ssz(txs, target, &mut Vec::new(), &mut Vec::new())
+        }
+    }
+
+    /// === VanillaBufferedSszTransactionProof ===
+    #[derive(Default)]
+    pub struct VanillaBufferedSszTxProof {
+        current_buf: Vec<B256>,
+        next_buf: Vec<B256>,
+    }
+
+    impl SszTransactionProof for VanillaBufferedSszTxProof {
+        fn description() -> &'static str {
+            "vanilla with buffers"
+        }
+
+        fn generate(&mut self, txs: &[Bytes], target: usize) -> Vec<B256> {
+            vanilla_transaction_proof_ssz(txs, target, &mut self.current_buf, &mut self.next_buf)
+        }
+    }
+
+    fn vanilla_transaction_proof_ssz(
+        txs: &[Bytes],
+        target: usize,
+        current_buf: &mut Vec<B256>,
+        next_buf: &mut Vec<B256>,
+    ) -> Vec<B256> {
+        current_buf.clear();
+        for idx in 0..MAX_CHUNK_COUNT {
+            let leaf = txs
+                .get(idx)
+                .map(|tx| tx_ssz_leaf_root(&tx))
+                .unwrap_or(B256::ZERO);
+            current_buf.insert(idx, leaf);
+        }
+
+        let mut branch = Vec::new();
+        let (current_level, next_level) = (current_buf, next_buf);
+        let mut current_index = target;
+
+        for _level in 0..TREE_DEPTH {
+            let sibling_index = current_index ^ 1;
+            branch.push(current_level[sibling_index]);
+
+            next_level.clear();
+            for i in (0..current_level.len()).step_by(2) {
+                let left = current_level[i];
+                let right = current_level[i + 1];
+                next_level.push(sha_pair(&left, &right));
+            }
+
+            std::mem::swap(current_level, next_level);
+            current_index /= 2;
+
+            if current_level.len() == 1 {
+                break;
+            }
+        }
+
+        branch
+    }
+
+    /// === CompactSszTxProof ===
+    #[derive(Default)]
+    pub struct CompactSszTxProof;
+
+    impl SszTransactionProof for CompactSszTxProof {
+        fn description() -> &'static str {
+            "compact"
+        }
+
+        fn generate(&mut self, txs: &[Bytes], target: usize) -> Vec<B256> {
+            let mut leaves = Vec::with_capacity(txs.len());
+            for tx in txs {
+                leaves.push(tx_ssz_leaf_root(tx));
+            }
+            CompactSszTransactionTree::from_leaves(leaves).proof(target)
+        }
+    }
+}

crates/rbuilder-primitives/src/mev_boost/ssz_roots.rs

@@ -3,6 +3,7 @@
 use alloy_primitives::{Address, Bytes, B256};
 use sha2::{Digest, Sha256};
 use ssz_types::{FixedVector, VariableList};
+use std::sync::LazyLock;
 use tree_hash::TreeHash as _;
 
 #[derive(tree_hash_derive::TreeHash)]
@@ -62,74 +63,83 @@ pub fn calculate_transactions_root_ssz(transactions: &[Bytes]) -> B256 {
 
 const TREE_DEPTH: usize = 20; // log₂(MAX_TRANSACTIONS_PER_PAYLOAD)
 
-const MAX_CHUNK_COUNT: usize = 1 << TREE_DEPTH;
-
-/// Generate SSZ proof for target transaction.
-pub fn generate_transaction_proof_ssz(transactions: &[Bytes], target: usize) -> Vec<B256> {
-    generate_transaction_proof_ssz_with_buffers(
-        transactions,
-        target,
-        &mut Vec::new(),
-        &mut Vec::new(),
-    )
-}
-
-/// Generate SSZ proof for target transaction with reusable buffer.
-pub fn generate_transaction_proof_ssz_with_buffers(
-    transactions: &[Bytes],
-    target: usize,
-    current_buf: &mut Vec<B256>,
-    next_buf: &mut Vec<B256>,
-) -> Vec<B256> {
-    // Compute all leaf hashes and fill remaining slots with 0 hashes.
-    // SSZ always pads to the maximum possible size defined by the type
-    current_buf.clear();
-    for idx in 0..MAX_CHUNK_COUNT {
-        let leaf = transactions
-            .get(idx)
-            .map(ssz_leaf_root)
-            .unwrap_or(B256::ZERO);
-        current_buf.insert(idx, leaf);
+// Precompute HASHES[k] = hash of a full-zero subtree at level k.
+static ZERO_SUBTREE: LazyLock<[B256; TREE_DEPTH + 1]> = LazyLock::new(|| {
+    let mut hashes = [B256::ZERO; TREE_DEPTH + 1];
+    for lvl in 0..TREE_DEPTH {
+        hashes[lvl + 1] = sha_pair(&hashes[lvl], &hashes[lvl]);
     }
+    hashes
+});
+
+#[derive(Debug)]
+pub struct CompactSszTransactionTree(Vec<Vec<B256>>);
+
+impl CompactSszTransactionTree {
+    /// Build a compact Merkle tree over `n = txs.len()` leaves.
+    /// Level 0 = leaves; Level k has len = ceil(prev_len/2).
+    /// Padding beyond n uses structural zeros Z[k].
+    pub fn from_leaves(mut leaves: Vec<B256>) -> Self {
+        // Degenerate case: treat as single zero leaf so we still have a root
+        if leaves.is_empty() {
+            leaves.push(ZERO_SUBTREE[0]);
+        }
 
-    // Build the merkle tree bottom-up and collect the proof
-    let mut branch = Vec::new();
-    let (current_level, next_level) = (current_buf, next_buf);
-    let mut current_index = target;
-
-    // Build the complete tree to depth TREE_DEPTH (20 levels)
-    for _level in 0..TREE_DEPTH {
-        // Get the sibling at this level
-        let sibling_index = current_index ^ 1;
-        branch.push(current_level[sibling_index]);
-
-        // Build next level up
-        next_level.clear();
-        for i in (0..current_level.len()).step_by(2) {
-            let left = current_level[i];
-            let right = current_level[i + 1];
-            next_level.push(sha_pair(&left, &right));
+        // Level 0: leaves
+        let mut levels: Vec<Vec<B256>> = Vec::new();
+        levels.push(leaves);
+
+        // Upper levels
+        for level in 0..TREE_DEPTH {
+            let prev = &levels[level];
+            if prev.len() == 1 {
+                break; // reached root
+            }
+            let parents = prev.len().div_ceil(2);
+            let mut next = Vec::with_capacity(parents);
+            for i in 0..parents {
+                // NOTE: left node should always be set
+                let l = prev.get(2 * i).copied().unwrap_or(ZERO_SUBTREE[level]);
+                let r = prev.get(2 * i + 1).copied().unwrap_or(ZERO_SUBTREE[level]);
+                next.push(sha_pair(&l, &r));
+            }
+            levels.push(next);
         }
 
-        std::mem::swap(current_level, next_level);
-        current_index /= 2;
+        Self(levels)
+    }
 
-        // Stop when we reach the root
-        if current_level.len() == 1 {
-            break;
+    pub fn proof(&self, target: usize) -> Vec<B256> {
+        let mut branch = Vec::with_capacity(TREE_DEPTH);
+        for level in 0..TREE_DEPTH {
+            if level >= self.0.len() || self.0[level].len() == 1 {
+                // Either level wasn't built or compact root reached - structural zero sibling.
+                branch.push(ZERO_SUBTREE[level]);
+                continue;
+            }
+
+            let segment_index = target >> level;
+            let sibling_index = segment_index ^ 1;
+            let sibling = self.0[level]
+                .get(sibling_index)
+                .copied()
+                .unwrap_or(ZERO_SUBTREE[level]); // structural zero if beyond built range
+            branch.push(sibling);
         }
-    }
 
-    branch
+        branch
+    }
 }
 
+/// Create the leaf root for transaction bytes.
 #[inline]
-fn ssz_leaf_root(data: &Bytes) -> B256 {
+pub fn tx_ssz_leaf_root(data: &[u8]) -> B256 {
     B256::from_slice(&BinaryTransaction::from(data.to_vec()).tree_hash_root()[..])
 }
 
+/// Compute a SHA-256 hash of the pair of 32 byte hashes.
 #[inline]
-fn sha_pair(a: &B256, b: &B256) -> B256 {
+pub fn sha_pair(a: &B256, b: &B256) -> B256 {
     let mut h = Sha256::new();
     h.update(a);
     h.update(b);