Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(forge): do not re-execute script on resume when possible #7361

Merged
merged 3 commits into from
Mar 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions crates/cheatcodes/src/script.rs
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,11 @@ impl ScriptWallets {
self.inner.lock().multi_wallet.add_signer(WalletSigner::from_private_key(private_key)?);
Ok(Default::default())
}

/// Locks inner Mutex and returns all signer addresses in the [MultiWallet].
pub fn signers(&self) -> Result<Vec<Address>> {
Ok(self.inner.lock().multi_wallet.signers()?.keys().cloned().collect())
}
}

/// Sets up broadcasting from a script using `new_origin` as the sender.
Expand Down
1 change: 1 addition & 0 deletions crates/forge/tests/cli/multi_script.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,5 +61,6 @@ forgetest_async!(can_resume_multi_chain_script, |prj, cmd| {
.broadcast(ScriptOutcome::MissingWallet)
.load_private_keys(&[0, 1])
.await
.arg("--multi")
.resume(ScriptOutcome::OkBroadcast);
});
11 changes: 2 additions & 9 deletions crates/script/src/broadcast.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,6 @@
use crate::{
build::LinkedBuildData,
execute::{ExecutionArtifacts, ExecutionData},
sequence::ScriptSequenceKind,
verify::BroadcastedState,
ScriptArgs, ScriptConfig,
build::LinkedBuildData, sequence::ScriptSequenceKind, verify::BroadcastedState, ScriptArgs,
ScriptConfig,
};

use super::receipts;
Expand Down Expand Up @@ -170,8 +167,6 @@ pub struct BundledState {
pub script_config: ScriptConfig,
pub script_wallets: ScriptWallets,
pub build_data: LinkedBuildData,
pub execution_data: ExecutionData,
pub execution_artifacts: ExecutionArtifacts,
pub sequence: ScriptSequenceKind,
}

Expand Down Expand Up @@ -408,8 +403,6 @@ impl BundledState {
args: self.args,
script_config: self.script_config,
build_data: self.build_data,
execution_data: self.execution_data,
execution_artifacts: self.execution_artifacts,
sequence: self.sequence,
})
}
Expand Down
112 changes: 110 additions & 2 deletions crates/script/src/build.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,20 @@
use crate::{execute::LinkedState, ScriptArgs, ScriptConfig};
use crate::{
broadcast::BundledState,
execute::LinkedState,
multi_sequence::MultiChainSequence,
sequence::{ScriptSequence, ScriptSequenceKind},
ScriptArgs, ScriptConfig,
};

use alloy_primitives::{Address, Bytes};
use ethers_providers::Middleware;
use eyre::{Context, OptionExt, Result};
use foundry_cheatcodes::ScriptWallets;
use foundry_cli::utils::get_cached_entry_by_name;
use foundry_common::{
compile::{self, ContractSources, ProjectCompiler},
provider::ethers::try_get_http_provider,
types::ToAlloy,
ContractsByArtifact,
};
use foundry_compilers::{
Expand All @@ -16,7 +25,7 @@ use foundry_compilers::{
ArtifactId,
};
use foundry_linking::{LinkOutput, Linker};
use std::str::FromStr;
use std::{str::FromStr, sync::Arc};

/// Container for the compiled contracts.
pub struct BuildData {
Expand Down Expand Up @@ -245,4 +254,103 @@ impl CompiledState {

Ok(LinkedState { args, script_config, script_wallets, build_data })
}

/// Tries loading the resumed state from the cache files, skipping simulation stage.
pub async fn resume(self) -> Result<BundledState> {
let chain = if self.args.multi {
None
} else {
let fork_url = self.script_config.evm_opts.fork_url.clone().ok_or_eyre("Missing --fork-url field, if you were trying to broadcast a multi-chain sequence, please use --multi flag")?;
let provider = Arc::new(try_get_http_provider(fork_url)?);
Some(provider.get_chainid().await?.as_u64())
};

let sequence = match self.try_load_sequence(chain, false) {
Ok(sequence) => sequence,
Err(_) => {
// If the script was simulated, but there was no attempt to broadcast yet,
// try to read the script sequence from the `dry-run/` folder
let mut sequence = self.try_load_sequence(chain, true)?;

// If sequence was in /dry-run, Update its paths so it is not saved into /dry-run
// this time as we are about to broadcast it.
sequence.update_paths_to_broadcasted(
&self.script_config.config,
&self.args.sig,
&self.build_data.target,
)?;

sequence.save(true, true)?;
sequence
}
};

let (args, build_data, script_wallets, script_config) = if !self.args.unlocked {
let mut froms = sequence.sequences().iter().flat_map(|s| {
s.transactions
.iter()
.skip(s.receipts.len())
.map(|t| t.transaction.from().expect("from is missing in script artifact"))
});

let available_signers = self
.script_wallets
.signers()
.map_err(|e| eyre::eyre!("Failed to get available signers: {}", e))?;

if !froms.all(|from| available_signers.contains(&from.to_alloy())) {
// IF we are missing required signers, execute script as we might need to collect
// private keys from the execution.
let executed = self.link()?.prepare_execution().await?.execute().await?;
(
executed.args,
executed.build_data.build_data,
executed.script_wallets,
executed.script_config,
)
} else {
(self.args, self.build_data, self.script_wallets, self.script_config)
}
} else {
(self.args, self.build_data, self.script_wallets, self.script_config)
};

// Collect libraries from sequence and link contracts with them.
let libraries = match sequence {
ScriptSequenceKind::Single(ref seq) => Libraries::parse(&seq.libraries)?,
// Library linking is not supported for multi-chain sequences
ScriptSequenceKind::Multi(_) => Libraries::default(),
};

let linked_build_data = build_data.link_with_libraries(libraries)?;

Ok(BundledState {
args,
script_config,
script_wallets,
build_data: linked_build_data,
sequence,
})
}

fn try_load_sequence(&self, chain: Option<u64>, dry_run: bool) -> Result<ScriptSequenceKind> {
if let Some(chain) = chain {
let sequence = ScriptSequence::load(
&self.script_config.config,
&self.args.sig,
&self.build_data.target,
chain,
dry_run,
)?;
Ok(ScriptSequenceKind::Single(sequence))
} else {
let sequence = MultiChainSequence::load(
&self.script_config.config,
&self.args.sig,
&self.build_data.target,
dry_run,
)?;
Ok(ScriptSequenceKind::Multi(sequence))
}
}
}
79 changes: 40 additions & 39 deletions crates/script/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,6 @@ mod execute;
mod multi_sequence;
mod providers;
mod receipts;
mod resume;
mod runner;
mod sequence;
mod simulate;
Expand Down Expand Up @@ -219,49 +218,51 @@ impl ScriptArgs {
pub async fn run_script(self) -> Result<()> {
trace!(target: "script", "executing script command");

// Drive state machine to point at which we have everything needed for simulation/resuming.
let pre_simulation = self
.preprocess()
.await?
.compile()?
.link()?
.prepare_execution()
.await?
.execute()
.await?
.prepare_simulation()
.await?;

if pre_simulation.args.debug {
pre_simulation.run_debugger()?;
}
let compiled = self.preprocess().await?.compile()?;

if pre_simulation.args.json {
pre_simulation.show_json()?;
// Move from `CompiledState` to `BundledState` either by resuming or executing and
// simulating script.
let bundled = if compiled.args.resume || (compiled.args.verify && !compiled.args.broadcast)
{
compiled.resume().await?
} else {
pre_simulation.show_traces().await?;
}
// Drive state machine to point at which we have everything needed for simulation.
let pre_simulation = compiled
.link()?
.prepare_execution()
.await?
.execute()
.await?
.prepare_simulation()
.await?;

if pre_simulation.args.debug {
pre_simulation.run_debugger()?;
}

// Ensure that we have transactions to simulate/broadcast, otherwise exit early to avoid
// hard error.
if pre_simulation.execution_result.transactions.as_ref().map_or(true, |txs| txs.is_empty())
{
return Ok(());
}
if pre_simulation.args.json {
pre_simulation.show_json()?;
} else {
pre_simulation.show_traces().await?;
}

// Check if there are any missing RPCs and exit early to avoid hard error.
if pre_simulation.execution_artifacts.rpc_data.missing_rpc {
shell::println("\nIf you wish to simulate on-chain transactions pass a RPC URL.")?;
return Ok(());
}
// Ensure that we have transactions to simulate/broadcast, otherwise exit early to avoid
// hard error.
if pre_simulation
.execution_result
.transactions
.as_ref()
.map_or(true, |txs| txs.is_empty())
{
return Ok(());
}

// Check if there are any missing RPCs and exit early to avoid hard error.
if pre_simulation.execution_artifacts.rpc_data.missing_rpc {
shell::println("\nIf you wish to simulate on-chain transactions pass a RPC URL.")?;
return Ok(());
}

// Move from `PreSimulationState` to `BundledState` either by resuming or simulating
// transactions.
let bundled = if pre_simulation.args.resume ||
(pre_simulation.args.verify && !pre_simulation.args.broadcast)
{
pre_simulation.resume().await?
} else {
pre_simulation.args.check_contract_sizes(
&pre_simulation.execution_result,
&pre_simulation.build_data.highlevel_known_contracts,
Expand Down
106 changes: 0 additions & 106 deletions crates/script/src/resume.rs

This file was deleted.

2 changes: 0 additions & 2 deletions crates/script/src/simulate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -400,8 +400,6 @@ impl FilledTransactionsState {
script_config: self.script_config,
script_wallets: self.script_wallets,
build_data: self.build_data,
execution_data: self.execution_data,
execution_artifacts: self.execution_artifacts,
sequence,
})
}
Expand Down
Loading
Loading