Refactor

Cargo.toml

@@ -28,7 +28,7 @@ resolver = "2"
 [workspace.package]
 version = "0.2.0"
 edition = "2021"
-rust-version = "1.74" # Remember to update clippy.toml as well
+rust-version = "1.74"  # Remember to update clippy.toml as well
 authors = ["Foundry Contributors"]
 license = "MIT OR Apache-2.0"
 homepage = "https://github.com/foundry-rs/foundry"
@@ -173,7 +173,10 @@ alloy-rlp = "0.3.3"
 solang-parser = "=0.3.3"
 
 ## misc
-chrono = { version = "0.4", default-features = false, features = ["clock", "std"] }
+chrono = { version = "0.4", default-features = false, features = [
+ "clock",
+ "std",
+] }
 color-eyre = "0.6"
 derive_more = "0.99"
 eyre = "0.6"
@@ -226,3 +229,5 @@ revm = { git = "https://github.com/bluealloy/revm", branch = "reth_freeze" }
 revm-primitives = { git = "https://github.com/bluealloy/revm", branch = "reth_freeze" }
 revm-interpreter = { git = "https://github.com/bluealloy/revm", branch = "reth_freeze" }
 revm-precompile = { git = "https://github.com/bluealloy/revm", branch = "reth_freeze" }
+
+foundry-compilers = { git = "https://github.com/foundry-rs/compilers", branch = "main" }

crates/forge/bin/cmd/script/build.rs

@@ -13,9 +13,9 @@ use foundry_compilers::{
  cache::SolFilesCache,
  contracts::ArtifactContracts,
  info::ContractInfo,
- Artifact, ArtifactId, Project, ProjectCompileOutput,
+ ArtifactId, Project, ProjectCompileOutput,
 };
-use std::{collections::BTreeMap, str::FromStr};
+use std::str::FromStr;
 
 impl ScriptArgs {
  /// Compiles the file or project and the verify metadata.
@@ -63,10 +63,12 @@ impl ScriptArgs {
  let target = self.find_target(&project, &contracts)?.clone();
  script_config.target_contract = Some(target.clone());
 
- let mut output = self.link(
+ let libraries = script_config.config.solc_settings()?.libraries;
+
+ let mut output = self.link_script_target(
  project,
  contracts,
- script_config.config.parsed_libraries()?,
+ libraries,
  script_config.evm_opts.sender,
  script_config.sender_nonce,
  target,
@@ -123,7 +125,7 @@ impl ScriptArgs {
  target.ok_or_eyre(format!("Could not find target contract: {}", target_fname))
  }
 
- pub fn link(
+ pub fn link_script_target(
  &self,
  project: Project,
  contracts: ArtifactContracts,
@@ -132,31 +134,13 @@ impl ScriptArgs {
  nonce: u64,
  target: ArtifactId,
  ) -> Result<BuildOutput> {
- let LinkOutput { libs_to_deploy, contracts: linked_contracts, predeployed_libs } =
- link_with_nonce_or_address(&contracts, &libraries, sender, nonce, &target)?;
-
- // Merge with user provided libraries
- let mut new_libraries = Libraries { libs: BTreeMap::new() };
- for (id, address) in libs_to_deploy.iter().chain(predeployed_libs.iter()) {
- new_libraries
- .libs
- .entry(id.source.clone())
- .or_default()
- .insert(id.name.split('.').next().unwrap().to_owned(), address.to_string());
- }
-
- let predeploy_libraries = libs_to_deploy
- .into_iter()
- .map(|(id, _)| {
- linked_contracts
- .get(id)
- .unwrap()
- .get_bytecode_bytes()
- .unwrap_or_else(|| panic!("bytecode for {} is unlinked", id.name))
- .into_owned()
- })
- .collect();
+ let LinkOutput {
+ libs_to_deploy: predeploy_libraries,
+ contracts: linked_contracts,
+ libraries,
+ } = link_with_nonce_or_address(&contracts, libraries, sender, nonce, &target)?;
 
+ // Get linked target artifact
  let contract = linked_contracts
  .get(&target)
  .ok_or_eyre("Target contract not found in artifacts")?
@@ -180,7 +164,7 @@ impl ScriptArgs {
  predeploy_libraries,
  sources: Default::default(),
  project,
- libraries: new_libraries,
+ libraries,
  })
  }
 

crates/forge/bin/cmd/script/cmd.rs

@@ -276,14 +276,19 @@ impl ScriptArgs {
  // We might have predeployed libraries from the broadcasting, so we need to
  // relink the contracts with them, since their mapping is
  // not included in the solc cache files.
- let BuildOutput { highlevel_known_contracts, .. } = self.link(
- project,
- default_known_contracts,
- Libraries::parse(&deployment_sequence.libraries)?,
- script_config.config.sender, // irrelevant, since we're not creating any
- 0, // irrelevant, since we're not creating any
- target,
- )?;
+ let BuildOutput { highlevel_known_contracts, predeploy_libraries, .. } = self
+ .link_script_target(
+ project,
+ default_known_contracts,
+ Libraries::parse(&deployment_sequence.libraries)?,
+ script_config.config.sender, // irrelevant, since we're not creating any
+ 0, // irrelevant, since we're not creating any
+ target,
+ )?;
+
+ if predeploy_libraries.len() > 0 {
+ eyre::bail!("Incomplete set of libraries in deployment artifact.");
+ }
 
  verify.known_contracts = flatten_contracts(&highlevel_known_contracts, false);
 
@@ -317,7 +322,7 @@ impl ScriptArgs {
 
  let BuildOutput {
  libraries, contract, highlevel_known_contracts, predeploy_libraries, ..
- } = self.link(
+ } = self.link_script_target(
  project,
  default_known_contracts,
  script_config.config.parsed_libraries()?,

crates/forge/src/link.rs

@@ -1,11 +1,8 @@
-use alloy_primitives::Address;
+use alloy_primitives::{Address, Bytes};
 use eyre::Result;
-use foundry_compilers::{artifacts::Libraries, contracts::ArtifactContracts, ArtifactId};
+use foundry_compilers::{artifacts::Libraries, contracts::ArtifactContracts, Artifact, ArtifactId};
 use semver::Version;
-use std::{
- collections::{HashMap, HashSet},
- str::FromStr,
-};
+use std::{collections::HashSet, path::PathBuf, str::FromStr};
 
 /// Finds an [ArtifactId] object in the given [ArtifactContracts] keys which corresponds to the
 /// library path in the form of "./path/to/Lib.sol:Lib"
@@ -28,7 +25,7 @@ fn find_artifact_id_by_library_path<'a>(
  continue;
  }
 
- if !(id.source.ends_with(file)) {
+ if id.source != PathBuf::from(file) {
  continue;
  }
 
@@ -56,76 +53,94 @@ pub fn collect_dependencies<'a>(
  }
 }
 
+/// Links given artifacts with given library addresses.
+///
+/// Artifacts returned by this function may still be partially unlinked if some of their
+/// dependencies weren't present in `libraries`.
+pub fn link(contracts: &ArtifactContracts, libraries: &Libraries) -> Result<ArtifactContracts> {
+ contracts
+ .iter()
+ .map(|(id, contract)| {
+ let mut contract = contract.clone();
+
+ for (file, libs) in &libraries.libs {
+ for (name, address) in libs {
+ let address = Address::from_str(address)?;
+ if let Some(bytecode) = contract.bytecode.as_mut() {
+ bytecode.link(file.to_string_lossy(), name, address);
+ }
+ if let Some(deployed_bytecode) =
+ contract.deployed_bytecode.as_mut().and_then(|b| b.bytecode.as_mut())
+ {
+ deployed_bytecode.link(file.to_string_lossy(), name, address);
+ }
+ }
+ }
+ Ok((id.clone(), contract))
+ })
+ .collect()
+}
+
 /// Output of the `link_with_nonce_or_address`
-pub struct LinkOutput<'a> {
+pub struct LinkOutput {
  /// [ArtifactContracts] object containing all artifacts linked with known libraries
  /// It is guaranteed to contain `target` and all it's dependencies fully linked, and any other
  /// contract may still be partially unlinked.
  pub contracts: ArtifactContracts,
  /// Vector of libraries predeployed by user (basically another form of
  /// `deployed_library_addresses`).
- pub predeployed_libs: Vec<(&'a ArtifactId, Address)>,
+ pub libraries: Libraries,
  /// Vector of libraries that need to be deployed from sender address.
  /// The order in which they appear in the vector is the order in which they should be deployed.
- pub libs_to_deploy: Vec<(&'a ArtifactId, Address)>,
+ pub libs_to_deploy: Vec<Bytes>,
 }
 
+/// Links given artifact with either given library addresses or address computed from sender and
+/// nonce.
+///
+/// Current linking implementation assumes that all link_references keys are matching the format in
+/// which sources appear in [ArtifactId] keys.
+///
+/// You should make sure to strip paths via `with_stripped_file_prefixes` for both
+/// `libraries` and `contracts` *before* invoking linker.
+///
+/// This function will always link target contract completely, however, in case of paths format
+/// mismatch it may not recognize some of the predeployed libraries, leading to more deployments
+/// than needed.
 pub fn link_with_nonce_or_address<'a>(
  contracts: &'a ArtifactContracts,
- deployed_library_addresses: &Libraries,
+ mut libraries: Libraries,
  sender: Address,
  mut nonce: u64,
  target: &'a ArtifactId,
-) -> Result<LinkOutput<'a>> {
+) -> Result<LinkOutput> {
  let mut needed_libraries = HashSet::new();
  collect_dependencies(target, contracts, &mut needed_libraries);
 
- let mut predeployed_libs = HashMap::new();
-
- // Populate predeployed libs firstly
- for (path, libs) in &deployed_library_addresses.libs {
- let path = path.to_string_lossy().to_string();
- for (name, address) in libs {
- let artifact_id = find_artifact_id_by_library_path(contracts, &path, name, None);
- if needed_libraries.contains(artifact_id) {
- let address = Address::from_str(address)?;
- predeployed_libs.insert(artifact_id, address);
- }
- }
- }
-
  let mut libs_to_deploy = Vec::new();
 
+ // If `libraries` does not contain needed dependency, compute its address and add to
+ // `libs_to_deploy`.
  for id in needed_libraries {
- if !predeployed_libs.contains_key(id) {
- libs_to_deploy.push((id, sender.create(nonce)));
+ let lib_name = id.name.split(".").next().unwrap().to_owned();
+ libraries.libs.entry(id.source.clone()).or_default().entry(lib_name).or_insert_with(|| {
+ let address = sender.create(nonce);
+ libs_to_deploy.push((id, address));
  nonce += 1;
- }
- }
 
- let predeployed_libs = predeployed_libs.into_iter().collect::<Vec<_>>();
+ address.to_checksum(None)
+ });
+ }
 
  // Link contracts
- let contracts = contracts
- .iter()
- .map(|(id, contract)| {
- let mut contract = contract.clone();
+ let contracts = link(contracts, &libraries)?;
 
- for (id, address) in libs_to_deploy.iter().chain(predeployed_libs.iter()) {
- if let Some(bytecode) = contract.bytecode.as_mut() {
- bytecode.link(id.source.to_string_lossy(), &id.name, *address);
- }
- if let Some(deployed_bytecode) =
- contract.deployed_bytecode.as_mut().and_then(|b| b.bytecode.as_mut())
- {
- deployed_bytecode.link(id.source.to_string_lossy(), &id.name, *address);
- }
- }
- (id.clone(), contract)
- })
- .collect::<ArtifactContracts>();
+ let libs_to_deploy = libs_to_deploy
+ .into_iter()
+ .map(|(id, _)| contracts.get(id).unwrap().get_bytecode_bytes().unwrap().into_owned())
+ .collect();
 
- Ok(LinkOutput { contracts, predeployed_libs, libs_to_deploy })
+ Ok(LinkOutput { contracts, libraries, libs_to_deploy })
 }
 
 #[cfg(test)]
@@ -185,7 +200,7 @@ mod tests {
 
  let LinkOutput { libs_to_deploy, .. } = link_with_nonce_or_address(
  &self.contracts,
- &Default::default(),
+ Default::default(),
  sender,
  initial_nonce,
  id,

crates/forge/src/multi_runner.rs

@@ -289,7 +289,7 @@ impl MultiContractRunnerBuilder {
 
  let LinkOutput { contracts, libs_to_deploy, .. } = link_with_nonce_or_address(
  &artifact_contracts,
- &Default::default(),
+ Default::default(),
  evm_opts.sender,
  1,
  &id,
@@ -314,18 +314,7 @@ impl MultiContractRunnerBuilder {
  if abi.constructor.as_ref().map(|c| c.inputs.is_empty()).unwrap_or(true) &&
  abi.functions().any(|func| func.name.is_test() || func.name.is_invariant_test())
  {
- deployable_contracts.insert(
- id.clone(),
- (
- abi.clone(),
- bytecode,
- libs_to_deploy
- .into_iter()
- .filter_map(|(id, _)| contracts.get(id).unwrap().bytecode.clone())
- .filter_map(|bcode| bcode.object.into_bytes())
- .collect::<Vec<_>>(),
- ),
- );
+ deployable_contracts.insert(id.clone(), (abi.clone(), bytecode, libs_to_deploy));
  }
 
  contract