Merge branch 'foundry-update-4af6cfa' into fix/verify-upstream

crates/evm/evm/src/executors/invariant/replay.rs

@@ -17,7 +17,7 @@ use indicatif::ProgressBar;
 use parking_lot::RwLock;
 use proptest::test_runner::TestError;
 use revm::primitives::U256;
-use std::sync::Arc;
+use std::{collections::HashMap, sync::Arc};
 
 /// Replays a call sequence for collecting logs and traces.
 /// Returns counterexample to be used when the call sequence is a failed scenario.
@@ -57,7 +57,11 @@ pub fn replay_run(
         }
 
         // Identify newly generated contracts, if they exist.
-        ided_contracts.extend(load_contracts(call_result.traces.as_slice(), known_contracts));
+        ided_contracts.extend(load_contracts(
+            call_result.traces.as_slice(),
+            known_contracts,
+            &HashMap::new(),
+        ));
 
         // Create counter example to be used in failed case.
         counterexample_sequence.push(BaseCounterExample::from_invariant_call(

crates/evm/evm/src/executors/mod.rs

@@ -401,7 +401,7 @@ impl Executor {
     pub fn transact_with_env(&mut self, mut env: EnvWithHandlerCfg) -> eyre::Result<RawCallResult> {
         let mut inspector = self.inspector.clone();
         let backend = &mut self.backend;
-        let result = match self.zk_tx.take() {
+        let result_and_state = match self.zk_tx.take() {
             None => backend.inspect(&mut env, &mut inspector)?,
             Some(zk_tx) => backend.inspect_ref_zk(
                 &mut env,
@@ -411,9 +411,24 @@ impl Executor {
                 Some(zk_tx.factory_deps),
             )?,
         };
+        let mut result = convert_executed_result(
+            env,
+            inspector,
+            result_and_state.clone(),
+            backend.has_snapshot_failure(),
+        )?;
+        let state = result_and_state.state.clone();
+        if let Some(traces) = &mut result.traces {
+            for trace_node in traces.nodes() {
+                if let Some(account_info) = state.get(&trace_node.trace.address) {
+                    result.deployments.insert(
+                        trace_node.trace.address,
+                        account_info.info.code.clone().unwrap_or_default().bytes(),
+                    );
+                }
+            }
+        }
 
-        let mut result =
-            convert_executed_result(env, inspector, result, backend.has_snapshot_failure())?;
         self.commit(&mut result);
         Ok(result)
     }
@@ -724,11 +739,14 @@ pub struct RawCallResult {
     pub out: Option<Output>,
     /// The chisel state
     pub chisel_state: Option<(Vec<U256>, Vec<u8>, InstructionResult)>,
+    /// The deployments generated during the call
+    pub deployments: HashMap<Address, Bytes>,
 }
 
 impl Default for RawCallResult {
     fn default() -> Self {
         Self {
+            deployments: HashMap::new(),
             exit_reason: InstructionResult::Continue,
             reverted: false,
             has_snapshot_failure: false,
@@ -886,6 +904,7 @@ fn convert_executed_result(
     };
 
     Ok(RawCallResult {
+        deployments: HashMap::new(),
         exit_reason,
         reverted: !matches!(exit_reason, return_ok!()),
         has_snapshot_failure,

crates/evm/traces/src/decoder/mod.rs

@@ -240,10 +240,9 @@ impl CallTraceDecoder {
             .nodes()
             .iter()
             .map(|node| {
-                (
-                    &node.trace.address,
-                    node.trace.kind.is_any_create().then_some(&node.trace.output[..]),
-                )
+                let address = &node.trace.address;
+                let output = node.trace.kind.is_any_create().then_some(&node.trace.output[..]);
+                (address, output)
             })
             .filter(|&(address, _)| {
                 !self.labels.contains_key(address) || !self.contracts.contains_key(address)

crates/evm/traces/src/identifier/local.rs

@@ -1,16 +1,18 @@
 use super::{AddressIdentity, TraceIdentifier};
 use alloy_json_abi::JsonAbi;
-use alloy_primitives::Address;
+use alloy_primitives::{Address, Bytes};
 use foundry_common::contracts::{bytecode_diff_score, ContractsByArtifact};
 use foundry_compilers::ArtifactId;
-use std::borrow::Cow;
+use std::{borrow::Cow, collections::HashMap};
 
 /// A trace identifier that tries to identify addresses using local contracts.
 pub struct LocalTraceIdentifier<'a> {
     /// Known contracts to search through.
     known_contracts: &'a ContractsByArtifact,
     /// Vector of pairs of artifact ID and the runtime code length of the given artifact.
     ordered_ids: Vec<(&'a ArtifactId, usize)>,
+    /// Deployments generated during the setup
+    pub deployments: HashMap<Address, Bytes>,
 }
 
 impl<'a> LocalTraceIdentifier<'a> {
@@ -23,7 +25,7 @@ impl<'a> LocalTraceIdentifier<'a> {
             .map(|(id, bytecode)| (id, bytecode.len()))
             .collect::<Vec<_>>();
         ordered_ids.sort_by_key(|(_, len)| *len);
-        Self { known_contracts, ordered_ids }
+        Self { known_contracts, ordered_ids, deployments: HashMap::new() }
     }
 
     /// Returns the known contracts.
@@ -116,9 +118,10 @@ impl TraceIdentifier for LocalTraceIdentifier<'_> {
         addresses
             .filter_map(|(address, code)| {
                 let _span = trace_span!(target: "evm::traces", "identify", %address).entered();
-
                 trace!(target: "evm::traces", "identifying");
-                let (id, abi) = self.identify_code(code?)?;
+                let (id, abi) = self.identify_code(code?).or_else(|| {
+                    self.deployments.get(address).and_then(|bytes| self.identify_code(bytes))
+                })?;
                 trace!(target: "evm::traces", id=%id.identifier(), "identified");
 
                 Some(AddressIdentity {

crates/evm/traces/src/lib.rs

@@ -8,12 +8,12 @@
 #[macro_use]
 extern crate tracing;
 
-use alloy_primitives::LogData;
+use alloy_primitives::{Address, Bytes, LogData};
 use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
 use foundry_evm_core::constants::CHEATCODE_ADDRESS;
 use futures::{future::BoxFuture, FutureExt};
 use serde::{Deserialize, Serialize};
-use std::fmt::Write;
+use std::{collections::HashMap, fmt::Write};
 use yansi::{Color, Paint};
 
 /// Call trace address identifiers.
@@ -298,13 +298,19 @@ fn trace_color(trace: &CallTrace) -> Color {
 pub fn load_contracts<'a>(
     traces: impl IntoIterator<Item = &'a CallTraceArena>,
     known_contracts: &ContractsByArtifact,
+    deployments: &HashMap<Address, Bytes>,
 ) -> ContractsByAddress {
     let mut local_identifier = LocalTraceIdentifier::new(known_contracts);
+    local_identifier.deployments = deployments.clone();
     let decoder = CallTraceDecoder::new();
     let mut contracts = ContractsByAddress::new();
     for trace in traces {
-        for address in local_identifier.identify_addresses(decoder.trace_addresses(trace)) {
-            if let (Some(contract), Some(abi)) = (address.contract, address.abi) {
+        let identified_addresses =
+            local_identifier.identify_addresses(decoder.trace_addresses(trace));
+        for address in identified_addresses {
+            let contract = address.contract;
+            let abi = address.abi;
+            if let (Some(contract), Some(abi)) = (contract, abi) {
                 contracts.insert(address.address, (contract, abi.into_owned()));
             }
         }

crates/forge/bin/cmd/coverage.rs

@@ -244,7 +244,7 @@ impl CoverageArgs {
                 ..Default::default()
             })
             .set_coverage(true)
-            .build(&root, output, None, env, evm_opts, DualCompiledContracts::default(), false)?;
+            .build(&root, output, None, env, evm_opts, DualCompiledContracts::default())?;
 
         let known_contracts = runner.known_contracts.clone();
 

crates/forge/bin/cmd/test/mod.rs

@@ -281,9 +281,12 @@ impl TestArgs {
             let zk_compiler = ProjectCompiler::new()
                 .quiet_if(self.json || self.opts.silent)
                 .files(sources_to_compile);
+
+            let zk_output =
+                zk_compiler.zksync_compile(&zk_project, config.zksync.avoid_contracts())?;
+            let dual_compiled_contracts =
+                DualCompiledContracts::new(&output, &zk_output, &project.paths);
 
-            let zk_output = zk_compiler.zksync_compile(&zk_project, None)?;
-            let dual_compiled_contracts = DualCompiledContracts::new(&output, &zk_output, &project.paths);
             (Some(zk_output), Some(dual_compiled_contracts))
         } else {
             (None, None)
@@ -331,7 +334,6 @@ impl TestArgs {
                 env,
                 evm_opts,
                 dual_compiled_contracts.unwrap_or_default(),
-                config.zksync.run_in_zk_mode(),
             )?;
 
         if let Some(debug_test_pattern) = &self.debug {

crates/forge/src/multi_runner.rs

@@ -9,8 +9,10 @@ use alloy_primitives::{Address, Bytes, U256};
 use eyre::Result;
 use foundry_common::{get_contract_name, ContractsByArtifact, TestFunctionExt};
 use foundry_compilers::{
-    artifacts::Libraries, compilers::Compiler, Artifact, ArtifactId, ProjectCompileOutput,
+    artifacts::{CompactBytecode, CompactContractBytecode, CompactDeployedBytecode, Libraries},
+    compilers::Compiler,
     zksync::compile::output::ProjectCompileOutput as ZkProjectCompileOutput,
+    Artifact, ArtifactId, ProjectCompileOutput,
 };
 use foundry_config::Config;
 use foundry_evm::{
@@ -376,9 +378,9 @@ impl MultiContractRunnerBuilder {
         env: revm::primitives::Env,
         evm_opts: EvmOpts,
         dual_compiled_contracts: DualCompiledContracts,
-        use_zk: bool,
     ) -> Result<MultiContractRunner> {
-        // TODO: Use zk_output
+        let use_zk = zk_output.is_some();
+        let mut known_contracts = ContractsByArtifact::default();
         let output = output.with_stripped_file_prefixes(root);
         let linker = Linker::new(root, output.artifact_ids().collect());
 
@@ -419,7 +421,47 @@ impl MultiContractRunnerBuilder {
             }
         }
 
-        let known_contracts = ContractsByArtifact::new(linked_contracts);
+        if !use_zk {
+            known_contracts = ContractsByArtifact::new(linked_contracts);
+        } else if let Some(zk_output) = zk_output {
+            let zk_contracts = zk_output.with_stripped_file_prefixes(root).into_artifacts();
+            let mut zk_contracts_map = BTreeMap::new();
+
+            for (id, contract) in zk_contracts {
+                if let Some(metadata) = contract.metadata {
+                    if let Some(solc_metadata_value) =
+                        metadata.get("solc_metadata").and_then(serde_json::Value::as_str)
+                    {
+                        if let Ok(solc_metadata_json) =
+                            serde_json::from_str::<serde_json::Value>(solc_metadata_value)
+                        {
+                            let abi: JsonAbi = JsonAbi::from_json_str(
+                                &solc_metadata_json["output"]["abi"].to_string(),
+                            )?;
+                            let bytecode = contract.bytecode.as_ref();
+
+                            if let Some(bytecode_object) = bytecode.map(|b| b.object.clone()) {
+                                let compact_bytecode = CompactBytecode {
+                                    object: bytecode_object.clone(),
+                                    source_map: None,
+                                    link_references: BTreeMap::new(),
+                                };
+                                let compact_contract = CompactContractBytecode {
+                                    abi: Some(abi),
+                                    bytecode: Some(compact_bytecode.clone()),
+                                    deployed_bytecode: Some(CompactDeployedBytecode {
+                                        bytecode: Some(compact_bytecode),
+                                        immutable_references: BTreeMap::new(),
+                                    }),
+                                };
+                                zk_contracts_map.insert(id.clone(), compact_contract);
+                            }
+                        }
+                    }
+                }
+            }
+            known_contracts = ContractsByArtifact::new(zk_contracts_map);
+        }
 
         Ok(MultiContractRunner {
             contracts: deployable_contracts,

crates/forge/src/result.rs

@@ -1,7 +1,7 @@
 //! Test outcomes.
 
 use crate::gas_report::GasReport;
-use alloy_primitives::{Address, Log};
+use alloy_primitives::{Address, Bytes, Log};
 use foundry_common::{evm::Breakpoints, get_contract_name, get_file_name, shell};
 use foundry_evm::{
     coverage::HitMaps,
@@ -508,6 +508,8 @@ impl TestKind {
 
 #[derive(Clone, Debug, Default)]
 pub struct TestSetup {
+    /// Deployments generated during the setup
+    pub deployments: HashMap<Address, Bytes>,
     /// The address at which the test contract was deployed
     pub address: Address,
     /// The logs emitted during setup
@@ -549,14 +551,24 @@ impl TestSetup {
     }
 
     pub fn success(
+        deployments: HashMap<Address, Bytes>,
         address: Address,
         logs: Vec<Log>,
         traces: Traces,
         labeled_addresses: HashMap<Address, String>,
         coverage: Option<HitMaps>,
         fuzz_fixtures: FuzzFixtures,
     ) -> Self {
-        Self { address, logs, traces, labeled_addresses, reason: None, coverage, fuzz_fixtures }
+        Self {
+            deployments,
+            address,
+            logs,
+            traces,
+            labeled_addresses,
+            reason: None,
+            coverage,
+            fuzz_fixtures,
+        }
     }
 
     pub fn failed_with(
@@ -566,6 +578,7 @@ impl TestSetup {
         reason: String,
     ) -> Self {
         Self {
+            deployments: HashMap::new(),
             address: Address::ZERO,
             logs,
             traces,

crates/forge/src/runner.rs

@@ -147,6 +147,8 @@ impl<'a> ContractRunner<'a> {
         // construction
         self.executor.set_balance(address, self.initial_balance)?;
 
+        let mut zk_setup_deployments = HashMap::new();
+
         // Deploy the test contract
         match self.executor.deploy(
             self.sender,
@@ -176,7 +178,8 @@ impl<'a> ContractRunner<'a> {
             trace!("calling setUp");
             let res = self.executor.setup(None, address, Some(self.revert_decoder));
             let (setup_logs, setup_traces, labeled_addresses, reason, coverage) = match res {
-                Ok(RawCallResult { traces, labels, logs, coverage, .. }) => {
+                Ok(RawCallResult { traces, labels, logs, coverage, deployments, .. }) => {
+                    zk_setup_deployments.extend(deployments);
                     trace!(%address, "successfully called setUp");
                     (logs, traces, labels, None, coverage)
                 }
@@ -195,6 +198,7 @@ impl<'a> ContractRunner<'a> {
             logs.extend(setup_logs);
 
             TestSetup {
+                deployments: zk_setup_deployments,
                 address,
                 logs,
                 traces,
@@ -205,6 +209,7 @@ impl<'a> ContractRunner<'a> {
             }
         } else {
             TestSetup::success(
+                zk_setup_deployments,
                 address,
                 logs,
                 traces,
@@ -384,8 +389,13 @@ impl<'a> ContractRunner<'a> {
             find_time,
         );
 
-        let identified_contracts = has_invariants
-            .then(|| load_contracts(setup.traces.iter().map(|(_, t)| t), &known_contracts));
+        let identified_contracts = has_invariants.then(|| {
+            load_contracts(
+                setup.traces.iter().map(|(_, t)| t),
+                &known_contracts,
+                &setup.deployments,
+            )
+        });
         let test_results = functions
             .par_iter()
             .map(|&func| {

crates/script/src/broadcast.rs

@@ -127,7 +127,8 @@ pub async fn send_transaction(
             debug!("sending transaction: {:?}", tx);
 
             let signed = if let Some(zk) = zk {
-                let signer = signer.signer_by_address(from).ok_or(eyre::eyre!("Signer not found"))?;
+                let signer =
+                    signer.signer_by_address(from).ok_or(eyre::eyre!("Signer not found"))?;
 
                 let (deploy_request, signable) = convert_to_zksync(&provider, tx, zk).await?;
                 let mut signable = signable.to_signable_tx();