feat(common): PathOrContractInfo arg type (foundry-rs#9770)

* feat(`common`): `PathOrContractInfo`

* fix

* docs

* nit

* test

* move find abi helper to ContractsByArtifact

* nit

* fix

* nit

* fix

* nit

* ensure sol file

* account for vyper contracts

* nit

Author: yash-atreya

crates/cast/bin/cmd/interface.rs

@@ -4,8 +4,10 @@ use clap::Parser;
 use eyre::{Context, Result};
 use foundry_block_explorers::Client;
 use foundry_cli::{opts::EtherscanOpts, utils::LoadConfig};
-use foundry_common::{compile::ProjectCompiler, fs, shell};
-use foundry_compilers::{info::ContractInfo, utils::canonicalize};
+use foundry_common::{
+    compile::{PathOrContractInfo, ProjectCompiler},
+    find_target_path, fs, shell, ContractsByArtifact,
+};
 use foundry_config::load_config;
 use itertools::Itertools;
 use serde_json::Value;
@@ -118,19 +120,20 @@ fn load_abi_from_artifact(path_or_contract: &str) -> Result<Vec<(JsonAbi, String
     let project = config.project()?;
     let compiler = ProjectCompiler::new().quiet(true);
 
-    let contract = ContractInfo::new(path_or_contract);
-    let target_path = if let Some(path) = &contract.path {
-        canonicalize(project.root().join(path))?
-    } else {
-        project.find_contract_path(&contract.name)?
-    };
-    let mut output = compiler.files([target_path.clone()]).compile(&project)?;
-
-    let artifact = output.remove(&target_path, &contract.name).ok_or_else(|| {
-        eyre::eyre!("Could not find artifact `{contract}` in the compiled artifacts")
-    })?;
-    let abi = artifact.abi.as_ref().ok_or_else(|| eyre::eyre!("Failed to fetch lossless ABI"))?;
-    Ok(vec![(abi.clone(), contract.name)])
+    let contract = PathOrContractInfo::from_str(path_or_contract)?;
+
+    let target_path = find_target_path(&project, &contract)?;
+    let output = compiler.files([target_path.clone()]).compile(&project)?;
+
+    let contracts_by_artifact = ContractsByArtifact::from(output);
+
+    let maybe_abi = contracts_by_artifact
+        .find_abi_by_name_or_src_path(contract.name().unwrap_or(&target_path.to_string_lossy()));
+
+    let (abi, name) =
+        maybe_abi.as_ref().ok_or_else(|| eyre::eyre!("Failed to fetch lossless ABI"))?;
+
+    Ok(vec![(abi.clone(), contract.name().unwrap_or(name).to_string())])
 }
 
 /// Fetches the ABI of a contract from Etherscan.

crates/common/src/compile.rs

@@ -15,6 +15,7 @@ use foundry_compilers::{
         solc::{Solc, SolcCompiler},
         Compiler,
     },
+    info::ContractInfo as CompilerContractInfo,
     report::{BasicStdoutReporter, NoReporter, Report},
     solc::SolcSettings,
     Artifact, Project, ProjectBuilder, ProjectCompileOutput, ProjectPathsConfig, SolcConfig,
@@ -25,6 +26,7 @@ use std::{
     fmt::Display,
     io::IsTerminal,
     path::{Path, PathBuf},
+    str::FromStr,
     time::Instant,
 };
 
@@ -540,3 +542,92 @@ pub fn with_compilation_reporter<O>(quiet: bool, f: impl FnOnce() -> O) -> O {
 
     foundry_compilers::report::with_scoped(&reporter, f)
 }
+
+/// Container type for parsing contract identifiers from CLI.
+///
+/// Passed string can be of the following forms:
+/// - `src/Counter.sol` - path to the contract file, in the case where it only contains one contract
+/// - `src/Counter.sol:Counter` - path to the contract file and the contract name
+/// - `Counter` - contract name only
+#[derive(Clone, PartialEq, Eq)]
+pub enum PathOrContractInfo {
+    /// Non-canoncalized path provided via CLI.
+    Path(PathBuf),
+    /// Contract info provided via CLI.
+    ContractInfo(CompilerContractInfo),
+}
+
+impl PathOrContractInfo {
+    /// Returns the path to the contract file if provided.
+    pub fn path(&self) -> Option<PathBuf> {
+        match self {
+            Self::Path(path) => Some(path.to_path_buf()),
+            Self::ContractInfo(info) => info.path.as_ref().map(PathBuf::from),
+        }
+    }
+
+    /// Returns the contract name if provided.
+    pub fn name(&self) -> Option<&str> {
+        match self {
+            Self::Path(_) => None,
+            Self::ContractInfo(info) => Some(&info.name),
+        }
+    }
+}
+
+impl FromStr for PathOrContractInfo {
+    type Err = eyre::Error;
+
+    fn from_str(s: &str) -> Result<Self> {
+        if let Ok(contract) = CompilerContractInfo::from_str(s) {
+            return Ok(Self::ContractInfo(contract));
+        }
+        let path = PathBuf::from(s);
+        if path.extension().is_some_and(|ext| ext == "sol" || ext == "vy") {
+            return Ok(Self::Path(path));
+        }
+        Err(eyre::eyre!("Invalid contract identifier, file is not *.sol or *.vy: {}", s))
+    }
+}
+
+impl std::fmt::Debug for PathOrContractInfo {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Path(path) => write!(f, "Path({})", path.display()),
+            Self::ContractInfo(info) => {
+                write!(f, "ContractInfo({info})")
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn parse_contract_identifiers() {
+        let t = ["src/Counter.sol", "src/Counter.sol:Counter", "Counter"];
+
+        let i1 = PathOrContractInfo::from_str(t[0]).unwrap();
+        assert_eq!(i1, PathOrContractInfo::Path(PathBuf::from(t[0])));
+
+        let i2 = PathOrContractInfo::from_str(t[1]).unwrap();
+        assert_eq!(
+            i2,
+            PathOrContractInfo::ContractInfo(CompilerContractInfo {
+                path: Some("src/Counter.sol".to_string()),
+                name: "Counter".to_string()
+            })
+        );
+
+        let i3 = PathOrContractInfo::from_str(t[2]).unwrap();
+        assert_eq!(
+            i3,
+            PathOrContractInfo::ContractInfo(CompilerContractInfo {
+                path: None,
+                name: "Counter".to_string()
+            })
+        );
+    }
+}

crates/common/src/contracts.rs

@@ -1,16 +1,24 @@
 //! Commonly used contract types and functions.
 
+use crate::compile::PathOrContractInfo;
 use alloy_json_abi::{Event, Function, JsonAbi};
 use alloy_primitives::{hex, Address, Bytes, Selector, B256};
-use eyre::Result;
+use eyre::{OptionExt, Result};
 use foundry_compilers::{
     artifacts::{
         BytecodeObject, CompactBytecode, CompactContractBytecode, CompactDeployedBytecode,
-        ContractBytecodeSome, Offsets,
+        ConfigurableContractArtifact, ContractBytecodeSome, Offsets,
     },
-    ArtifactId,
+    utils::canonicalized,
+    ArtifactId, Project, ProjectCompileOutput,
+};
+use std::{
+    collections::BTreeMap,
+    ops::Deref,
+    path::{Path, PathBuf},
+    str::FromStr,
+    sync::Arc,
 };
-use std::{collections::BTreeMap, ops::Deref, str::FromStr, sync::Arc};
 
 /// Libraries' runtime code always starts with the following instruction:
 /// `PUSH20 0x0000000000000000000000000000000000000000`
@@ -268,6 +276,17 @@ impl ContractsByArtifact {
             .map(|(_, contract)| contract.abi.clone())
     }
 
+    /// Finds abi by name or source path
+    ///
+    /// Returns the abi and the contract name.
+    pub fn find_abi_by_name_or_src_path(&self, name_or_path: &str) -> Option<(JsonAbi, String)> {
+        self.iter()
+            .find(|(artifact, _)| {
+                artifact.name == name_or_path || artifact.source == PathBuf::from(name_or_path)
+            })
+            .map(|(_, contract)| (contract.abi.clone(), contract.name.clone()))
+    }
+
     /// Flattens the contracts into functions, events and errors.
     pub fn flatten(&self) -> (BTreeMap<Selector, Function>, BTreeMap<B256, Event>, JsonAbi) {
         let mut funcs = BTreeMap::new();
@@ -288,6 +307,21 @@ impl ContractsByArtifact {
     }
 }
 
+impl From<ProjectCompileOutput> for ContractsByArtifact {
+    fn from(value: ProjectCompileOutput) -> Self {
+        Self::new(value.into_artifacts().map(|(id, ar)| {
+            (
+                id,
+                CompactContractBytecode {
+                    abi: ar.abi,
+                    bytecode: ar.bytecode,
+                    deployed_bytecode: ar.deployed_bytecode,
+                },
+            )
+        }))
+    }
+}
+
 impl Deref for ContractsByArtifact {
     type Target = BTreeMap<ArtifactId, ContractData>;
 
@@ -399,6 +433,61 @@ pub fn compact_to_contract(contract: CompactContractBytecode) -> Result<Contract
     })
 }
 
+/// Returns the canonicalized target path for the given identifier.
+pub fn find_target_path(project: &Project, identifier: &PathOrContractInfo) -> Result<PathBuf> {
+    match identifier {
+        PathOrContractInfo::Path(path) => Ok(canonicalized(project.root().join(path))),
+        PathOrContractInfo::ContractInfo(info) => {
+            let path = project.find_contract_path(&info.name)?;
+            Ok(path)
+        }
+    }
+}
+
+/// Returns the target artifact given the path and name.
+pub fn find_matching_contract_artifact(
+    output: &mut ProjectCompileOutput,
+    target_path: &Path,
+    target_name: Option<&str>,
+) -> eyre::Result<ConfigurableContractArtifact> {
+    if let Some(name) = target_name {
+        output
+            .remove(target_path, name)
+            .ok_or_eyre(format!("Could not find artifact `{name}` in the compiled artifacts"))
+    } else {
+        let possible_targets = output
+            .artifact_ids()
+            .filter(|(id, _artifact)| id.source == target_path)
+            .collect::<Vec<_>>();
+
+        if possible_targets.is_empty() {
+            eyre::bail!("Could not find artifact linked to source `{target_path:?}` in the compiled artifacts");
+        }
+
+        let (target_id, target_artifact) = possible_targets[0].clone();
+        if possible_targets.len() == 1 {
+            return Ok(target_artifact.clone());
+        }
+
+        // If all artifact_ids in `possible_targets` have the same name (without ".", indicates
+        // additional compiler profiles), it means that there are multiple contracts in the
+        // same file.
+        if !target_id.name.contains(".") &&
+            possible_targets.iter().any(|(id, _)| id.name != target_id.name)
+        {
+            eyre::bail!("Multiple contracts found in the same file, please specify the target <path>:<contract> or <contract>");
+        }
+
+        // Otherwise, we're dealing with additional compiler profiles wherein `id.source` is the
+        // same but `id.path` is different.
+        let artifact = possible_targets
+            .iter()
+            .find_map(|(id, artifact)| if id.profile == "default" { Some(*artifact) } else { None })
+            .unwrap_or(target_artifact);
+
+        Ok(artifact.clone())
+    }
+}
 #[cfg(test)]
 mod tests {
     use super::*;

crates/common/src/utils.rs

@@ -1,7 +1,6 @@
 //! Uncategorised utilities.
 
 use alloy_primitives::{keccak256, B256, U256};
-
 /// Block on a future using the current tokio runtime on the current thread.
 pub fn block_on<F: std::future::Future>(future: F) -> F::Output {
     block_on_handle(&tokio::runtime::Handle::current(), future)

crates/config/src/lib.rs

@@ -122,7 +122,7 @@ mod bind_json;
 use bind_json::BindJsonConfig;
 
 mod compilation;
-use compilation::{CompilationRestrictions, SettingsOverrides};
+pub use compilation::{CompilationRestrictions, SettingsOverrides};
 
 /// Foundry configuration
 ///

crates/forge/bin/cmd/inspect.rs

@@ -5,27 +5,29 @@ use comfy_table::{modifiers::UTF8_ROUND_CORNERS, Cell, Table};
 use eyre::{Context, Result};
 use forge::revm::primitives::Eof;
 use foundry_cli::opts::{BuildOpts, CompilerOpts};
-use foundry_common::{compile::ProjectCompiler, fmt::pretty_eof, shell};
-use foundry_compilers::{
-    artifacts::{
-        output_selection::{
-            BytecodeOutputSelection, ContractOutputSelection, DeployedBytecodeOutputSelection,
-            EvmOutputSelection, EwasmOutputSelection,
-        },
-        CompactBytecode, StorageLayout,
+use foundry_common::{
+    compile::{PathOrContractInfo, ProjectCompiler},
+    find_matching_contract_artifact, find_target_path,
+    fmt::pretty_eof,
+    shell,
+};
+use foundry_compilers::artifacts::{
+    output_selection::{
+        BytecodeOutputSelection, ContractOutputSelection, DeployedBytecodeOutputSelection,
+        EvmOutputSelection, EwasmOutputSelection,
     },
-    info::ContractInfo,
-    utils::canonicalize,
+    CompactBytecode, StorageLayout,
 };
 use regex::Regex;
 use serde_json::{Map, Value};
-use std::{collections::BTreeMap, fmt, sync::LazyLock};
+use std::{collections::BTreeMap, fmt, str::FromStr, sync::LazyLock};
 
 /// CLI arguments for `forge inspect`.
 #[derive(Clone, Debug, Parser)]
 pub struct InspectArgs {
     /// The identifier of the contract to inspect in the form `(<path>:)?<contractname>`.
-    pub contract: ContractInfo,
+    #[arg(value_parser = PathOrContractInfo::from_str)]
+    pub contract: PathOrContractInfo,
 
     /// The contract artifact field to inspect.
     #[arg(value_enum)]
@@ -64,17 +66,11 @@ impl InspectArgs {
         // Build the project
         let project = modified_build_args.project()?;
         let compiler = ProjectCompiler::new().quiet(true);
-        let target_path = if let Some(path) = &contract.path {
-            canonicalize(project.root().join(path))?
-        } else {
-            project.find_contract_path(&contract.name)?
-        };
+        let target_path = find_target_path(&project, &contract)?;
         let mut output = compiler.files([target_path.clone()]).compile(&project)?;
 
         // Find the artifact
-        let artifact = output.remove(&target_path, &contract.name).ok_or_else(|| {
-            eyre::eyre!("Could not find artifact `{contract}` in the compiled artifacts")
-        })?;
+        let artifact = find_matching_contract_artifact(&mut output, &target_path, contract.name())?;
 
         // Match on ContractArtifactFields and pretty-print
         match field {