refactor: extract shared get_artifact_source

Author: zacharyftw

crates/cheatcodes/src/fs.rs

@@ -5,13 +5,13 @@ use crate::{
     Cheatcode, Cheatcodes, CheatcodesExecutor, CheatsCtxt, Result, Vm::*, inspector::exec_create,
 };
 use alloy_dyn_abi::DynSolType;
-use alloy_json_abi::{ContractObject, JsonAbi};
+use alloy_json_abi::ContractObject;
 use alloy_network::{Network, ReceiptResponse};
 use alloy_primitives::{Bytes, FixedBytes, U256, hex, map::Entry};
 use alloy_sol_types::SolValue;
 use dialoguer::{Input, Password};
 use forge_script_sequence::{BroadcastReader, TransactionWithMetadata};
-use foundry_common::fs;
+use foundry_common::{contracts::ContractData, fs};
 use foundry_config::fs_permissions::FsAccessKind;
 use foundry_evm_core::evm::FoundryEvmNetwork;
 use revm::{
@@ -313,9 +313,20 @@ impl Cheatcode for getDeployedCodeCall {
 impl Cheatcode for getSelectorsCall {
     fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
         let Self { artifactPath: path } = self;
-        let abi = get_artifact_abi(state, path)?;
-        let selectors: Vec<FixedBytes<4>> =
-            abi.functions().map(|func| func.selector().into()).collect();
+        let selectors: Vec<FixedBytes<4>> = match get_artifact_source(state, path)? {
+            ArtifactSource::InMemory(data) => data.abi.functions().map(|f| f.selector()).collect(),
+            ArtifactSource::Disk(path) => {
+                let data = read_artifact_file(state, &path)?;
+                // Parse as raw JSON rather than `ContractObject` so we can still read selectors
+                // from artifacts with unlinked bytecode (which `ContractObject` rejects).
+                let json: serde_json::Value = serde_json::from_str(&data)?;
+                let abi =
+                    json.get("abi").ok_or_else(|| fmt_err!("no `abi` field in artifact JSON"))?;
+                let abi: alloy_json_abi::JsonAbi =
+                    serde_json::from_value(abi.clone()).map_err(|e| fmt_err!("{e}"))?;
+                abi.functions().map(|f| f.selector()).collect()
+            }
+        };
         Ok(selectors.abi_encode())
     }
 }
@@ -457,175 +468,31 @@ fn deploy_code<FEN: FoundryEvmNetwork>(
     Ok(address.abi_encode())
 }
 
-/// Returns the bytecode from a JSON artifact file.
+/// Resolved location of an artifact referenced by a cheatcode path argument.
+enum ArtifactSource<'a> {
+    /// The artifact was matched in the in-memory `available_artifacts` list.
+    InMemory(&'a ContractData),
+    /// The artifact must be read from the given path on disk.
+    Disk(PathBuf),
+}
+
+/// Resolves a cheatcode artifact reference to its source.
 ///
-/// Can parse following input formats:
+/// Can parse the following input formats:
 /// - `path/to/artifact.json`
 /// - `path/to/contract.sol`
 /// - `path/to/contract.sol:ContractName`
 /// - `path/to/contract.sol:ContractName:0.8.23`
 /// - `path/to/contract.sol:0.8.23`
 /// - `ContractName`
 /// - `ContractName:0.8.23`
-///
-/// This function is safe to use with contracts that have library dependencies.
-/// `alloy_json_abi::ContractObject` validates bytecode during JSON parsing and will
-/// reject artifacts with unlinked library placeholders.
-fn get_artifact_code<FEN: FoundryEvmNetwork>(
-    state: &Cheatcodes<FEN>,
-    path: &str,
-    deployed: bool,
-) -> Result<Bytes> {
-    let path = if path.ends_with(".json") {
-        PathBuf::from(path)
-    } else {
-        let mut parts = path.split(':');
-
-        let mut file = None;
-        let mut contract_name = None;
-        let mut version = None;
-
-        let path_or_name = parts.next().unwrap();
-        if path_or_name.contains('.') {
-            file = Some(PathBuf::from(path_or_name));
-            if let Some(name_or_version) = parts.next() {
-                if name_or_version.contains('.') {
-                    version = Some(name_or_version);
-                } else {
-                    contract_name = Some(name_or_version);
-                    version = parts.next();
-                }
-            }
-        } else {
-            contract_name = Some(path_or_name);
-            version = parts.next();
-        }
-
-        let version = if let Some(version) = version {
-            Some(Version::parse(version).map_err(|e| fmt_err!("failed parsing version: {e}"))?)
-        } else {
-            None
-        };
-
-        // Use available artifacts list if present
-        if let Some(artifacts) = &state.config.available_artifacts {
-            let filtered = artifacts
-                .iter()
-                .filter(|(id, _)| {
-                    // name might be in the form of "Counter.0.8.23"
-                    let id_name = id.name.split('.').next().unwrap();
-
-                    if let Some(path) = &file
-                        && !id.source.ends_with(path)
-                    {
-                        return false;
-                    }
-                    if let Some(name) = contract_name
-                        && id_name != name
-                    {
-                        return false;
-                    }
-                    if let Some(ref version) = version
-                        && (id.version.minor != version.minor
-                            || id.version.major != version.major
-                            || id.version.patch != version.patch)
-                    {
-                        return false;
-                    }
-                    true
-                })
-                .collect::<Vec<_>>();
-
-            let artifact = match &filtered[..] {
-                [] => None,
-                [artifact] => Some(Ok(*artifact)),
-                filtered => {
-                    let mut filtered = filtered.to_vec();
-                    // If we know the current script/test contract solc version, try to filter by it
-                    Some(
-                        state
-                            .config
-                            .running_artifact
-                            .as_ref()
-                            .and_then(|running| {
-                                // Firstly filter by version
-                                filtered.retain(|(id, _)| id.version == running.version);
-
-                                // Return artifact if only one matched
-                                if filtered.len() == 1 {
-                                    return Some(filtered[0]);
-                                }
-
-                                // Try filtering by profile as well
-                                filtered.retain(|(id, _)| id.profile == running.profile);
-
-                                (filtered.len() == 1).then(|| filtered[0])
-                            })
-                            .ok_or_else(|| fmt_err!("multiple matching artifacts found")),
-                    )
-                }
-            };
-
-            if let Some(artifact) = artifact {
-                let artifact = artifact?;
-                let maybe_bytecode = if deployed {
-                    artifact.1.deployed_bytecode().cloned()
-                } else {
-                    artifact.1.bytecode().cloned()
-                };
-
-                return maybe_bytecode.ok_or_else(|| {
-                    fmt_err!("no bytecode for contract; is it abstract or unlinked?")
-                });
-            }
-        }
-
-        // Fallback: construct path manually when no artifacts list or no match found
-        let path_in_artifacts = match (file.map(|f| f.to_string_lossy().to_string()), contract_name)
-        {
-            (Some(file), Some(contract_name)) => {
-                PathBuf::from(format!("{file}/{contract_name}.json"))
-            }
-            (None, Some(contract_name)) => {
-                PathBuf::from(format!("{contract_name}.sol/{contract_name}.json"))
-            }
-            (Some(file), None) => {
-                let name = file.replace(".sol", "");
-                PathBuf::from(format!("{file}/{name}.json"))
-            }
-            _ => bail!("invalid artifact path"),
-        };
-
-        state.config.paths.artifacts.join(path_in_artifacts)
-    };
-
-    let path = state.config.ensure_path_allowed(path, FsAccessKind::Read)?;
-    let data = fs::read_to_string(path).map_err(|e| {
-        if state.config.available_artifacts.is_some() {
-            fmt_err!("no matching artifact found")
-        } else {
-            e.into()
-        }
-    })?;
-    let artifact = serde_json::from_str::<ContractObject>(&data)?;
-    let maybe_bytecode = if deployed { artifact.deployed_bytecode } else { artifact.bytecode };
-    maybe_bytecode.ok_or_else(|| fmt_err!("no bytecode for contract; is it abstract or unlinked?"))
-}
-
-/// Returns the ABI of a matching artifact from the given path.
-///
-/// See [`get_artifact_code`] for the supported path formats.
-fn get_artifact_abi<FEN: FoundryEvmNetwork>(
-    state: &Cheatcodes<FEN>,
+fn get_artifact_source<'a, FEN: FoundryEvmNetwork>(
+    state: &'a Cheatcodes<FEN>,
     path: &str,
-) -> Result<JsonAbi> {
+) -> Result<ArtifactSource<'a>> {
     if path.ends_with(".json") {
-        // Read JSON artifact directly from disk.
         let path = state.config.ensure_path_allowed(path, FsAccessKind::Read)?;
-        let data = fs::read_to_string(path)?;
-        let json: serde_json::Value = serde_json::from_str(&data)?;
-        let abi = json.get("abi").ok_or_else(|| fmt_err!("no `abi` field in artifact JSON"))?;
-        return serde_json::from_value(abi.clone()).map_err(|e| fmt_err!("{e}"));
+        return Ok(ArtifactSource::Disk(path));
     }
 
     let mut parts = path.split(':');
@@ -661,6 +528,7 @@ fn get_artifact_abi<FEN: FoundryEvmNetwork>(
         let filtered = artifacts
             .iter()
             .filter(|(id, _)| {
+                // name might be in the form of "Counter.0.8.23"
                 let id_name = id.name.split('.').next().unwrap();
 
                 if let Some(path) = &file
@@ -689,17 +557,24 @@ fn get_artifact_abi<FEN: FoundryEvmNetwork>(
             [artifact] => Some(Ok(*artifact)),
             filtered => {
                 let mut filtered = filtered.to_vec();
+                // If we know the current script/test contract solc version, try to filter by it.
                 Some(
                     state
                         .config
                         .running_artifact
                         .as_ref()
                         .and_then(|running| {
+                            // Firstly filter by version.
                             filtered.retain(|(id, _)| id.version == running.version);
+
+                            // Return artifact if only one matched.
                             if filtered.len() == 1 {
                                 return Some(filtered[0]);
                             }
+
+                            // Try filtering by profile as well.
                             filtered.retain(|(id, _)| id.profile == running.profile);
+
                             (filtered.len() == 1).then(|| filtered[0])
                         })
                         .ok_or_else(|| fmt_err!("multiple matching artifacts found")),
@@ -708,16 +583,13 @@ fn get_artifact_abi<FEN: FoundryEvmNetwork>(
         };
 
         if let Some(artifact) = artifact {
-            let artifact = artifact?;
-            return Ok(artifact.1.abi.clone());
+            return Ok(ArtifactSource::InMemory(artifact?.1));
         }
     }
 
-    // Fallback: construct path manually and read JSON from disk.
+    // Fallback: construct path manually when no artifacts list or no match found.
     let path_in_artifacts = match (file.map(|f| f.to_string_lossy().to_string()), contract_name) {
-        (Some(file), Some(contract_name)) => {
-            PathBuf::from(format!("{file}/{contract_name}.json"))
-        }
+        (Some(file), Some(contract_name)) => PathBuf::from(format!("{file}/{contract_name}.json")),
         (None, Some(contract_name)) => {
             PathBuf::from(format!("{contract_name}.sol/{contract_name}.json"))
         }
@@ -730,16 +602,47 @@ fn get_artifact_abi<FEN: FoundryEvmNetwork>(
 
     let path = state.config.paths.artifacts.join(path_in_artifacts);
     let path = state.config.ensure_path_allowed(path, FsAccessKind::Read)?;
-    let data = fs::read_to_string(path).map_err(|e| {
+    Ok(ArtifactSource::Disk(path))
+}
+
+/// Reads an artifact JSON file, mapping I/O errors to a helpful message when the
+/// lookup fell through the in-memory artifacts list.
+fn read_artifact_file<FEN: FoundryEvmNetwork>(
+    state: &Cheatcodes<FEN>,
+    path: &Path,
+) -> Result<String> {
+    fs::read_to_string(path).map_err(|e| {
         if state.config.available_artifacts.is_some() {
             fmt_err!("no matching artifact found")
         } else {
             e.into()
         }
-    })?;
-    let json: serde_json::Value = serde_json::from_str(&data)?;
-    let abi = json.get("abi").ok_or_else(|| fmt_err!("no `abi` field in artifact JSON"))?;
-    serde_json::from_value(abi.clone()).map_err(|e| fmt_err!("{e}"))
+    })
+}
+
+/// Returns the bytecode from a JSON artifact file.
+///
+/// See [`get_artifact_source`] for the supported path formats.
+///
+/// This function is safe to use with contracts that have library dependencies.
+/// `alloy_json_abi::ContractObject` validates bytecode during JSON parsing and will
+/// reject artifacts with unlinked library placeholders.
+fn get_artifact_code<FEN: FoundryEvmNetwork>(
+    state: &Cheatcodes<FEN>,
+    path: &str,
+    deployed: bool,
+) -> Result<Bytes> {
+    let maybe_bytecode = match get_artifact_source(state, path)? {
+        ArtifactSource::InMemory(data) => {
+            if deployed { data.deployed_bytecode() } else { data.bytecode() }.cloned()
+        }
+        ArtifactSource::Disk(path) => {
+            let data = read_artifact_file(state, &path)?;
+            let artifact = serde_json::from_str::<ContractObject>(&data)?;
+            if deployed { artifact.deployed_bytecode } else { artifact.bytecode }
+        }
+    };
+    maybe_bytecode.ok_or_else(|| fmt_err!("no bytecode for contract; is it abstract or unlinked?"))
 }
 
 impl Cheatcode for ffiCall {