executor.rs

use std::sync::Arc;

use alloy_consensus::{BlockHeader, Header};
use alloy_primitives::B256;
use itertools::Itertools;
use primitives::is_goat_testnet;
use reth_chainspec::ChainSpec;
use reth_errors::BlockExecutionError;
use reth_evm::{
    execute::{BasicBlockExecutor, Executor},
    ConfigureEvm, OnStateHook,
};
use reth_evm_ethereum::EthEvmConfig;
use reth_execution_types::ExecutionOutcome;
use reth_primitives_traits::Block;
use reth_trie::{KeccakKeyHasher, TrieAccount, EMPTY_ROOT_HASH};
use revm::{database::WrapDatabaseRef, install_crypto, DatabaseRef};
use revm_primitives::{Address, HashMap, U256};

use crate::{
    custom::{CustomCrypto, CustomEvmFactory},
    error::ClientError,
    into_primitives::FromInput,
    io::{ClientExecutorInput, TrieDB, WitnessInput},
    tracking::OpCodesTrackingBlockExecutor,
    BlockValidator,
};

pub const DESERIALZE_INPUTS: &str = "deserialize inputs";
pub const INIT_WITNESS_DB: &str = "initialize witness db";
pub const RECOVER_SENDERS: &str = "recover senders";
pub const BLOCK_EXECUTION: &str = "block execution";
pub const VALIDATE_HEADER: &str = "validate header";
pub const VALIDATE_EXECUTION: &str = "validate block post-execution";
pub const COMPUTE_STATE_ROOT: &str = "compute state root";
pub const CHECK_SLOT_AND_VALUE: &str = "check slot and value";

pub type EthClientExecutor = ClientExecutor<EthEvmConfig<ChainSpec, CustomEvmFactory>, ChainSpec>;

#[cfg(feature = "optimism")]
pub type OpClientExecutor =
    ClientExecutor<reth_optimism_evm::OpEvmConfig, reth_optimism_chainspec::OpChainSpec>;

/// An executor that executes a block inside a zkVM.
#[derive(Debug, Clone)]
pub struct ClientExecutor<C: ConfigureEvm, CS> {
    evm_config: C,
    chain_spec: Arc<CS>,
}

impl<C, CS> ClientExecutor<C, CS>
where
    C: ConfigureEvm,
    C::Primitives: FromInput + BlockValidator<CS>,
{
    pub fn execute(
        &self,
        mut input: ClientExecutorInput<C::Primitives>,
        storage_info: Vec<(Address, U256, U256)>,
    ) -> Result<(Header, B256), ClientError> {
        let chain_id: u64 = (&input.genesis).try_into().expect("convert chain id err");

        let sealed_headers = input.sealed_headers().collect::<Vec<_>>();

        // Initialize the witnessed database with verified storage proofs.
        let db = profile_report!(INIT_WITNESS_DB, {
            let trie_db = input.witness_db(&sealed_headers).unwrap();
            WrapDatabaseRef(trie_db)
        });

        let block_executor: BlockExecutor<'_, C> =
            BlockExecutor::new(self.evm_config.clone(), db, input.opcode_tracking, chain_id);

        let block = profile_report!(RECOVER_SENDERS, {
            C::Primitives::from_input_block(input.current_block.clone())
                .try_into_recovered()
                .map_err(|_| ClientError::SignatureRecoveryFailed)
        })?;

        // Validate the blocks.
        profile_report!(VALIDATE_HEADER, {
            C::Primitives::validate_block(&block, self.chain_spec.clone())
                .expect("The block is invalid");

            for (header, parent) in sealed_headers.iter().tuple_windows() {
                C::Primitives::validate_header(parent, self.chain_spec.clone())
                    .expect("A parent header is invalid");

                C::Primitives::validate_header_against_parent(
                    header,
                    parent,
                    self.chain_spec.clone(),
                )
                .expect("The header is invalid against its parent");
            }
        });

        let execution_output =
            profile_report!(BLOCK_EXECUTION, { block_executor.execute(&block) })?;

        // Validate the block post execution.
        profile_report!(VALIDATE_EXECUTION, {
            C::Primitives::validate_block_post_execution(
                &block,
                self.chain_spec.clone(),
                &execution_output,
                is_goat_testnet(chain_id),
            )
        })?;

        // Convert the output to an execution outcome.
        let executor_outcome = ExecutionOutcome::new(
            execution_output.state,
            vec![execution_output.result.receipts],
            input.current_block.header().number(),
            vec![execution_output.result.requests],
        );

        let parent_state_root = input.parent_state.state_root();

        // Verify the state root.
        let state_root = profile_report!(COMPUTE_STATE_ROOT, {
            input.parent_state.update(&executor_outcome.hash_state_slow::<KeccakKeyHasher>());
            input.parent_state.state_root()
        });

        if state_root != input.current_block.header().state_root() {
            return Err(ClientError::MismatchedStateRoot);
        }

        // Derive the block header.
        // Note: the receipts root and gas used are verified by `validate_block_post_execution`.
        let header = Header {
            parent_hash: input.current_block.header().parent_hash(),
            ommers_hash: input.current_block.header().ommers_hash(),
            beneficiary: input.current_block.header().beneficiary(),
            state_root,
            transactions_root: input.current_block.header().transactions_root(),
            receipts_root: input.current_block.header().receipts_root(),
            logs_bloom: input.current_block.logs_bloom,
            difficulty: input.current_block.header().difficulty(),
            number: input.current_block.header().number(),
            gas_limit: input.current_block.header().gas_limit(),
            gas_used: input.current_block.header().gas_used(),
            timestamp: input.current_block.header().timestamp(),
            extra_data: input.current_block.header().extra_data().clone(),
            mix_hash: input.current_block.header().mix_hash().unwrap(),
            nonce: input.current_block.header().nonce().unwrap(),
            base_fee_per_gas: input.current_block.header().base_fee_per_gas(),
            withdrawals_root: input.current_block.header().withdrawals_root(),
            blob_gas_used: input.current_block.header().blob_gas_used(),
            excess_blob_gas: input.current_block.header().excess_blob_gas(),
            parent_beacon_block_root: input.current_block.header().parent_beacon_block_root(),
            requests_hash: input.current_block.header().requests_hash(),
        };

        if !storage_info.is_empty() {
            let check_result: Result<(), ClientError> = profile_report!(CHECK_SLOT_AND_VALUE, {
                let state = input.state();
                let db = {
                    for (hashed_address, storage_trie) in state.storage_tries.iter() {
                        let account = state
                            .state_trie
                            .get_rlp::<TrieAccount>(hashed_address.as_slice())
                            .unwrap();
                        let storage_root = account.map_or(EMPTY_ROOT_HASH, |a| a.storage_root);
                        if storage_root != storage_trie.hash() {
                            return Err(ClientError::MismatchedStorageRoot);
                        }
                    }

                    let bytecodes_by_hash = input
                        .bytecodes()
                        .map(|code| (code.hash_slow(), code))
                        .collect::<HashMap<_, _>>();

                    // Verify and build block hashes
                    let mut block_hashes: HashMap<u64, B256> =
                        HashMap::with_hasher(Default::default());
                    for (child_header, parent_header) in input.sealed_headers().tuple_windows() {
                        if parent_header.number() != child_header.number() - 1 {
                            return Err(ClientError::InvalidHeaderBlockNumber(
                                parent_header.number() + 1,
                                child_header.number(),
                            ));
                        }

                        let parent_header_hash = parent_header.hash_slow();
                        if parent_header_hash != child_header.parent_hash() {
                            return Err(ClientError::InvalidHeaderParentHash(
                                parent_header_hash,
                                child_header.parent_hash(),
                            ));
                        }

                        block_hashes.insert(parent_header.number(), child_header.parent_hash());
                    }

                    TrieDB::new(state, block_hashes, bytecodes_by_hash)
                };

                for (contract_address, slot_id, expected_value) in storage_info {
                    match db.storage_ref(contract_address, slot_id) {
                        Ok(actual_value) => {
                            if actual_value != expected_value {
                                return Err(ClientError::FailedToCheckSlotAndValue(slot_id));
                            }
                        }
                        _ => {
                            return Err(ClientError::FailedToFetchSlotAndValue(slot_id));
                        }
                    }
                }
                Ok(())
            });
            if check_result.is_err() {
                return Err(check_result.err().unwrap());
            }
        }
        Ok((header, parent_state_root))
    }
}

impl EthClientExecutor {
    pub fn eth(chain_spec: Arc<ChainSpec>, custom_beneficiary: Option<Address>) -> Self {
        install_crypto(CustomCrypto::default());

        Self {
            evm_config: EthEvmConfig::new_with_evm_factory(
                chain_spec.clone(),
                CustomEvmFactory::new(custom_beneficiary),
            ),
            chain_spec,
        }
    }
}

#[cfg(feature = "optimism")]
impl OpClientExecutor {
    pub fn optimism(chain_spec: Arc<reth_optimism_chainspec::OpChainSpec>) -> Self {
        install_crypto(CustomCrypto::default());

        Self {
            evm_config: reth_optimism_evm::OpEvmConfig::optimism(chain_spec.clone()),
            chain_spec,
        }
    }
}

enum BlockExecutor<'a, C> {
    Basic(BasicBlockExecutor<C, WrapDatabaseRef<TrieDB<'a>>>),
    OpcodeTracking(OpCodesTrackingBlockExecutor<C, WrapDatabaseRef<TrieDB<'a>>>),
}

impl<'a, C: ConfigureEvm> BlockExecutor<'a, C> {
    fn new(
        strategy_factory: C,
        db: WrapDatabaseRef<TrieDB<'a>>,
        opcode_tracking: bool,
        chain_id: u64,
    ) -> Self {
        if opcode_tracking {
            Self::OpcodeTracking(OpCodesTrackingBlockExecutor::new(strategy_factory, db))
        } else {
            Self::Basic(BasicBlockExecutor::new(strategy_factory, db, Some(chain_id)))
        }
    }
}

impl<'a, C> Executor<WrapDatabaseRef<TrieDB<'a>>> for BlockExecutor<'a, C>
where
    C: ConfigureEvm,
{
    type Primitives = C::Primitives;
    type Error = BlockExecutionError;

    fn execute_one(
        &mut self,
        block: &reth_primitives_traits::RecoveredBlock<
            <Self::Primitives as reth_primitives_traits::NodePrimitives>::Block,
        >,
    ) -> Result<
        reth_execution_types::BlockExecutionResult<
            <Self::Primitives as reth_primitives_traits::NodePrimitives>::Receipt,
        >,
        Self::Error,
    > {
        match self {
            BlockExecutor::Basic(basic_block_executor) => basic_block_executor.execute_one(block),
            BlockExecutor::OpcodeTracking(op_codes_tracking_block_executor) => {
                op_codes_tracking_block_executor.execute_one(block)
            }
        }
    }

    fn execute_one_with_state_hook<H>(
        &mut self,
        block: &reth_primitives_traits::RecoveredBlock<
            <Self::Primitives as reth_primitives_traits::NodePrimitives>::Block,
        >,
        state_hook: H,
    ) -> Result<
        reth_execution_types::BlockExecutionResult<
            <Self::Primitives as reth_primitives_traits::NodePrimitives>::Receipt,
        >,
        Self::Error,
    >
    where
        H: OnStateHook + 'static,
    {
        match self {
            BlockExecutor::Basic(basic_block_executor) => {
                basic_block_executor.execute_one_with_state_hook(block, state_hook)
            }
            BlockExecutor::OpcodeTracking(op_codes_tracking_block_executor) => {
                op_codes_tracking_block_executor.execute_one_with_state_hook(block, state_hook)
            }
        }
    }

    fn into_state(self) -> revm::database::State<WrapDatabaseRef<TrieDB<'a>>> {
        match self {
            BlockExecutor::Basic(basic_block_executor) => basic_block_executor.into_state(),
            BlockExecutor::OpcodeTracking(op_codes_tracking_block_executor) => {
                op_codes_tracking_block_executor.into_state()
            }
        }
    }

    fn size_hint(&self) -> usize {
        match self {
            BlockExecutor::Basic(basic_block_executor) => basic_block_executor.size_hint(),
            BlockExecutor::OpcodeTracking(op_codes_tracking_block_executor) => {
                op_codes_tracking_block_executor.size_hint()
            }
        }
    }
}