refactor: inject call to CREATE2 factory through custom revm handler (foundry-rs#7653)

* wip

* wip

* add docs

* clippy

* update doc

* simplify logic

* review fixes

* doc

* review fixes

* fix

Author: klkvr

crates/anvil/src/eth/backend/mem/inspector.rs

@@ -6,13 +6,14 @@ use foundry_evm::{
     call_inspectors,
     decode::decode_console_logs,
     inspectors::{LogCollector, TracingInspector},
-    revm,
     revm::{
+        self,
         interpreter::{CallInputs, CallOutcome, CreateInputs, CreateOutcome, Interpreter},
         primitives::U256,
         EvmContext,
     },
     traces::TracingInspectorConfig,
+    InspectorExt,
 };
 
 /// The [`revm::Inspector`] used when transacting in the evm
@@ -136,6 +137,8 @@ impl<DB: Database> revm::Inspector<DB> for Inspector {
     }
 }
 
+impl<DB: Database> InspectorExt<DB> for Inspector {}
+
 /// Prints all the logs
 #[inline]
 pub fn print_logs(logs: &[Log]) {

crates/anvil/src/eth/backend/mem/mod.rs

@@ -76,6 +76,7 @@ use foundry_evm::{
         },
     },
     utils::new_evm_with_inspector_ref,
+    InspectorExt,
 };
 use futures::channel::mpsc::{unbounded, UnboundedSender};
 use parking_lot::{Mutex, RwLock};
@@ -819,7 +820,7 @@ impl Backend {
     ) -> revm::Evm<'_, I, WrapDatabaseRef<DB>>
     where
         DB: revm::DatabaseRef,
-        I: revm::Inspector<WrapDatabaseRef<DB>>,
+        I: InspectorExt<WrapDatabaseRef<DB>>,
     {
         let mut evm = new_evm_with_inspector_ref(db, env, inspector);
         if let Some(ref factory) = self.precompile_factory {

crates/cheatcodes/src/inspector.rs

@@ -21,8 +21,9 @@ use alloy_sol_types::{SolInterface, SolValue};
 use foundry_common::{evm::Breakpoints, provider::alloy::RpcUrl, SELECTOR_LEN};
 use foundry_evm_core::{
     abi::Vm::stopExpectSafeMemoryCall,
-    backend::{DatabaseError, DatabaseExt, RevertDiagnostic},
-    constants::{CHEATCODE_ADDRESS, DEFAULT_CREATE2_DEPLOYER, HARDHAT_CONSOLE_ADDRESS},
+    backend::{DatabaseExt, RevertDiagnostic},
+    constants::{CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS},
+    InspectorExt,
 };
 use itertools::Itertools;
 use revm::{
@@ -1327,22 +1328,19 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
                 ecx.env.tx.caller = broadcast.new_origin;
 
                 if ecx.journaled_state.depth() == broadcast.depth {
-                    let (bytecode, to, nonce) = process_broadcast_create(
-                        broadcast.new_origin,
-                        call.init_code.clone(),
-                        ecx,
-                        call,
-                    );
+                    call.caller = broadcast.new_origin;
                     let is_fixed_gas_limit = check_if_fixed_gas_limit(ecx, call.gas_limit);
 
+                    let account = &ecx.journaled_state.state()[&broadcast.new_origin];
+
                     self.broadcastable_transactions.push_back(BroadcastableTransaction {
                         rpc: ecx.db.active_fork_url(),
                         transaction: TransactionRequest {
                             from: Some(broadcast.new_origin),
-                            to,
+                            to: None,
                             value: Some(call.value),
-                            input: TransactionInput::new(bytecode),
-                            nonce: Some(nonce),
+                            input: TransactionInput::new(call.init_code.clone()),
+                            nonce: Some(account.info.nonce),
                             gas: if is_fixed_gas_limit {
                                 Some(call.gas_limit as u128)
                             } else {
@@ -1351,6 +1349,7 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
                             ..Default::default()
                         },
                     });
+
                     let kind = match call.scheme {
                         CreateScheme::Create => "create",
                         CreateScheme::Create2 { .. } => "create2",
@@ -1360,29 +1359,6 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
             }
         }
 
-        // Apply the Create2 deployer
-        if self.broadcast.is_some() || self.config.always_use_create_2_factory {
-            match apply_create2_deployer(
-                ecx,
-                call,
-                self.prank.as_ref(),
-                self.broadcast.as_ref(),
-                self.recorded_account_diffs_stack.as_mut(),
-            ) {
-                Ok(_) => {}
-                Err(err) => {
-                    return Some(CreateOutcome {
-                        result: InterpreterResult {
-                            result: InstructionResult::Revert,
-                            output: Error::encode(err),
-                            gas,
-                        },
-                        address: None,
-                    })
-                }
-            };
-        }
-
         // allow cheatcodes from the address of the new contract
         // Compute the address *after* any possible broadcast updates, so it's based on the updated
         // call inputs
@@ -1526,6 +1502,29 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
     }
 }
 
+impl<DB: DatabaseExt> InspectorExt<DB> for Cheatcodes {
+    fn should_use_create2_factory(
+        &mut self,
+        ecx: &mut EvmContext<DB>,
+        inputs: &mut CreateInputs,
+    ) -> bool {
+        if let CreateScheme::Create2 { .. } = inputs.scheme {
+            let target_depth = if let Some(prank) = &self.prank {
+                prank.depth
+            } else if let Some(broadcast) = &self.broadcast {
+                broadcast.depth
+            } else {
+                1
+            };
+
+            ecx.journaled_state.depth() == target_depth &&
+                (self.broadcast.is_some() || self.config.always_use_create_2_factory)
+        } else {
+            false
+        }
+    }
+}
+
 /// Helper that expands memory, stores a revert string pertaining to a disallowed memory write,
 /// and sets the return range to the revert string's location in memory.
 ///
@@ -1554,110 +1553,6 @@ fn disallowed_mem_write(
     };
 }
 
-/// Applies the default CREATE2 deployer for contract creation.
-///
-/// This function is invoked during the contract creation process and updates the caller of the
-/// contract creation transaction to be the `DEFAULT_CREATE2_DEPLOYER` if the `CreateScheme` is
-/// `Create2` and the current execution depth matches the depth at which the `prank` or `broadcast`
-/// was started, or the default depth of 1 if no prank or broadcast is currently active.
-///
-/// Returns a `DatabaseError::MissingCreate2Deployer` if the `DEFAULT_CREATE2_DEPLOYER` account is
-/// not found or if it does not have any associated bytecode.
-fn apply_create2_deployer<DB: DatabaseExt>(
-    ecx: &mut InnerEvmContext<DB>,
-    call: &mut CreateInputs,
-    prank: Option<&Prank>,
-    broadcast: Option<&Broadcast>,
-    diffs_stack: Option<&mut Vec<Vec<AccountAccess>>>,
-) -> Result<(), DB::Error> {
-    if let CreateScheme::Create2 { salt } = call.scheme {
-        let mut base_depth = 1;
-        if let Some(prank) = &prank {
-            base_depth = prank.depth;
-        } else if let Some(broadcast) = &broadcast {
-            base_depth = broadcast.depth;
-        }
-
-        // If the create scheme is Create2 and the depth equals the broadcast/prank/default
-        // depth, then use the default create2 factory as the deployer
-        if ecx.journaled_state.depth() == base_depth {
-            // Record the call to the create2 factory in the state diff
-            if let Some(recorded_account_diffs_stack) = diffs_stack {
-                let calldata = [&salt.to_be_bytes::<32>()[..], &call.init_code[..]].concat();
-                recorded_account_diffs_stack.push(vec![AccountAccess {
-                    chainInfo: crate::Vm::ChainInfo {
-                        forkId: ecx.db.active_fork_id().unwrap_or_default(),
-                        chainId: U256::from(ecx.env.cfg.chain_id),
-                    },
-                    accessor: call.caller,
-                    account: DEFAULT_CREATE2_DEPLOYER,
-                    kind: crate::Vm::AccountAccessKind::Call,
-                    initialized: true,
-                    oldBalance: U256::ZERO, // updated on create_end
-                    newBalance: U256::ZERO, // updated on create_end
-                    value: call.value,
-                    data: calldata.into(),
-                    reverted: false,
-                    deployedCode: Bytes::new(), // updated on create_end
-                    storageAccesses: vec![],    // updated on create_end
-                    depth: ecx.journaled_state.depth(),
-                }])
-            }
-
-            // Sanity checks for our CREATE2 deployer
-            // TODO: use ecx.load_account
-            let info =
-                &ecx.journaled_state.load_account(DEFAULT_CREATE2_DEPLOYER, &mut ecx.db)?.0.info;
-            match &info.code {
-                Some(code) if code.is_empty() => return Err(DatabaseError::MissingCreate2Deployer),
-                None if ecx.db.code_by_hash(info.code_hash)?.is_empty() => {
-                    return Err(DatabaseError::MissingCreate2Deployer)
-                }
-                _ => {}
-            }
-
-            call.caller = DEFAULT_CREATE2_DEPLOYER;
-        }
-    }
-    Ok(())
-}
-
-/// Processes the creation of a new contract when broadcasting, preparing the necessary data for the
-/// transaction to deploy the contract.
-///
-/// Returns the transaction calldata and the target address.
-///
-/// If the CreateScheme is Create, then this function returns the input bytecode without
-/// modification and no address since it will be filled in later. If the CreateScheme is Create2,
-/// then this function returns the calldata for the call to the create2 deployer which must be the
-/// salt and init code concatenated.
-fn process_broadcast_create<DB: DatabaseExt>(
-    broadcast_sender: Address,
-    bytecode: Bytes,
-    ecx: &mut InnerEvmContext<DB>,
-    call: &mut CreateInputs,
-) -> (Bytes, Option<Address>, u64) {
-    call.caller = broadcast_sender;
-    match call.scheme {
-        CreateScheme::Create => {
-            (bytecode, None, ecx.journaled_state.account(broadcast_sender).info.nonce)
-        }
-        CreateScheme::Create2 { salt } => {
-            // We have to increment the nonce of the user address, since this create2 will be done
-            // by the create2_deployer
-            let account = ecx.journaled_state.state().get_mut(&broadcast_sender).unwrap();
-            let prev = account.info.nonce;
-            // Touch account to ensure that incremented nonce is committed
-            account.mark_touch();
-            account.info.nonce += 1;
-            debug!(target: "cheatcodes", address=%broadcast_sender, nonce=prev+1, prev, "incremented nonce in create2");
-            // Proxy deployer requires the data to be `salt ++ init_code`
-            let calldata = [&salt.to_be_bytes::<32>()[..], &bytecode[..]].concat();
-            (calldata.into(), Some(DEFAULT_CREATE2_DEPLOYER), prev)
-        }
-    }
-}
-
 // Determines if the gas limit on a given call was manually set in the script and should therefore
 // not be overwritten by later estimations
 fn check_if_fixed_gas_limit<DB: DatabaseExt>(

crates/evm/core/src/backend/cow.rs

@@ -6,6 +6,7 @@ use crate::{
         RevertSnapshotAction,
     },
     fork::{CreateFork, ForkId},
+    InspectorExt,
 };
 use alloy_genesis::GenesisAccount;
 use alloy_primitives::{Address, B256, U256};
@@ -16,7 +17,7 @@ use revm::{
         Account, AccountInfo, Bytecode, Env, EnvWithHandlerCfg, HashMap as Map, ResultAndState,
         SpecId,
     },
-    Database, DatabaseCommit, Inspector, JournaledState,
+    Database, DatabaseCommit, JournaledState,
 };
 use std::{borrow::Cow, collections::BTreeMap};
 
@@ -58,7 +59,7 @@ impl<'a> CowBackend<'a> {
     ///
     /// Note: in case there are any cheatcodes executed that modify the environment, this will
     /// update the given `env` with the new values.
-    pub fn inspect<'b, I: Inspector<&'b mut Self>>(
+    pub fn inspect<'b, I: InspectorExt<&'b mut Self>>(
         &'b mut self,
         env: &mut EnvWithHandlerCfg,
         inspector: I,
@@ -176,7 +177,7 @@ impl<'a> DatabaseExt for CowBackend<'a> {
         self.backend_mut(env).roll_fork_to_transaction(id, transaction, env, journaled_state)
     }
 
-    fn transact<I: Inspector<Backend>>(
+    fn transact<I: InspectorExt<Backend>>(
         &mut self,
         id: Option<LocalForkId>,
         transaction: B256,

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

@@ -5,6 +5,7 @@ use crate::{
     fork::{CreateFork, ForkId, MultiFork, SharedBackend},
     snapshot::Snapshots,
     utils::configure_tx_env,
+    InspectorExt,
 };
 use alloy_genesis::GenesisAccount;
 use alloy_primitives::{b256, keccak256, Address, B256, U256};
@@ -19,7 +20,7 @@ use revm::{
         Account, AccountInfo, Bytecode, CreateScheme, Env, EnvWithHandlerCfg, HashMap as Map, Log,
         ResultAndState, SpecId, State, StorageSlot, TransactTo, KECCAK_EMPTY,
     },
-    Database, DatabaseCommit, Inspector, JournaledState,
+    Database, DatabaseCommit, JournaledState,
 };
 use std::{
     collections::{BTreeMap, HashMap, HashSet},
@@ -188,7 +189,7 @@ pub trait DatabaseExt: Database<Error = DatabaseError> {
     ) -> eyre::Result<()>;
 
     /// Fetches the given transaction for the fork and executes it, committing the state in the DB
-    fn transact<I: Inspector<Backend>>(
+    fn transact<I: InspectorExt<Backend>>(
         &mut self,
         id: Option<LocalForkId>,
         transaction: B256,
@@ -780,7 +781,7 @@ impl Backend {
     ///
     /// Note: in case there are any cheatcodes executed that modify the environment, this will
     /// update the given `env` with the new values.
-    pub fn inspect<'a, I: Inspector<&'a mut Self>>(
+    pub fn inspect<'a, I: InspectorExt<&'a mut Self>>(
         &'a mut self,
         env: &mut EnvWithHandlerCfg,
         inspector: I,
@@ -1228,7 +1229,7 @@ impl DatabaseExt for Backend {
         Ok(())
     }
 
-    fn transact<I: Inspector<Backend>>(
+    fn transact<I: InspectorExt<Backend>>(
         &mut self,
         maybe_id: Option<LocalForkId>,
         transaction: B256,
@@ -1868,7 +1869,7 @@ fn update_env_block(env: &mut Env, fork_block: u64, block: &Block) {
 
 /// Executes the given transaction and commits state changes to the database _and_ the journaled
 /// state, with an optional inspector
-fn commit_transaction<I: Inspector<Backend>>(
+fn commit_transaction<I: InspectorExt<Backend>>(
     tx: WithOtherFields<Transaction>,
     mut env: EnvWithHandlerCfg,
     journaled_state: &mut JournaledState,

crates/evm/core/src/lib.rs

@@ -4,6 +4,10 @@
 
 #![warn(unused_crate_dependencies)]
 
+use auto_impl::auto_impl;
+use revm::{inspectors::NoOpInspector, interpreter::CreateInputs, Database, EvmContext, Inspector};
+use revm_inspectors::access_list::AccessListInspector;
+
 #[macro_use]
 extern crate tracing;
 
@@ -19,3 +23,23 @@ pub mod opcodes;
 pub mod opts;
 pub mod snapshot;
 pub mod utils;
+
+/// An extension trait that allows us to add additional hooks to Inspector for later use in
+/// handlers.
+#[auto_impl(&mut, Box)]
+pub trait InspectorExt<DB: Database>: Inspector<DB> {
+    /// Determines whether the `DEFAULT_CREATE2_DEPLOYER` should be used for a CREATE2 frame.
+    ///
+    /// If this function returns true, we'll replace CREATE2 frame with a CALL frame to CREATE2
+    /// factory.
+    fn should_use_create2_factory(
+        &mut self,
+        _context: &mut EvmContext<DB>,
+        _inputs: &mut CreateInputs,
+    ) -> bool {
+        false
+    }
+}
+
+impl<DB: Database> InspectorExt<DB> for NoOpInspector {}
+impl<DB: Database> InspectorExt<DB> for AccessListInspector {}