refactor(cheatcodes): mv ScriptWallets into Cheatcode (foundry-rs#9106)

* refactor(`cheatcodes`): mv `ScriptWallets` into `Cheatcode` from `CheatsConfig`

* nit

* rename `ScriptWallets` to `Wallets`

* rename cheatcode

* doc nits

Author: yash-atreya

crates/cheatcodes/assets/cheatcodes.json

@@ -1912,7 +1912,7 @@ interface Vm {
 
     /// Returns addresses of available unlocked wallets in the script environment.
     #[cheatcode(group = Scripting)]
-    function getScriptWallets() external returns (address[] memory wallets);
+    function getWallets() external returns (address[] memory wallets);
 
     // ======== Utilities ========
 

crates/cheatcodes/spec/src/vm.rs

@@ -1,5 +1,5 @@
 use super::Result;
-use crate::{script::ScriptWallets, Vm::Rpc};
+use crate::Vm::Rpc;
 use alloy_primitives::{map::AddressHashMap, U256};
 use foundry_common::{fs::normalize_path, ContractsByArtifact};
 use foundry_compilers::{utils::canonicalize, ProjectPathsConfig};
@@ -43,8 +43,6 @@ pub struct CheatsConfig {
     pub evm_opts: EvmOpts,
     /// Address labels from config
     pub labels: AddressHashMap<String>,
-    /// Script wallets
-    pub script_wallets: Option<ScriptWallets>,
     /// Artifacts which are guaranteed to be fresh (either recompiled or cached).
     /// If Some, `vm.getDeployedCode` invocations are validated to be in scope of this list.
     /// If None, no validation is performed.
@@ -65,7 +63,6 @@ impl CheatsConfig {
         config: &Config,
         evm_opts: EvmOpts,
         available_artifacts: Option<ContractsByArtifact>,
-        script_wallets: Option<ScriptWallets>,
         running_contract: Option<String>,
         running_version: Option<Version>,
     ) -> Self {
@@ -93,7 +90,6 @@ impl CheatsConfig {
             allowed_paths,
             evm_opts,
             labels: config.labels.clone(),
-            script_wallets,
             available_artifacts,
             running_contract,
             running_version,
@@ -223,7 +219,6 @@ impl Default for CheatsConfig {
             allowed_paths: vec![],
             evm_opts: Default::default(),
             labels: Default::default(),
-            script_wallets: None,
             available_artifacts: Default::default(),
             running_contract: Default::default(),
             running_version: Default::default(),
@@ -245,7 +240,6 @@ mod tests {
             None,
             None,
             None,
-            None,
         )
     }
 

crates/cheatcodes/src/config.rs

@@ -1,6 +1,6 @@
 //! Implementations of [`Crypto`](spec::Group::Crypto) Cheatcodes.
 
-use crate::{Cheatcode, Cheatcodes, Result, ScriptWallets, Vm::*};
+use crate::{Cheatcode, Cheatcodes, Result, Vm::*, Wallets};
 use alloy_primitives::{keccak256, Address, B256, U256};
 use alloy_signer::{Signer, SignerSync};
 use alloy_signer_local::{
@@ -17,7 +17,6 @@ use k256::{
     elliptic_curve::{bigint::ArrayEncoding, sec1::ToEncodedPoint},
 };
 use p256::ecdsa::{signature::hazmat::PrehashSigner, Signature, SigningKey as P256SigningKey};
-use std::sync::Arc;
 
 /// The BIP32 default derivation path prefix.
 const DEFAULT_DERIVATION_PATH_PREFIX: &str = "m/44'/60'/0'/0/";
@@ -134,9 +133,9 @@ fn inject_wallet(state: &mut Cheatcodes, wallet: LocalSigner<SigningKey>) -> Add
         script_wallets.add_local_signer(wallet);
     } else {
         // This is needed in case of testing scripts, wherein script wallets are not set on setup.
-        let script_wallets = ScriptWallets::new(MultiWallet::default(), None);
+        let script_wallets = Wallets::new(MultiWallet::default(), None);
         script_wallets.add_local_signer(wallet);
-        Arc::make_mut(&mut state.config).script_wallets = Some(script_wallets);
+        state.set_wallets(script_wallets);
     }
     address
 }

crates/cheatcodes/src/crypto.rs

@@ -8,7 +8,7 @@ use crate::{
         DealRecord, GasRecord, RecordAccess,
     },
     inspector::utils::CommonCreateInput,
-    script::{Broadcast, ScriptWallets},
+    script::{Broadcast, Wallets},
     test::{
         assume::AssumeNoRevert,
         expect::{
@@ -476,6 +476,8 @@ pub struct Cheatcodes {
 
     /// Deprecated cheatcodes mapped to the reason. Used to report warnings on test results.
     pub deprecated: HashMap<&'static str, Option<&'static str>>,
+    /// Unlocked wallets used in scripts and testing of scripts.
+    pub wallets: Option<Wallets>,
 }
 
 // This is not derived because calling this in `fn new` with `..Default::default()` creates a second
@@ -523,12 +525,18 @@ impl Cheatcodes {
             ignored_traces: Default::default(),
             arbitrary_storage: Default::default(),
             deprecated: Default::default(),
+            wallets: Default::default(),
         }
     }
 
     /// Returns the configured script wallets.
-    pub fn script_wallets(&self) -> Option<&ScriptWallets> {
-        self.config.script_wallets.as_ref()
+    pub fn script_wallets(&self) -> Option<&Wallets> {
+        self.wallets.as_ref()
+    }
+
+    /// Sets the unlocked wallets.
+    pub fn set_wallets(&mut self, wallets: Wallets) {
+        self.wallets = Some(wallets);
     }
 
     /// Decodes the input data and applies the cheatcode.

crates/cheatcodes/src/inspector.rs

@@ -45,7 +45,7 @@ mod inspector;
 mod json;
 
 mod script;
-pub use script::{ScriptWallets, ScriptWalletsInner};
+pub use script::{Wallets, WalletsInner};
 
 mod string;
 

crates/cheatcodes/src/lib.rs

@@ -61,7 +61,7 @@ impl Cheatcode for stopBroadcastCall {
     }
 }
 
-impl Cheatcode for getScriptWalletsCall {
+impl Cheatcode for getWalletsCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let script_wallets =
             ccx.state.script_wallets().cloned().map(|sw| sw.signers().unwrap_or_default());
@@ -91,29 +91,29 @@ pub struct Broadcast {
 
 /// Contains context for wallet management.
 #[derive(Debug)]
-pub struct ScriptWalletsInner {
+pub struct WalletsInner {
     /// All signers in scope of the script.
     pub multi_wallet: MultiWallet,
     /// Optional signer provided as `--sender` flag.
     pub provided_sender: Option<Address>,
 }
 
-/// Clonable wrapper around [`ScriptWalletsInner`].
+/// Clonable wrapper around [`WalletsInner`].
 #[derive(Debug, Clone)]
-pub struct ScriptWallets {
+pub struct Wallets {
     /// Inner data.
-    pub inner: Arc<Mutex<ScriptWalletsInner>>,
+    pub inner: Arc<Mutex<WalletsInner>>,
 }
 
-impl ScriptWallets {
+impl Wallets {
     #[allow(missing_docs)]
     pub fn new(multi_wallet: MultiWallet, provided_sender: Option<Address>) -> Self {
-        Self { inner: Arc::new(Mutex::new(ScriptWalletsInner { multi_wallet, provided_sender })) }
+        Self { inner: Arc::new(Mutex::new(WalletsInner { multi_wallet, provided_sender })) }
     }
 
-    /// Consumes [ScriptWallets] and returns [MultiWallet].
+    /// Consumes [Wallets] and returns [MultiWallet].
     ///
-    /// Panics if [ScriptWallets] is still in use.
+    /// Panics if [Wallets] is still in use.
     pub fn into_multi_wallet(self) -> MultiWallet {
         Arc::into_inner(self.inner)
             .map(|m| m.into_inner().multi_wallet)

crates/cheatcodes/src/script.rs

@@ -335,7 +335,6 @@ impl SessionSource {
                         self.config.evm_opts.clone(),
                         None,
                         None,
-                        None,
                         Some(self.solc.version.clone()),
                     )
                     .into(),

crates/chisel/src/executor.rs

@@ -3,7 +3,7 @@ use super::{
     TracingInspector,
 };
 use alloy_primitives::{map::AddressHashMap, Address, Bytes, Log, TxKind, U256};
-use foundry_cheatcodes::CheatcodesExecutor;
+use foundry_cheatcodes::{CheatcodesExecutor, Wallets};
 use foundry_evm_core::{backend::DatabaseExt, InspectorExt};
 use foundry_evm_coverage::HitMaps;
 use foundry_evm_traces::{SparsedTraceArena, TraceMode};
@@ -57,6 +57,8 @@ pub struct InspectorStackBuilder {
     pub enable_isolation: bool,
     /// Whether to enable Alphanet features.
     pub alphanet: bool,
+    /// The wallets to set in the cheatcodes context.
+    pub wallets: Option<Wallets>,
 }
 
 impl InspectorStackBuilder {
@@ -87,6 +89,13 @@ impl InspectorStackBuilder {
         self
     }
 
+    /// Set the wallets.
+    #[inline]
+    pub fn wallets(mut self, wallets: Wallets) -> Self {
+        self.wallets = Some(wallets);
+        self
+    }
+
     /// Set the fuzzer inspector.
     #[inline]
     pub fn fuzzer(mut self, fuzzer: Fuzzer) -> Self {
@@ -161,13 +170,20 @@ impl InspectorStackBuilder {
             chisel_state,
             enable_isolation,
             alphanet,
+            wallets,
         } = self;
         let mut stack = InspectorStack::new();
 
         // inspectors
         if let Some(config) = cheatcodes {
-            stack.set_cheatcodes(Cheatcodes::new(config));
+            let mut cheatcodes = Cheatcodes::new(config);
+            // Set wallets if they are provided
+            if let Some(wallets) = wallets {
+                cheatcodes.set_wallets(wallets);
+            }
+            stack.set_cheatcodes(cheatcodes);
         }
+
         if let Some(fuzzer) = fuzzer {
             stack.set_fuzzer(fuzzer);
         }

crates/evm/evm/src/inspectors/stack.rs

@@ -242,7 +242,6 @@ impl MultiContractRunner {
             &self.config,
             self.evm_opts.clone(),
             Some(self.known_contracts.clone()),
-            None,
             Some(artifact_id.name.clone()),
             Some(artifact_id.version.clone()),
         );