forge_script/
simulate.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
use super::{
    multi_sequence::MultiChainSequence, providers::ProvidersManager, runner::ScriptRunner,
    sequence::ScriptSequenceKind, transaction::ScriptTransactionBuilder,
};
use crate::{
    broadcast::{estimate_gas, BundledState},
    build::LinkedBuildData,
    execute::{ExecutionArtifacts, ExecutionData},
    sequence::get_commit_hash,
    ScriptArgs, ScriptConfig, ScriptResult,
};
use alloy_network::TransactionBuilder;
use alloy_primitives::{map::HashMap, utils::format_units, Address, Bytes, TxKind, U256};
use dialoguer::Confirm;
use eyre::{Context, Result};
use forge_script_sequence::{ScriptSequence, TransactionWithMetadata};
use foundry_cheatcodes::Wallets;
use foundry_cli::utils::{has_different_gas_calc, now};
use foundry_common::{get_contract_name, ContractData};
use foundry_evm::traces::{decode_trace_arena, render_trace_arena};
use futures::future::{join_all, try_join_all};
use parking_lot::RwLock;
use std::{
    collections::{BTreeMap, VecDeque},
    sync::Arc,
};

/// Same as [ExecutedState](crate::execute::ExecutedState), but also contains [ExecutionArtifacts]
/// which are obtained from [ScriptResult].
///
/// Can be either converted directly to [BundledState] or driven to it through
/// [FilledTransactionsState].
pub struct PreSimulationState {
    pub args: ScriptArgs,
    pub script_config: ScriptConfig,
    pub script_wallets: Wallets,
    pub build_data: LinkedBuildData,
    pub execution_data: ExecutionData,
    pub execution_result: ScriptResult,
    pub execution_artifacts: ExecutionArtifacts,
}

impl PreSimulationState {
    /// If simulation is enabled, simulates transactions against fork and fills gas estimation and
    /// metadata. Otherwise, metadata (e.g. additional contracts, created contract names) is
    /// left empty.
    ///
    /// Both modes will panic if any of the transactions have None for the `rpc` field.
    pub async fn fill_metadata(self) -> Result<FilledTransactionsState> {
        let address_to_abi = self.build_address_to_abi_map();

        let mut transactions = self
            .execution_result
            .transactions
            .clone()
            .unwrap_or_default()
            .into_iter()
            .map(|tx| {
                let rpc = tx.rpc.expect("missing broadcastable tx rpc url");
                let sender = tx.transaction.from().expect("all transactions should have a sender");
                let nonce = tx.transaction.nonce().expect("all transactions should have a sender");
                let to = tx.transaction.to();

                let mut builder = ScriptTransactionBuilder::new(tx.transaction, rpc);

                if let Some(TxKind::Call(_)) = to {
                    builder.set_call(&address_to_abi, &self.execution_artifacts.decoder)?;
                } else {
                    builder.set_create(false, sender.create(nonce), &address_to_abi)?;
                }

                Ok(builder.build())
            })
            .collect::<Result<VecDeque<_>>>()?;

        if self.args.skip_simulation {
            sh_println!("\nSKIPPING ON CHAIN SIMULATION.")?;
        } else {
            transactions = self.simulate_and_fill(transactions).await?;
        }

        Ok(FilledTransactionsState {
            args: self.args,
            script_config: self.script_config,
            script_wallets: self.script_wallets,
            build_data: self.build_data,
            execution_artifacts: self.execution_artifacts,
            transactions,
        })
    }

    /// Builds separate runners and environments for each RPC used in script and executes all
    /// transactions in those environments.
    ///
    /// Collects gas usage and metadata for each transaction.
    pub async fn simulate_and_fill(
        &self,
        transactions: VecDeque<TransactionWithMetadata>,
    ) -> Result<VecDeque<TransactionWithMetadata>> {
        trace!(target: "script", "executing onchain simulation");

        let runners = Arc::new(
            self.build_runners()
                .await?
                .into_iter()
                .map(|(rpc, runner)| (rpc, Arc::new(RwLock::new(runner))))
                .collect::<HashMap<_, _>>(),
        );

        let mut final_txs = VecDeque::new();

        // Executes all transactions from the different forks concurrently.
        let futs = transactions
            .into_iter()
            .map(|mut transaction| async {
                let mut runner = runners.get(&transaction.rpc).expect("invalid rpc url").write();

                let tx = transaction.tx_mut();
                let to = if let Some(TxKind::Call(to)) = tx.to() { Some(to) } else { None };
                let result = runner
                    .simulate(
                        tx.from()
                            .expect("transaction doesn't have a `from` address at execution time"),
                        to,
                        tx.input().map(Bytes::copy_from_slice),
                        tx.value(),
                    )
                    .wrap_err("Internal EVM error during simulation")?;

                if !result.success {
                    return Ok((None, false, result.traces));
                }

                // Simulate mining the transaction if the user passes `--slow`.
                if self.args.slow {
                    runner.executor.env_mut().block.number += U256::from(1);
                }

                let is_noop_tx = if let Some(to) = to {
                    runner.executor.is_empty_code(to)? && tx.value().unwrap_or_default().is_zero()
                } else {
                    false
                };

                let transaction = ScriptTransactionBuilder::from(transaction)
                    .with_execution_result(&result, self.args.gas_estimate_multiplier)
                    .build();

                eyre::Ok((Some(transaction), is_noop_tx, result.traces))
            })
            .collect::<Vec<_>>();

        if self.script_config.evm_opts.verbosity > 3 {
            println!("==========================");
            println!("Simulated On-chain Traces:\n");
        }

        let mut abort = false;
        for res in join_all(futs).await {
            let (tx, is_noop_tx, mut traces) = res?;

            // Transaction will be `None`, if execution didn't pass.
            if tx.is_none() || self.script_config.evm_opts.verbosity > 3 {
                for (_, trace) in &mut traces {
                    decode_trace_arena(trace, &self.execution_artifacts.decoder).await?;
                    println!("{}", render_trace_arena(trace));
                }
            }

            if let Some(tx) = tx {
                if is_noop_tx {
                    let to = tx.contract_address.unwrap();
                    sh_warn!(
                        "Script contains a transaction to {to} which does not contain any code."
                    )?;

                    // Only prompt if we're broadcasting and we've not disabled interactivity.
                    if self.args.should_broadcast() &&
                        !self.args.non_interactive &&
                        !Confirm::new()
                            .with_prompt("Do you wish to continue?".to_string())
                            .interact()?
                    {
                        eyre::bail!("User canceled the script.");
                    }
                }

                final_txs.push_back(tx);
            } else {
                abort = true;
            }
        }

        if abort {
            eyre::bail!("Simulated execution failed.")
        }

        Ok(final_txs)
    }

    /// Build mapping from contract address to its ABI, code and contract name.
    fn build_address_to_abi_map(&self) -> BTreeMap<Address, &ContractData> {
        self.execution_artifacts
            .decoder
            .contracts
            .iter()
            .filter_map(move |(addr, contract_id)| {
                let contract_name = get_contract_name(contract_id);
                if let Ok(Some((_, data))) =
                    self.build_data.known_contracts.find_by_name_or_identifier(contract_name)
                {
                    return Some((*addr, data));
                }
                None
            })
            .collect()
    }

    /// Build [ScriptRunner] forking given RPC for each RPC used in the script.
    async fn build_runners(&self) -> Result<Vec<(String, ScriptRunner)>> {
        let rpcs = self.execution_artifacts.rpc_data.total_rpcs.clone();

        let n = rpcs.len();
        let s = if n != 1 { "s" } else { "" };
        sh_println!("\n## Setting up {n} EVM{s}.")?;

        let futs = rpcs.into_iter().map(|rpc| async move {
            let mut script_config = self.script_config.clone();
            script_config.evm_opts.fork_url = Some(rpc.clone());
            let runner = script_config.get_runner().await?;
            Ok((rpc.clone(), runner))
        });
        try_join_all(futs).await
    }
}

/// At this point we have converted transactions collected during script execution to
/// [TransactionWithMetadata] objects which contain additional metadata needed for broadcasting and
/// verification.
pub struct FilledTransactionsState {
    pub args: ScriptArgs,
    pub script_config: ScriptConfig,
    pub script_wallets: Wallets,
    pub build_data: LinkedBuildData,
    pub execution_artifacts: ExecutionArtifacts,
    pub transactions: VecDeque<TransactionWithMetadata>,
}

impl FilledTransactionsState {
    /// Bundles all transactions of the [`TransactionWithMetadata`] type in a list of
    /// [`ScriptSequence`]. List length will be higher than 1, if we're dealing with a multi
    /// chain deployment.
    ///
    /// Each transaction will be added with the correct transaction type and gas estimation.
    pub async fn bundle(self) -> Result<BundledState> {
        let is_multi_deployment = self.execution_artifacts.rpc_data.total_rpcs.len() > 1;

        if is_multi_deployment && !self.build_data.libraries.is_empty() {
            eyre::bail!("Multi-chain deployment is not supported with libraries.");
        }

        let mut total_gas_per_rpc: HashMap<String, u128> = HashMap::default();

        // Batches sequence of transactions from different rpcs.
        let mut new_sequence = VecDeque::new();
        let mut manager = ProvidersManager::default();
        let mut sequences = vec![];

        // Peeking is used to check if the next rpc url is different. If so, it creates a
        // [`ScriptSequence`] from all the collected transactions up to this point.
        let mut txes_iter = self.transactions.clone().into_iter().peekable();

        while let Some(mut tx) = txes_iter.next() {
            let tx_rpc = tx.rpc.to_owned();
            let provider_info = manager.get_or_init_provider(&tx.rpc, self.args.legacy).await?;

            if let Some(tx) = tx.tx_mut().as_unsigned_mut() {
                // Handles chain specific requirements for unsigned transactions.
                tx.set_chain_id(provider_info.chain);
            }

            if !self.args.skip_simulation {
                let tx = tx.tx_mut();

                if has_different_gas_calc(provider_info.chain) {
                    // only estimate gas for unsigned transactions
                    if let Some(tx) = tx.as_unsigned_mut() {
                        trace!("estimating with different gas calculation");
                        let gas = tx.gas.expect("gas is set by simulation.");

                        // We are trying to show the user an estimation of the total gas usage.
                        //
                        // However, some transactions might depend on previous ones. For
                        // example, tx1 might deploy a contract that tx2 uses. That
                        // will result in the following `estimate_gas` call to fail,
                        // since tx1 hasn't been broadcasted yet.
                        //
                        // Not exiting here will not be a problem when actually broadcasting,
                        // because for chains where `has_different_gas_calc`
                        // returns true, we await each transaction before
                        // broadcasting the next one.
                        if let Err(err) = estimate_gas(
                            tx,
                            &provider_info.provider,
                            self.args.gas_estimate_multiplier,
                        )
                        .await
                        {
                            trace!("gas estimation failed: {err}");

                            // Restore gas value, since `estimate_gas` will remove it.
                            tx.set_gas_limit(gas);
                        }
                    }
                }

                let total_gas = total_gas_per_rpc.entry(tx_rpc.clone()).or_insert(0);
                *total_gas += tx.gas().expect("gas is set");
            }

            new_sequence.push_back(tx);
            // We only create a [`ScriptSequence`] object when we collect all the rpc related
            // transactions.
            if let Some(next_tx) = txes_iter.peek() {
                if next_tx.rpc == tx_rpc {
                    continue;
                }
            }

            let sequence =
                self.create_sequence(is_multi_deployment, provider_info.chain, new_sequence)?;

            sequences.push(sequence);

            new_sequence = VecDeque::new();
        }

        if !self.args.skip_simulation {
            // Present gas information on a per RPC basis.
            for (rpc, total_gas) in total_gas_per_rpc {
                let provider_info = manager.get(&rpc).expect("provider is set.");

                // We don't store it in the transactions, since we want the most updated value.
                // Right before broadcasting.
                let per_gas = if let Some(gas_price) = self.args.with_gas_price {
                    gas_price.to()
                } else {
                    provider_info.gas_price()?
                };

                sh_println!("\n==========================")?;
                sh_println!("\nChain {}", provider_info.chain)?;

                sh_println!(
                    "\nEstimated gas price: {} gwei",
                    format_units(per_gas, 9)
                        .unwrap_or_else(|_| "[Could not calculate]".to_string())
                        .trim_end_matches('0')
                        .trim_end_matches('.')
                )?;
                sh_println!("\nEstimated total gas used for script: {total_gas}")?;
                sh_println!(
                    "\nEstimated amount required: {} ETH",
                    format_units(total_gas.saturating_mul(per_gas), 18)
                        .unwrap_or_else(|_| "[Could not calculate]".to_string())
                        .trim_end_matches('0')
                )?;
                sh_println!("\n==========================")?;
            }
        }

        let sequence = if sequences.len() == 1 {
            ScriptSequenceKind::Single(sequences.pop().expect("empty sequences"))
        } else {
            ScriptSequenceKind::Multi(MultiChainSequence::new(
                sequences,
                &self.args.sig,
                &self.build_data.build_data.target,
                &self.script_config.config,
                !self.args.broadcast,
            )?)
        };

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

    /// Creates a [ScriptSequence] object from the given transactions.
    fn create_sequence(
        &self,
        multi: bool,
        chain: u64,
        transactions: VecDeque<TransactionWithMetadata>,
    ) -> Result<ScriptSequence> {
        // Paths are set to None for multi-chain sequences parts, because they don't need to be
        // saved to a separate file.
        let paths = if multi {
            None
        } else {
            Some(ScriptSequence::get_paths(
                &self.script_config.config,
                &self.args.sig,
                &self.build_data.build_data.target,
                chain,
                !self.args.broadcast,
            )?)
        };

        let commit = get_commit_hash(&self.script_config.config.root.0);

        let libraries = self
            .build_data
            .libraries
            .libs
            .iter()
            .flat_map(|(file, libs)| {
                libs.iter()
                    .map(|(name, address)| format!("{}:{name}:{address}", file.to_string_lossy()))
            })
            .collect();

        let sequence = ScriptSequence {
            transactions,
            returns: self.execution_artifacts.returns.clone(),
            receipts: vec![],
            pending: vec![],
            paths,
            timestamp: now().as_secs(),
            libraries,
            chain,
            commit,
        };
        Ok(sequence)
    }
}