chore(cheatcodes): wallet nits (foundry-rs#9118)

Author: yash-atreya

crates/cheatcodes/src/crypto.rs

@@ -1,6 +1,6 @@
 //! Implementations of [`Crypto`](spec::Group::Crypto) Cheatcodes.
 
-use crate::{Cheatcode, Cheatcodes, Result, Vm::*, Wallets};
+use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
 use alloy_primitives::{keccak256, Address, B256, U256};
 use alloy_signer::{Signer, SignerSync};
 use alloy_signer_local::{
@@ -11,7 +11,6 @@ use alloy_signer_local::{
     LocalSigner, MnemonicBuilder, PrivateKeySigner,
 };
 use alloy_sol_types::SolValue;
-use foundry_wallets::multi_wallet::MultiWallet;
 use k256::{
     ecdsa::SigningKey,
     elliptic_curve::{bigint::ArrayEncoding, sec1::ToEncodedPoint},
@@ -98,11 +97,7 @@ impl Cheatcode for rememberKeyCall {
 impl Cheatcode for rememberKeys_0Call {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self { mnemonic, derivationPath, count } = self;
-        tracing::info!("Remembering {} keys", count);
         let wallets = derive_wallets::<English>(mnemonic, derivationPath, *count)?;
-
-        tracing::info!("Adding {} keys to script wallets", count);
-
         let mut addresses = Vec::<Address>::with_capacity(wallets.len());
         for wallet in wallets {
             let addr = inject_wallet(state, wallet);
@@ -129,14 +124,7 @@ impl Cheatcode for rememberKeys_1Call {
 
 fn inject_wallet(state: &mut Cheatcodes, wallet: LocalSigner<SigningKey>) -> Address {
     let address = wallet.address();
-    if let Some(script_wallets) = state.script_wallets() {
-        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 = Wallets::new(MultiWallet::default(), None);
-        script_wallets.add_local_signer(wallet);
-        state.set_wallets(script_wallets);
-    }
+    state.wallets().add_local_signer(wallet);
     address
 }
 
@@ -256,13 +244,13 @@ fn sign_with_wallet(
     signer: Option<Address>,
     digest: &B256,
 ) -> Result<alloy_primitives::Signature> {
-    let Some(script_wallets) = state.script_wallets() else {
-        bail!("no wallets are available");
-    };
+    if state.wallets().is_empty() {
+        bail!("no wallets available");
+    }
 
-    let mut script_wallets = script_wallets.inner.lock();
-    let maybe_provided_sender = script_wallets.provided_sender;
-    let signers = script_wallets.multi_wallet.signers()?;
+    let mut wallets = state.wallets().inner.lock();
+    let maybe_provided_sender = wallets.provided_sender;
+    let signers = wallets.multi_wallet.signers()?;
 
     let signer = if let Some(signer) = signer {
         signer

crates/cheatcodes/src/inspector.rs

@@ -37,6 +37,7 @@ use foundry_evm_core::{
     InspectorExt,
 };
 use foundry_evm_traces::{TracingInspector, TracingInspectorConfig};
+use foundry_wallets::multi_wallet::MultiWallet;
 use itertools::Itertools;
 use proptest::test_runner::{RngAlgorithm, TestRng, TestRunner};
 use rand::Rng;
@@ -529,9 +530,9 @@ impl Cheatcodes {
         }
     }
 
-    /// Returns the configured script wallets.
-    pub fn script_wallets(&self) -> Option<&Wallets> {
-        self.wallets.as_ref()
+    /// Returns the configured wallets if available, else creates a new instance.
+    pub fn wallets(&mut self) -> &Wallets {
+        self.wallets.get_or_insert(Wallets::new(MultiWallet::default(), None))
     }
 
     /// Sets the unlocked wallets.

crates/cheatcodes/src/script.rs

@@ -63,15 +63,8 @@ impl Cheatcode for stopBroadcastCall {
 
 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());
-
-        if let Some(script_wallets) = script_wallets {
-            let script_wallets: Vec<Address> = script_wallets.into_iter().collect();
-            Ok(script_wallets.abi_encode())
-        } else {
-            Ok(Default::default())
-        }
+        let wallets = ccx.state.wallets().signers().unwrap_or_default();
+        Ok(wallets.abi_encode())
     }
 }
 
@@ -135,6 +128,21 @@ impl Wallets {
     pub fn signers(&self) -> Result<Vec<Address>> {
         Ok(self.inner.lock().multi_wallet.signers()?.keys().cloned().collect())
     }
+
+    /// Number of signers in the [MultiWallet].
+    pub fn len(&self) -> usize {
+        let mut inner = self.inner.lock();
+        let signers = inner.multi_wallet.signers();
+        if signers.is_err() {
+            return 0;
+        }
+        signers.unwrap().len()
+    }
+
+    /// Whether the [MultiWallet] is empty.
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
 }
 
 /// Sets up broadcasting from a script using `new_origin` as the sender.
@@ -148,16 +156,14 @@ fn broadcast(ccx: &mut CheatsCtxt, new_origin: Option<&Address>, single_call: bo
     let mut new_origin = new_origin.cloned();
 
     if new_origin.is_none() {
-        if let Some(script_wallets) = ccx.state.script_wallets() {
-            let mut script_wallets = script_wallets.inner.lock();
-            if let Some(provided_sender) = script_wallets.provided_sender {
-                new_origin = Some(provided_sender);
-            } else {
-                let signers = script_wallets.multi_wallet.signers()?;
-                if signers.len() == 1 {
-                    let address = signers.keys().next().unwrap();
-                    new_origin = Some(*address);
-                }
+        let mut wallets = ccx.state.wallets().inner.lock();
+        if let Some(provided_sender) = wallets.provided_sender {
+            new_origin = Some(provided_sender);
+        } else {
+            let signers = wallets.multi_wallet.signers()?;
+            if signers.len() == 1 {
+                let address = signers.keys().next().unwrap();
+                new_origin = Some(*address);
             }
         }
     }
@@ -175,17 +181,16 @@ fn broadcast(ccx: &mut CheatsCtxt, new_origin: Option<&Address>, single_call: bo
 }
 
 /// Sets up broadcasting from a script with the sender derived from `private_key`.
-/// Adds this private key to `state`'s `script_wallets` vector to later be used for signing
+/// Adds this private key to `state`'s `wallets` vector to later be used for signing
 /// if broadcast is successful.
 fn broadcast_key(ccx: &mut CheatsCtxt, private_key: &U256, single_call: bool) -> Result {
     let wallet = super::crypto::parse_wallet(private_key)?;
     let new_origin = wallet.address();
 
     let result = broadcast(ccx, Some(&new_origin), single_call);
     if result.is_ok() {
-        if let Some(script_wallets) = ccx.state.script_wallets() {
-            script_wallets.add_local_signer(wallet);
-        }
+        let wallets = ccx.state.wallets();
+        wallets.add_local_signer(wallet);
     }
     result
 }