perf(katana): optimisations + benchmark setup

crates/katana/core/Cargo.toml

@@ -11,7 +11,7 @@ katana-chain-spec.workspace = true
 katana-db.workspace = true
 katana-executor = { workspace = true, features = [ "blockifier" ] }
 katana-pool.workspace = true
-katana-primitives.workspace = true
+katana-primitives = { workspace = true, features = [ "arbitrary" ]}
 katana-provider.workspace = true
 katana-tasks.workspace = true
 katana-trie.workspace = true
@@ -26,6 +26,7 @@ metrics.workspace = true
 num-traits.workspace = true
 parking_lot.workspace = true
 reqwest.workspace = true
+rayon.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 starknet.workspace = true
@@ -47,8 +48,15 @@ alloy-transport = { workspace = true, default-features = false }
 
 [dev-dependencies]
 assert_matches.workspace = true
+criterion.workspace = true
 hex.workspace = true
 tempfile.workspace = true
+arbitrary.workspace = true
+rand.workspace = true
 
 [features]
 starknet-messaging = [ "dep:starknet-crypto" ]
+
+[[bench]]
+name = "commit"
+harness = false

crates/katana/core/benches/commit.rs

@@ -0,0 +1,130 @@
+use std::collections::BTreeMap;
+
+use arbitrary::{Arbitrary, Unstructured};
+use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion};
+use katana_core::backend::UncommittedBlock;
+use katana_primitives::block::{GasPrices, PartialHeader};
+use katana_primitives::da::L1DataAvailabilityMode;
+use katana_primitives::receipt::{Receipt, ReceiptWithTxHash};
+use katana_primitives::state::StateUpdates;
+use katana_primitives::transaction::{Tx, TxWithHash};
+use katana_primitives::version::CURRENT_STARKNET_VERSION;
+use katana_primitives::{ContractAddress, Felt};
+use katana_provider::providers::db::DbProvider;
+
+const NB_OF_TXS: usize = 20;
+const NB_OF_RECEIPTS: usize = 20;
+const NB_OF_NONCES: usize = 100;
+const NB_OF_STORAGE_KEYS: usize = 100;
+const NB_OF_STORAGE_VALUES: usize = 100;
+const NB_OF_CLASSES: usize = 100;
+const NB_OF_CONTRACTS: usize = 100;
+
+pub fn commit(block: UncommittedBlock<'_, DbProvider>) {
+    let _ = block.commit();
+}
+
+pub fn commit_parallel(block: UncommittedBlock<'_, DbProvider>) {
+    let _ = block.commit_parallel();
+}
+
+#[inline(always)]
+pub fn random_array(size: usize) -> Vec<u8> {
+    (0..size).map(|_| rand::random::<u8>()).collect()
+}
+
+#[inline(always)]
+pub fn random_felt() -> Felt {
+    Felt::arbitrary(&mut Unstructured::new(&random_array(Felt::size_hint(0).0))).unwrap()
+}
+
+#[inline(always)]
+pub fn random_tx() -> Tx {
+    Tx::arbitrary(&mut Unstructured::new(&random_array(Tx::size_hint(0).0))).unwrap()
+}
+
+#[inline(always)]
+pub fn random_tx_with_hash() -> TxWithHash {
+    TxWithHash { hash: random_felt(), transaction: random_tx() }
+}
+
+#[inline(always)]
+pub fn random_receipt() -> Receipt {
+    Receipt::arbitrary(&mut Unstructured::new(&random_array(Receipt::size_hint(0).0))).unwrap()
+}
+
+#[inline(always)]
+pub fn random_receipt_with_hash() -> ReceiptWithTxHash {
+    ReceiptWithTxHash { tx_hash: random_felt(), receipt: random_receipt() }
+}
+
+#[inline(always)]
+pub fn random_felt_to_felt_map(size: usize) -> BTreeMap<Felt, Felt> {
+    (0..size).map(|_| (random_felt(), random_felt())).collect()
+}
+
+#[inline(always)]
+pub fn random_address_to_felt_map(size: usize) -> BTreeMap<ContractAddress, Felt> {
+    (0..size).map(|_| (ContractAddress::new(random_felt()), random_felt())).collect()
+}
+
+pub fn commit_benchmark(c: &mut Criterion) {
+    let provider = DbProvider::new_ephemeral();
+
+    let gas_prices = GasPrices { eth: 100 * u128::pow(10, 9), strk: 100 * u128::pow(10, 9) };
+    let sequencer_address = ContractAddress(1u64.into());
+
+    let header = PartialHeader {
+        protocol_version: CURRENT_STARKNET_VERSION,
+        number: 1,
+        timestamp: 100,
+        sequencer_address,
+        parent_hash: 123u64.into(),
+        l1_gas_prices: gas_prices.clone(),
+        l1_data_gas_prices: gas_prices.clone(),
+        l1_da_mode: L1DataAvailabilityMode::Calldata,
+    };
+
+    let transactions: Vec<TxWithHash> = (0..NB_OF_TXS).map(|_| random_tx_with_hash()).collect();
+    let receipts: Vec<ReceiptWithTxHash> =
+        (0..NB_OF_RECEIPTS).map(|_| random_receipt_with_hash()).collect();
+
+    let nonce_updates: BTreeMap<ContractAddress, Felt> =
+        (0..NB_OF_NONCES).map(|_| (ContractAddress::new(random_felt()), random_felt())).collect();
+
+    let storage_updates: BTreeMap<ContractAddress, BTreeMap<Felt, Felt>> = (0..NB_OF_STORAGE_KEYS)
+        .map(|_| {
+            (ContractAddress::new(random_felt()), random_felt_to_felt_map(NB_OF_STORAGE_VALUES))
+        })
+        .collect();
+
+    let declared_classes: BTreeMap<Felt, Felt> = random_felt_to_felt_map(NB_OF_CLASSES);
+    let deployed_contracts: BTreeMap<ContractAddress, Felt> =
+        random_address_to_felt_map(NB_OF_CONTRACTS);
+
+    let state_updates = StateUpdates {
+        nonce_updates,
+        storage_updates,
+        declared_classes,
+        deployed_contracts,
+        ..Default::default()
+    };
+
+    let block =
+        UncommittedBlock::new(header, transactions, receipts.as_slice(), &state_updates, provider);
+
+    c.bench_function("commit", |b| {
+        b.iter_batched(|| block.clone(), |input| commit(black_box(input)), BatchSize::SmallInput);
+    });
+
+    c.bench_function("commit_parallel", |b| {
+        b.iter_batched(
+            || block.clone(),
+            |input| commit_parallel(black_box(input)),
+            BatchSize::SmallInput,
+        );
+    });
+}
+
+criterion_group!(benches, commit_benchmark);
+criterion_main!(benches);

crates/katana/core/src/backend/mod.rs

@@ -18,6 +18,7 @@
 use katana_provider::traits::trie::TrieWriter;
 use katana_trie::compute_merkle_root;
 use parking_lot::RwLock;
+use rayon::prelude::*;
 use starknet::macros::short_string;
 use starknet_types_core::hash::{self, StarkHash};
 use tracing::info;
@@ -221,6 +222,7 @@
         let transaction_count = self.transactions.len() as u32;
         let state_diff_length = self.state_updates.len() as u32;
 
+        // optimisation 1
         let state_root = self.compute_new_state_root();
         let transactions_commitment = self.compute_transaction_commitment();
         let events_commitment = self.compute_event_commitment();
@@ -251,6 +253,51 @@
         SealedBlock { hash, header, body: self.transactions }
     }
 
+    pub fn commit_parallel(self) -> SealedBlock {
+        // get the hash of the latest committed block
+        let parent_hash = self.header.parent_hash;
+        let events_count = self.receipts.iter().map(|r| r.events().len() as u32).sum::<u32>();
+        let transaction_count = self.transactions.len() as u32;
+        let state_diff_length = self.state_updates.len() as u32;
+
+        let mut state_root = Felt::default();
+        let mut transactions_commitment = Felt::default();
+        let mut events_commitment = Felt::default();
+        let mut receipts_commitment = Felt::default();
+        let mut state_diff_commitment = Felt::default();
+
+        rayon::scope(|s| {
+            s.spawn(|_| state_root = self.compute_new_state_root());
+            s.spawn(|_| transactions_commitment = self.compute_transaction_commitment());
+            s.spawn(|_| events_commitment = self.compute_event_commitment_parallel());
+            s.spawn(|_| receipts_commitment = self.compute_receipt_commitment_parallel());
+            s.spawn(|_| state_diff_commitment = self.compute_state_diff_commitment());
+        });
+
+        let header = Header {
+            state_root,
+            parent_hash,
+            events_count,
+            state_diff_length,
+            transaction_count,
+            events_commitment,
+            receipts_commitment,
+            state_diff_commitment,
+            transactions_commitment,
+            number: self.header.number,
+            timestamp: self.header.timestamp,
+            l1_da_mode: self.header.l1_da_mode,
+            l1_gas_prices: self.header.l1_gas_prices,
+            l1_data_gas_prices: self.header.l1_data_gas_prices,
+            sequencer_address: self.header.sequencer_address,
+            protocol_version: self.header.protocol_version,
+        };
+
+        let hash = header.compute_hash();
+
+        SealedBlock { hash, header, body: self.transactions }
+    }
+
     fn compute_transaction_commitment(&self) -> Felt {
         let tx_hashes = self.transactions.iter().map(|t| t.hash).collect::<Vec<TxHash>>();
         compute_merkle_root::<hash::Poseidon>(&tx_hashes).unwrap()
@@ -261,6 +308,12 @@
         compute_merkle_root::<hash::Poseidon>(&receipt_hashes).unwrap()
     }
 
+    fn compute_receipt_commitment_parallel(&self) -> Felt {
+        let receipt_hashes =
+            self.receipts.par_iter().map(|r| r.compute_hash()).collect::<Vec<Felt>>();
+        compute_merkle_root::<hash::Poseidon>(&receipt_hashes).unwrap()
+    }
+
     fn compute_state_diff_commitment(&self) -> Felt {
         compute_state_diff_hash(self.state_updates.clone())
     }
@@ -276,7 +329,6 @@
         // the iterator will yield all events from all the receipts, each one paired with the
         // transaction hash that emitted it: (tx hash, event).
         let events = self.receipts.iter().flat_map(|r| r.events().iter().map(|e| (r.tx_hash, e)));
-
         let mut hashes = Vec::new();
         for (tx, event) in events {
             let event_hash = event_hash(tx, event);
@@ -287,6 +339,27 @@
         compute_merkle_root::<hash::Poseidon>(&hashes).unwrap()
     }
 
+    fn compute_event_commitment_parallel(&self) -> Felt {
+        // h(emitter_address, tx_hash, h(keys), h(data))
+        fn event_hash(tx: TxHash, event: &Event) -> Felt {
+            let keys_hash = hash::Poseidon::hash_array(&event.keys);
+            let data_hash = hash::Poseidon::hash_array(&event.data);
+            hash::Poseidon::hash_array(&[tx, event.from_address.into(), keys_hash, data_hash])
+        }
+
+        // the iterator will yield all events from all the receipts, each one paired with the
+        // transaction hash that emitted it: (tx hash, event).
+        let events = self.receipts.iter().flat_map(|r| r.events().iter().map(|e| (r.tx_hash, e)));
+        let hashes = events
+            .par_bridge()
+            .into_par_iter()
+            .map(|(tx, event)| event_hash(tx, event))
+            .collect::<Vec<_>>();
+
+        // compute events commitment
+        compute_merkle_root::<hash::Poseidon>(&hashes).unwrap()
+    }
+
     // state_commitment = hPos("STARKNET_STATE_V0", contract_trie_root, class_trie_root)
     fn compute_new_state_root(&self) -> Felt {
         let class_trie_root = self