refactor(cheatcodes): extract fork env helper to reduce duplication (foundry-rs#13848)

Extract `fork_env_op` helper that handles the common pattern shared by
all fork-switching cheatcodes: clone EVM/tx env → run db operation →
write env back. Deduplicates 7 call sites (rollFork × 4, selectFork,
createSelectFork × 2).

Author: figtracer

crates/cheatcodes/src/evm/fork.rs

@@ -10,8 +10,11 @@ use alloy_provider::Provider;
 use alloy_rpc_types::Filter;
 use alloy_sol_types::SolValue;
 use foundry_common::provider::ProviderBuilder;
-use foundry_evm_core::{backend::FoundryJournalExt, fork::CreateFork};
-use revm::context::ContextTr;
+use foundry_evm_core::{
+    backend::{FoundryJournalExt, JournaledState},
+    fork::CreateFork,
+};
+use revm::context::{ContextTr, TxEnv};
 
 impl Cheatcode for activeForkCall {
     fn apply_stateful<CTX: ContextTr<Db: DatabaseExt>>(
@@ -73,55 +76,39 @@ impl Cheatcode for rollFork_0Call {
     fn apply_stateful<CTX: EthCheatCtx>(&self, ccx: &mut CheatsCtxt<'_, CTX>) -> Result {
         let Self { blockNumber } = self;
         persist_caller(ccx);
-        let mut evm_env = ccx.ecx.evm_clone();
-        let mut tx_env = ccx.ecx.tx_clone();
-        let (db, inner) = ccx.ecx.journal_mut().as_db_and_inner();
-        db.roll_fork(None, (*blockNumber).to(), &mut evm_env, &mut tx_env, inner)?;
-        ccx.ecx.set_evm(evm_env);
-        ccx.ecx.set_tx(tx_env);
-        Ok(Default::default())
+        fork_env_op(ccx, |db, evm_env, tx_env, inner| {
+            db.roll_fork(None, (*blockNumber).to(), evm_env, tx_env, inner)
+        })
     }
 }
 
 impl Cheatcode for rollFork_1Call {
     fn apply_stateful<CTX: EthCheatCtx>(&self, ccx: &mut CheatsCtxt<'_, CTX>) -> Result {
         let Self { txHash } = self;
         persist_caller(ccx);
-        let mut evm_env = ccx.ecx.evm_clone();
-        let mut tx_env = ccx.ecx.tx_clone();
-        let (db, inner) = ccx.ecx.journal_mut().as_db_and_inner();
-        db.roll_fork_to_transaction(None, *txHash, &mut evm_env, &mut tx_env, inner)?;
-        ccx.ecx.set_evm(evm_env);
-        ccx.ecx.set_tx(tx_env);
-        Ok(Default::default())
+        fork_env_op(ccx, |db, evm_env, tx_env, inner| {
+            db.roll_fork_to_transaction(None, *txHash, evm_env, tx_env, inner)
+        })
     }
 }
 
 impl Cheatcode for rollFork_2Call {
     fn apply_stateful<CTX: EthCheatCtx>(&self, ccx: &mut CheatsCtxt<'_, CTX>) -> Result {
         let Self { forkId, blockNumber } = self;
         persist_caller(ccx);
-        let mut evm_env = ccx.ecx.evm_clone();
-        let mut tx_env = ccx.ecx.tx_clone();
-        let (db, inner) = ccx.ecx.journal_mut().as_db_and_inner();
-        db.roll_fork(Some(*forkId), (*blockNumber).to(), &mut evm_env, &mut tx_env, inner)?;
-        ccx.ecx.set_evm(evm_env);
-        ccx.ecx.set_tx(tx_env);
-        Ok(Default::default())
+        fork_env_op(ccx, |db, evm_env, tx_env, inner| {
+            db.roll_fork(Some(*forkId), (*blockNumber).to(), evm_env, tx_env, inner)
+        })
     }
 }
 
 impl Cheatcode for rollFork_3Call {
     fn apply_stateful<CTX: EthCheatCtx>(&self, ccx: &mut CheatsCtxt<'_, CTX>) -> Result {
         let Self { forkId, txHash } = self;
         persist_caller(ccx);
-        let mut evm_env = ccx.ecx.evm_clone();
-        let mut tx_env = ccx.ecx.tx_clone();
-        let (db, inner) = ccx.ecx.journal_mut().as_db_and_inner();
-        db.roll_fork_to_transaction(Some(*forkId), *txHash, &mut evm_env, &mut tx_env, inner)?;
-        ccx.ecx.set_evm(evm_env);
-        ccx.ecx.set_tx(tx_env);
-        Ok(Default::default())
+        fork_env_op(ccx, |db, evm_env, tx_env, inner| {
+            db.roll_fork_to_transaction(Some(*forkId), *txHash, evm_env, tx_env, inner)
+        })
     }
 }
 
@@ -130,13 +117,9 @@ impl Cheatcode for selectForkCall {
         let Self { forkId } = self;
         persist_caller(ccx);
         check_broadcast(ccx.state)?;
-        let mut evm_env = ccx.ecx.evm_clone();
-        let mut tx_env = ccx.ecx.tx_clone();
-        let (db, inner) = ccx.ecx.journal_mut().as_db_and_inner();
-        db.select_fork(*forkId, &mut evm_env, &mut tx_env, inner)?;
-        ccx.ecx.set_evm(evm_env);
-        ccx.ecx.set_tx(tx_env);
-        Ok(Default::default())
+        fork_env_op(ccx, |db, evm_env, tx_env, inner| {
+            db.select_fork(*forkId, evm_env, tx_env, inner)
+        })
     }
 }
 
@@ -355,13 +338,9 @@ fn create_select_fork<CTX: EthCheatCtx>(
     check_broadcast(ccx.state)?;
 
     let fork = create_fork_request(ccx, url_or_alias, block)?;
-    let mut evm_env = ccx.ecx.evm_clone();
-    let mut tx_env = ccx.ecx.tx_clone();
-    let (db, inner) = ccx.ecx.journal_mut().as_db_and_inner();
-    let id = db.create_select_fork(fork, &mut evm_env, &mut tx_env, inner)?;
-    ccx.ecx.set_evm(evm_env);
-    ccx.ecx.set_tx(tx_env);
-    Ok(id.abi_encode())
+    fork_env_op(ccx, |db, evm_env, tx_env, inner| {
+        db.create_select_fork(fork, evm_env, tx_env, inner)
+    })
 }
 
 /// Creates a new fork
@@ -384,14 +363,9 @@ fn create_select_fork_at_transaction<CTX: EthCheatCtx>(
     check_broadcast(ccx.state)?;
 
     let fork = create_fork_request(ccx, url_or_alias, None)?;
-    let mut evm_env = ccx.ecx.evm_clone();
-    let mut tx_env = ccx.ecx.tx_clone();
-    let (db, inner) = ccx.ecx.journal_mut().as_db_and_inner();
-    let id =
-        db.create_select_fork_at_transaction(fork, &mut evm_env, &mut tx_env, inner, *transaction)?;
-    ccx.ecx.set_evm(evm_env);
-    ccx.ecx.set_tx(tx_env);
-    Ok(id.abi_encode())
+    fork_env_op(ccx, |db, evm_env, tx_env, inner| {
+        db.create_select_fork_at_transaction(fork, evm_env, tx_env, inner, *transaction)
+    })
 }
 
 /// Creates a new fork at the given transaction
@@ -435,6 +409,27 @@ fn create_fork_request<CTX: EthCheatCtx>(
     Ok(fork)
 }
 
+/// Clones the EVM and tx environments, runs a fork operation that may modify them, then writes
+/// them back. This is the common pattern for all fork-switching cheatcodes (rollFork, selectFork,
+/// createSelectFork).
+fn fork_env_op<CTX: EthCheatCtx, T: SolValue>(
+    ccx: &mut CheatsCtxt<'_, CTX>,
+    f: impl FnOnce(
+        &mut dyn DatabaseExt,
+        &mut EvmEnv,
+        &mut TxEnv,
+        &mut JournaledState,
+    ) -> eyre::Result<T>,
+) -> Result {
+    let mut evm_env = ccx.ecx.evm_clone();
+    let mut tx_env = ccx.ecx.tx_clone();
+    let (db, inner) = ccx.ecx.journal_mut().as_db_and_inner();
+    let result = f(db, &mut evm_env, &mut tx_env, inner)?;
+    ccx.ecx.set_evm(evm_env);
+    ccx.ecx.set_tx(tx_env);
+    Ok(result.abi_encode())
+}
+
 fn check_broadcast(state: &Cheatcodes) -> Result<()> {
     if state.broadcast.is_none() {
         Ok(())