From 11bd7715d5c16252f58c47b6e418930cf1501576 Mon Sep 17 00:00:00 2001 From: Yash Atreya <44857776+yash-atreya@users.noreply.github.com> Date: Tue, 15 Oct 2024 14:08:09 +0530 Subject: [PATCH] chore(`cheatcodes`): wallet nits (#9118) --- crates/cheatcodes/src/crypto.rs | 28 +++++----------- crates/cheatcodes/src/inspector.rs | 7 ++-- crates/cheatcodes/src/script.rs | 51 ++++++++++++++++-------------- 3 files changed, 40 insertions(+), 46 deletions(-) diff --git a/crates/cheatcodes/src/crypto.rs b/crates/cheatcodes/src/crypto.rs index f8f24e70b..cdb07720c 100644 --- a/crates/cheatcodes/src/crypto.rs +++ b/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::(mnemonic, derivationPath, *count)?; - - tracing::info!("Adding {} keys to script wallets", count); - let mut addresses = Vec::
::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) -> 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
, digest: &B256, ) -> Result { - 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 diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index a5a2a474f..043f8af2f 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/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. diff --git a/crates/cheatcodes/src/script.rs b/crates/cheatcodes/src/script.rs index 960d6a550..29a804efd 100644 --- a/crates/cheatcodes/src/script.rs +++ b/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
= 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> { 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,7 +181,7 @@ 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)?; @@ -183,9 +189,8 @@ fn broadcast_key(ccx: &mut CheatsCtxt, private_key: &U256, single_call: bool) -> 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 }