chore(l2): granulize SP1 zkVM cycles report

crates/l2/prover/src/guest_program/Cargo.toml

@@ -39,5 +39,9 @@ sp1 = ["dep:sp1-build", "dep:sp1-sdk"]
 l2 = []
 c-kzg = ["ethrex-vm/c-kzg", "ethrex-common/c-kzg"]
 
+# When this feature is enabled, the cycles used in the code inside the 
+# report_cycles! macro will be reported to stdout when running inside SP1 zkVM.
+sp1-cycles = []
+
 [lib]
 path = "./src/lib.rs"

crates/l2/prover/src/guest_program/src/execution.rs

@@ -1,9 +1,11 @@
 use crate::input::ProgramInput;
 use crate::output::ProgramOutput;
+use crate::report_cycles;
 
 use ethrex_blockchain::error::ChainError;
 use ethrex_blockchain::{
     validate_block, validate_gas_used, validate_receipts_root, validate_requests_hash,
+    validate_state_root,
 };
 use ethrex_common::types::AccountUpdate;
 use ethrex_common::types::block_execution_witness::ExecutionWitness;
@@ -17,7 +19,7 @@ use ethrex_common::{H256, types::Block};
 use ethrex_l2_common::l1_messages::L1Message;
 use ethrex_rlp::encode::RLPEncode;
 use ethrex_vm::{Evm, EvmError, GuestProgramStateWrapper, VmDatabase};
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 
 #[cfg(feature = "l2")]
 use ethrex_common::types::{
@@ -116,35 +118,162 @@ pub fn execution_program(input: ProgramInput) -> Result<ProgramOutput, Stateless
         );
     }
 
-    stateless_validation_l1(&blocks, execution_witness, elasticity_multiplier, chain_id)
+    stateless_validation_l1(blocks, execution_witness, elasticity_multiplier, chain_id)
 }
 
 pub fn stateless_validation_l1(
-    blocks: &[Block],
+    blocks: Vec<Block>,
     execution_witness: ExecutionWitness,
     elasticity_multiplier: u64,
     chain_id: u64,
 ) -> Result<ProgramOutput, StatelessExecutionError> {
-    let StatelessResult {
-        initial_state_hash,
-        final_state_hash,
-        last_block_hash,
-        non_privileged_count,
-        ..
-    } = execute_stateless(blocks, execution_witness, elasticity_multiplier, None)?;
+    let guest_program_state: GuestProgramState =
+        report_cycles("guest_program_state_initialization", || {
+            execution_witness
+                .try_into()
+                .map_err(StatelessExecutionError::GuestProgramState)
+        })?;
+
+    let mut wrapped_db = GuestProgramStateWrapper::new(guest_program_state);
+
+    let chain_config = wrapped_db.get_chain_config().map_err(|_| {
+        StatelessExecutionError::Internal("No chain config in execution witness".to_string())
+    })?;
+
+    // Hashing is an expensive operation in zkVMs, this way we avoid hashing twice
+    // (once in get_first_invalid_block_hash(), later in validate_block()).
+    report_cycles("initialize_block_header_hashes", || {
+        wrapped_db.initialize_block_header_hashes(&blocks)
+    })?;
+
+    // Validate execution witness' block hashes, except parent block hash (latest block hash).
+    report_cycles("get_first_invalid_block_hash", || {
+        if let Ok(Some(invalid_block_header)) = wrapped_db.get_first_invalid_block_hash() {
+            return Err(StatelessExecutionError::InvalidBlockHash(
+                invalid_block_header,
+            ));
+        }
+        Ok(())
+    })?;
+
+    // Validate the initial state
+    let parent_block_header = wrapped_db
+        .get_block_parent_header(
+            blocks
+                .first()
+                .ok_or(StatelessExecutionError::EmptyBatchError)?
+                .header
+                .number,
+        )
+        .map_err(StatelessExecutionError::GuestProgramState)?;
+
+    let initial_state_hash = report_cycles("state_trie_root", || {
+        wrapped_db
+            .state_trie_root()
+            .map_err(StatelessExecutionError::GuestProgramState)
+    })?;
+
+    if initial_state_hash != parent_block_header.state_root {
+        return Err(StatelessExecutionError::InvalidInitialStateTrie);
+    }
+
+    // Execute blocks
+    let mut parent_block_header = &parent_block_header;
+    let mut acc_account_updates: BTreeMap<Address, AccountUpdate> = BTreeMap::new();
+    let mut acc_receipts = Vec::new();
+    let mut non_privileged_count = 0;
+
+    for block in blocks.iter() {
+        // Validate the block
+        report_cycles("validate_block", || {
+            validate_block(
+                block,
+                parent_block_header,
+                &chain_config,
+                elasticity_multiplier,
+            )
+            .map_err(StatelessExecutionError::BlockValidationError)
+        })?;
+
+        let mut vm = report_cycles("setup_evm", || {
+            let vm = Evm::new_for_l1(wrapped_db.clone());
+            Ok::<_, StatelessExecutionError>(vm)
+        })?;
+
+        let result = report_cycles("execute_block", || {
+            vm.execute_block(block)
+                .map_err(StatelessExecutionError::EvmError)
+        })?;
+
+        let account_updates = report_cycles("get_state_transitions", || {
+            vm.get_state_transitions()
+                .map_err(StatelessExecutionError::EvmError)
+        })?;
+
+        // Update db for the next block
+        report_cycles("apply_account_updates", || {
+            wrapped_db
+                .apply_account_updates(&account_updates)
+                .map_err(StatelessExecutionError::GuestProgramState)
+        })?;
+        // Update acc_account_updates
+        for account in account_updates {
+            let address = account.address;
+            if let Some(existing) = acc_account_updates.get_mut(&address) {
+                existing.merge(account);
+            } else {
+                acc_account_updates.insert(address, account);
+            }
+        }
+
+        report_cycles("validate_gas_and_receipts", || {
+            validate_gas_used(&result.receipts, &block.header)
+                .map_err(StatelessExecutionError::GasValidationError)
+        })?;
+
+        report_cycles("validate_receipts_root", || {
+            validate_receipts_root(&block.header, &result.receipts)
+                .map_err(StatelessExecutionError::ReceiptsRootValidationError)
+        })?;
+
+        // validate_requests_hash doesn't do anything for l2 blocks as this verifies l1 requests (messages, privileged transactions and consolidations)
+        report_cycles("validate_requests_hash", || {
+            validate_requests_hash(&block.header, &chain_config, &result.requests)
+                .map_err(StatelessExecutionError::RequestsRootValidationError)
+        })?;
+
+        non_privileged_count += block.body.transactions.len();
+        acc_receipts.push(result.receipts);
+        parent_block_header = &block.header;
+    }
+
+    let final_state_root = report_cycles("get_final_state_root", || {
+        wrapped_db
+            .state_trie_root()
+            .map_err(StatelessExecutionError::GuestProgramState)
+    })?;
+
+    let last_block = blocks
+        .last()
+        .ok_or(StatelessExecutionError::EmptyBatchError)?;
+
+    report_cycles("validate_state_root", || {
+        validate_state_root(&last_block.header, final_state_root)
+            .map_err(|_chain_err| StatelessExecutionError::InvalidFinalStateTrie)
+    })?;
 
     Ok(ProgramOutput {
         initial_state_hash,
-        final_state_hash,
+        final_state_hash: final_state_root,
         #[cfg(feature = "l2")]
         l1messages_merkle_root: H256::zero(),
         #[cfg(feature = "l2")]
         privileged_transactions_hash: H256::zero(),
         #[cfg(feature = "l2")]
         blob_versioned_hash: H256::zero(),
-        last_block_hash,
+        last_block_hash: last_block.header.hash(),
         chain_id: chain_id.into(),
-        non_privileged_count,
+        non_privileged_count: non_privileged_count.into(),
     })
 }
 

crates/l2/prover/src/guest_program/src/lib.rs

@@ -16,3 +16,16 @@ pub static ZKVM_RISC0_PROGRAM_VK: &str = include_str!(concat!("./risc0/out/riscv
 // To avoid compilation errors, we override it with an empty slice.
 #[cfg(any(clippy, not(feature = "risc0")))]
 pub const ZKVM_RISC0_PROGRAM_VK: &str = "";
+
+/// Report cycles used in a code block when running inside SP1 zkVM.
+///
+/// When the feature "sp1-cycles" is enabled, it will print start and end cycle
+/// tracking messages that are compatible with SP1's cycle tracking system.
+pub fn report_cycles<T, E>(_label: &str, block: impl FnOnce() -> Result<T, E>) -> Result<T, E> {
+    #[cfg(feature = "sp1-cycles")]
+    println!("cycle-tracker-report-start: {_label}");
+    let result = block();
+    #[cfg(feature = "sp1-cycles")]
+    println!("cycle-tracker-report-end: {_label}");
+    result
+}

crates/l2/prover/src/guest_program/src/sp1/Cargo.toml

@@ -9,7 +9,7 @@ edition = "2024"
 sp1-zkvm = { version = "=5.0.8" }
 rkyv = { version = "0.8.10", features = ["std", "unaligned"] }
 
-guest_program = { path = "../../" }
+guest_program = { path = "../../", features = ["sp1-cycles"] }
 
 ethrex-common = { path = "../../../../../../common", default-features = false }
 ethrex-storage = { path = "../../../../../../storage", default-features = false }