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

crates/cheatcodes/spec/src/vm.rs

@@ -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/src/config.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/crypto.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/inspector.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/lib.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/script.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/chisel/src/executor.rs

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

crates/evm/evm/src/inspectors/stack.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/forge/src/multi_runner.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()),
         );

crates/forge/tests/cli/script.rs

@@ -2077,12 +2077,12 @@ forgetest_init!(can_get_script_wallets, |prj, cmd| {
 import "forge-std/Script.sol";
 
 interface Vm {
-    function getScriptWallets() external returns (address[] memory wallets);
+    function getWallets() external returns (address[] memory wallets);
 }
 
 contract WalletScript is Script {
     function run() public {
-        address[] memory wallets = Vm(address(vm)).getScriptWallets();
+        address[] memory wallets = Vm(address(vm)).getWallets();
         console.log(wallets[0]);
     }
 }"#,

crates/script/src/broadcast.rs

@@ -17,7 +17,7 @@ use alloy_serde::WithOtherFields;
 use alloy_transport::Transport;
 use eyre::{bail, Context, Result};
 use forge_verify::provider::VerificationProviderType;
-use foundry_cheatcodes::ScriptWallets;
+use foundry_cheatcodes::Wallets;
 use foundry_cli::utils::{has_batch_support, has_different_gas_calc};
 use foundry_common::{
     provider::{get_http_provider, try_get_http_provider, RetryProvider},
@@ -154,7 +154,7 @@ impl SendTransactionsKind {
 pub struct BundledState {
     pub args: ScriptArgs,
     pub script_config: ScriptConfig,
-    pub script_wallets: ScriptWallets,
+    pub script_wallets: Wallets,
     pub build_data: LinkedBuildData,
     pub sequence: ScriptSequenceKind,
 }

crates/script/src/build.rs

@@ -8,7 +8,7 @@ use crate::{
 use alloy_primitives::{Bytes, B256};
 use alloy_provider::Provider;
 use eyre::{OptionExt, Result};
-use foundry_cheatcodes::ScriptWallets;
+use foundry_cheatcodes::Wallets;
 use foundry_common::{
     compile::ProjectCompiler, provider::try_get_http_provider, ContractData, ContractsByArtifact,
 };
@@ -156,7 +156,7 @@ impl LinkedBuildData {
 pub struct PreprocessedState {
     pub args: ScriptArgs,
     pub script_config: ScriptConfig,
-    pub script_wallets: ScriptWallets,
+    pub script_wallets: Wallets,
 }
 
 impl PreprocessedState {
@@ -242,7 +242,7 @@ impl PreprocessedState {
 pub struct CompiledState {
     pub args: ScriptArgs,
     pub script_config: ScriptConfig,
-    pub script_wallets: ScriptWallets,
+    pub script_wallets: Wallets,
     pub build_data: BuildData,
 }
 

crates/script/src/execute.rs

@@ -14,7 +14,7 @@ use alloy_provider::Provider;
 use alloy_rpc_types::TransactionInput;
 use async_recursion::async_recursion;
 use eyre::{OptionExt, Result};
-use foundry_cheatcodes::ScriptWallets;
+use foundry_cheatcodes::Wallets;
 use foundry_cli::utils::{ensure_clean_constructor, needs_setup};
 use foundry_common::{
     fmt::{format_token, format_token_raw},
@@ -41,7 +41,7 @@ use yansi::Paint;
 pub struct LinkedState {
     pub args: ScriptArgs,
     pub script_config: ScriptConfig,
-    pub script_wallets: ScriptWallets,
+    pub script_wallets: Wallets,
     pub build_data: LinkedBuildData,
 }
 
@@ -92,7 +92,7 @@ impl LinkedState {
 pub struct PreExecutionState {
     pub args: ScriptArgs,
     pub script_config: ScriptConfig,
-    pub script_wallets: ScriptWallets,
+    pub script_wallets: Wallets,
     pub build_data: LinkedBuildData,
     pub execution_data: ExecutionData,
 }
@@ -274,7 +274,7 @@ pub struct ExecutionArtifacts {
 pub struct ExecutedState {
     pub args: ScriptArgs,
     pub script_config: ScriptConfig,
-    pub script_wallets: ScriptWallets,
+    pub script_wallets: Wallets,
     pub build_data: LinkedBuildData,
     pub execution_data: ExecutionData,
     pub execution_result: ScriptResult,