`kona-executor` Extensions

use alloy_evm::{Database, EvmEnv, EvmFactory};
use alloy_op_evm::OpEvm;
use alloy_primitives::{Address, Bytes, u64_to_address};
use kona_executor::StatelessL2Builder;
use kona_genesis::RollupConfig;
use op_revm::{
    DefaultOp, OpContext, OpEvm as RevmOpEvm, OpHaltReason, OpSpecId, OpTransaction,
    OpTransactionError,
};
use revm::{
    Context, Inspector,
    context::{Evm as RevmEvm, FrameStack, TxEnv, result::EVMError},
    handler::instructions::EthInstructions,
    inspector::NoOpInspector,
    precompile::{PrecompileResult, PrecompileOutput, Precompiles},
};

const MY_PRECOMPILE_ADDRESS: Address = u64_to_address(0xFF);

fn my_precompile(input: &Bytes, gas_limit: u64) -> PrecompileResult {
   Ok(PrecompileOutput::new(50, "hello, world!".as_bytes().into()))
}

#[derive(Debug, Clone)]
pub struct CustomEvmFactory;

impl EvmFactory for CustomEvmFactory {
    type Evm<DB: Database, I: Inspector<OpContext<DB>>> = OpEvm<DB, I, CustomPrecompiles>;
    type Context<DB: Database> = OpContext<DB>;
    type Tx = OpTransaction<TxEnv>;
    type Error<DBError: core::error::Error + Send + Sync + 'static> =
        EVMError<DBError, OpTransactionError>;
    type HaltReason = OpHaltReason;
    type Spec = OpSpecId;
    type Precompiles = CustomPrecompiles;

    fn create_evm<DB: Database>(
        &self,
        db: DB,
        input: EvmEnv<OpSpecId>,
    ) -> Self::Evm<DB, NoOpInspector> {
        let spec_id = *input.spec_id();
        let ctx = Context::op().with_db(db).with_block(input.block_env).with_cfg(input.cfg_env);
        
        // Create custom precompiles with our added precompile
        let mut precompiles = op_revm::precompiles::granite();
        precompiles.insert(MY_PRECOMPILE_ADDRESS, my_precompile);
        let custom_precompiles = CustomPrecompiles { precompiles };

        let revm_evm = RevmOpEvm(RevmEvm {
            ctx,
            inspector: NoOpInspector {},
            instruction: EthInstructions::new_mainnet(),
            precompiles: custom_precompiles,
            frame_stack: FrameStack::new(),
        });

        OpEvm::new(revm_evm, false)
    }

    fn create_evm_with_inspector<DB: Database, I: Inspector<Self::Context<DB>>>(
        &self,
        db: DB,
        input: EvmEnv<OpSpecId>,
        inspector: I,
    ) -> Self::Evm<DB, I> {
        let spec_id = *input.spec_id();
        let ctx = Context::op().with_db(db).with_block(input.block_env).with_cfg(input.cfg_env);
        
        // Create custom precompiles with our added precompile
        let mut precompiles = op_revm::precompiles::granite();
        precompiles.insert(MY_PRECOMPILE_ADDRESS, my_precompile);
        let custom_precompiles = CustomPrecompiles { precompiles };

        let revm_evm = RevmOpEvm(RevmEvm {
            ctx,
            inspector,
            instruction: EthInstructions::new_mainnet(),
            precompiles: custom_precompiles,
            frame_stack: FrameStack::new(),
        });

        OpEvm::new(revm_evm, true)
    }
}

// Custom precompiles wrapper
#[derive(Debug)]
pub struct CustomPrecompiles {
    precompiles: Precompiles,
}

impl<CTX> revm::handler::PrecompileProvider<CTX> for CustomPrecompiles
where
    CTX: revm::context::ContextTr<Cfg: revm::context::Cfg<Spec = OpSpecId>>,
{
    type Output = revm::interpreter::InterpreterResult;

    fn set_spec(&mut self, spec: OpSpecId) -> bool {
        // Update precompiles based on spec if needed
        false
    }

    fn run(
        &mut self,
        _context: &mut CTX,
        address: &Address,
        inputs: &revm::interpreter::InputsImpl,
        _is_static: bool,
        gas_limit: u64,
    ) -> Result<Option<Self::Output>, String> {
        use revm::interpreter::{Gas, InstructionResult, InterpreterResult};
        
        let input = match &inputs.input {
            revm::interpreter::CallInput::Bytes(bytes) => bytes.clone(),
            revm::interpreter::CallInput::SharedBuffer(range) => {
                // Handle shared buffer case - simplified for example
                Bytes::new()
            }
        };

        if let Some(precompile) = self.precompiles.get(address) {
            let result = (*precompile)(&input, gas_limit);
            match result {
                Ok(output) => Ok(Some(InterpreterResult {
                    result: InstructionResult::Return,
                    gas: Gas::new(gas_limit - output.gas_used),
                    output: output.bytes,
                })),
                Err(_) => Ok(Some(InterpreterResult {
                    result: InstructionResult::PrecompileError,
                    gas: Gas::new(0),
                    output: Bytes::new(),
                })),
            }
        } else {
            Ok(None)
        }
    }
}

// - snip -

let cfg = RollupConfig::default();
let provider = ...;
let hinter = ...;
let parent_header = ...;

let executor = StatelessL2Builder::new(
    &cfg,
    CustomEvmFactory,
    provider,
    hinter,
    parent_header,
);