Skip to content

Commit

Permalink
chore: Arc ContractsByArtifact internally
Browse files Browse the repository at this point in the history
  • Loading branch information
DaniPopes committed Jun 1, 2024
1 parent a3071e5 commit eab4717
Show file tree
Hide file tree
Showing 8 changed files with 26 additions and 42 deletions.
5 changes: 2 additions & 3 deletions crates/cheatcodes/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ use semver::Version;
use std::{
collections::HashMap,
path::{Path, PathBuf},
sync::Arc,
time::Duration,
};

Expand Down Expand Up @@ -50,7 +49,7 @@ pub struct CheatsConfig {
/// Artifacts which are guaranteed to be fresh (either recompiled or cached).
/// If Some, `vm.getDeployedCode` invocations are validated to be in scope of this list.
/// If None, no validation is performed.
pub available_artifacts: Option<Arc<ContractsByArtifact>>,
pub available_artifacts: Option<ContractsByArtifact>,
/// Version of the script/test contract which is currently running.
pub running_version: Option<Version>,
}
Expand All @@ -60,7 +59,7 @@ impl CheatsConfig {
pub fn new(
config: &Config,
evm_opts: EvmOpts,
available_artifacts: Option<Arc<ContractsByArtifact>>,
available_artifacts: Option<ContractsByArtifact>,
script_wallets: Option<ScriptWallets>,
running_version: Option<Version>,
) -> Self {
Expand Down
39 changes: 14 additions & 25 deletions crates/common/src/contracts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,7 @@ use foundry_compilers::{
},
ArtifactId,
};
use std::{
collections::BTreeMap,
ops::{Deref, DerefMut},
str::FromStr,
};
use std::{collections::BTreeMap, ops::Deref, str::FromStr, sync::Arc};

/// Libraries' runtime code always starts with the following instruction:
/// `PUSH20 0x0000000000000000000000000000000000000000`
Expand All @@ -24,7 +20,7 @@ const CALL_PROTECTION_BYTECODE_PREFIX: [u8; 21] =
hex!("730000000000000000000000000000000000000000");

/// Container for commonly used contract data.
#[derive(Debug, Clone)]
#[derive(Debug)]
pub struct ContractData {
/// Contract name.
pub name: String,
Expand Down Expand Up @@ -52,26 +48,25 @@ type ArtifactWithContractRef<'a> = (&'a ArtifactId, &'a ContractData);

/// Wrapper type that maps an artifact to a contract ABI and bytecode.
#[derive(Clone, Default, Debug)]
pub struct ContractsByArtifact(pub BTreeMap<ArtifactId, ContractData>);
pub struct ContractsByArtifact(Arc<BTreeMap<ArtifactId, ContractData>>);

impl ContractsByArtifact {
/// Creates a new instance by collecting all artifacts with present bytecode from an iterator.
///
/// It is recommended to use this method with an output of
/// [foundry_linking::Linker::get_linked_artifacts].
pub fn new(artifacts: impl IntoIterator<Item = (ArtifactId, CompactContractBytecode)>) -> Self {
Self(
artifacts
.into_iter()
.filter_map(|(id, artifact)| {
let name = id.name.clone();

let CompactContractBytecode { abi, bytecode, deployed_bytecode } = artifact;

Some((id, ContractData { name, abi: abi?, bytecode, deployed_bytecode }))
})
.collect(),
)
let map = artifacts
.into_iter()
.filter_map(|(id, artifact)| {
let name = id.name.clone();

let CompactContractBytecode { abi, bytecode, deployed_bytecode } = artifact;

Some((id, ContractData { name, abi: abi?, bytecode, deployed_bytecode }))
})
.collect();
Self(Arc::new(map))
}

/// Finds a contract which has a similar bytecode as `code`.
Expand Down Expand Up @@ -232,12 +227,6 @@ impl Deref for ContractsByArtifact {
}
}

impl DerefMut for ContractsByArtifact {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}

/// Wrapper type that maps an address to a contract identifier and contract ABI.
pub type ContractsByAddress = BTreeMap<Address, (String, JsonAbi)>;

Expand Down
2 changes: 1 addition & 1 deletion crates/forge/src/multi_runner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ impl MultiContractRunner {
self.output.artifact_ids().collect(),
);
let linked_contracts = linker.get_linked_artifacts(&contract.libraries).unwrap_or_default();
let known_contracts = Arc::new(ContractsByArtifact::new(linked_contracts));
let known_contracts = ContractsByArtifact::new(linked_contracts);

let cheats_config = CheatsConfig::new(
&self.config,
Expand Down
8 changes: 3 additions & 5 deletions crates/forge/src/result.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
//! Test outcomes.

use crate::gas_report::GasReport;
use alloy_primitives::{Address, Log};
use foundry_common::{
evm::Breakpoints, get_contract_name, get_file_name, shell, ContractsByArtifact,
Expand All @@ -16,13 +17,10 @@ use serde::{Deserialize, Serialize};
use std::{
collections::{BTreeMap, HashMap},
fmt::{self, Write},
sync::Arc,
time::Duration,
};
use yansi::Paint;

use crate::gas_report::GasReport;

/// The aggregated result of a test run.
#[derive(Clone, Debug)]
pub struct TestOutcome {
Expand Down Expand Up @@ -202,7 +200,7 @@ pub struct SuiteResult {
pub libraries: Libraries,
/// Contracts linked with correct libraries.
#[serde(skip)]
pub known_contracts: Arc<ContractsByArtifact>,
pub known_contracts: ContractsByArtifact,
}

impl SuiteResult {
Expand All @@ -211,7 +209,7 @@ impl SuiteResult {
test_results: BTreeMap<String, TestResult>,
warnings: Vec<String>,
libraries: Libraries,
known_contracts: Arc<ContractsByArtifact>,
known_contracts: ContractsByArtifact,
) -> Self {
Self { duration, test_results, warnings, libraries, known_contracts }
}
Expand Down
3 changes: 1 addition & 2 deletions crates/forge/src/runner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@ use std::{
borrow::Cow,
cmp::min,
collections::{BTreeMap, HashMap},
sync::Arc,
time::Instant,
};

Expand Down Expand Up @@ -258,7 +257,7 @@ impl<'a> ContractRunner<'a> {
mut self,
filter: &dyn TestFilter,
test_options: &TestOptions,
known_contracts: Arc<ContractsByArtifact>,
known_contracts: ContractsByArtifact,
handle: &tokio::runtime::Handle,
) -> SuiteResult {
info!("starting tests");
Expand Down
3 changes: 1 addition & 2 deletions crates/script/src/build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -146,10 +146,9 @@ impl LinkedBuildData {
}

/// Fetches target bytecode from linked contracts.
pub fn get_target_contract(&self) -> Result<ContractData> {
pub fn get_target_contract(&self) -> Result<&ContractData> {
self.known_contracts
.get(&self.build_data.target)
.cloned()
.ok_or_eyre("target not found in linked artifacts")
}
}
Expand Down
4 changes: 2 additions & 2 deletions crates/script/src/execute.rs
Original file line number Diff line number Diff line change
Expand Up @@ -73,13 +73,13 @@ impl LinkedState {
args,
script_config,
script_wallets,
build_data,
execution_data: ExecutionData {
func,
calldata,
bytecode: bytecode.clone(),
abi: target_contract.abi,
abi: target_contract.abi.clone(),
},
build_data,
})
}
}
Expand Down
4 changes: 2 additions & 2 deletions crates/script/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ use foundry_evm::{
};
use foundry_wallets::MultiWalletOpts;
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, sync::Arc};
use std::collections::HashMap;
use yansi::Paint;

mod broadcast;
Expand Down Expand Up @@ -591,7 +591,7 @@ impl ScriptConfig {
CheatsConfig::new(
&self.config,
self.evm_opts.clone(),
Some(Arc::new(known_contracts)),
Some(known_contracts),
Some(script_wallets),
Some(target.version),
)
Expand Down

0 comments on commit eab4717

Please sign in to comment.