chore: avoid cloning CompileOutput to parse inline config

DaniPopes · DaniPopes · commit 11a36f99d954 · 2023-08-20T23:44:11.000+02:00
diff --git a/crates/config/src/inline/natspec.rs b/crates/config/src/inline/natspec.rs
@@ -1,8 +1,11 @@
-use std::{collections::BTreeMap, path::Path};
+use std::{
+    collections::BTreeMap,
+    path::{Path, PathBuf},
+};
 
 use ethers_solc::{
     artifacts::{ast::NodeType, Node},
-    ProjectCompileOutput,
+    ArtifactId, ArtifactOutput, Artifacts, ConfigurableArtifacts, ProjectCompileOutput,
 };
 use serde_json::Value;
 
@@ -26,21 +29,21 @@ impl NatSpec {
     /// Factory function that extracts a vector of [`NatSpec`] instances from
     /// a solc compiler output. The root path is to express contract base dirs.
     /// That is essential to match per-test configs at runtime.
-    pub fn parse<P>(output: &ProjectCompileOutput, root: &P) -> Vec<Self>
-    where
-        P: AsRef<Path>,
-    {
+    pub fn parse(output: &ProjectCompileOutput, root: &Path) -> Vec<Self> {
         let mut natspecs: Vec<Self> = vec![];
 
-        let output = output.clone();
-        for artifact in output.with_stripped_file_prefixes(root).into_artifacts() {
-            if let Some(ast) = artifact.1.ast.as_ref() {
-                let contract: String = artifact.0.identifier();
-                if let Some(node) = contract_root_node(&ast.nodes, &contract) {
-                    apply(&mut natspecs, &contract, node)
-                }
-            }
+        let artifacts = artifacts::<_, ConfigurableArtifacts>(output.cached_artifacts())
+            .chain(artifacts::<_, ConfigurableArtifacts>(output.compiled_artifacts()));
+        for (id, artifact) in artifacts {
+            let Some(ast) = &artifact.ast else { continue };
+            let path = id.source.as_path();
+            let path = path.strip_prefix(root).unwrap_or(path);
+            // id.identifier
+            let contract = format!("{}:{}", path.display(), id.name);
+            let Some(node) = contract_root_node(&ast.nodes, &contract) else { continue };
+            apply(&mut natspecs, &contract, node)
         }
+
         natspecs
     }
 
@@ -52,24 +55,21 @@ impl NatSpec {
     }
 
     /// Returns a list of configuration lines that match the current profile
-    pub fn current_profile_configs(&self) -> Vec<String> {
-        let prefix: &str = INLINE_CONFIG_PREFIX_SELECTED_PROFILE.as_ref();
-        self.config_lines_with_prefix(prefix)
+    pub fn current_profile_configs(&self) -> impl Iterator<Item = String> + '_ {
+        self.config_lines_with_prefix(INLINE_CONFIG_PREFIX_SELECTED_PROFILE.as_str())
     }
 
     /// Returns a list of configuration lines that match a specific string prefix
-    pub fn config_lines_with_prefix<S: Into<String>>(&self, prefix: S) -> Vec<String> {
-        let prefix: String = prefix.into();
-        self.config_lines().into_iter().filter(|l| l.starts_with(&prefix)).collect()
+    pub fn config_lines_with_prefix<'a>(
+        &'a self,
+        prefix: &'a str,
+    ) -> impl Iterator<Item = String> + 'a {
+        self.config_lines().into_iter().filter(move |l| l.starts_with(prefix))
     }
 
     /// Returns a list of all the configuration lines available in the natspec
-    pub fn config_lines(&self) -> Vec<String> {
-        self.docs
-            .split('\n')
-            .map(remove_whitespaces)
-            .filter(|line| line.contains(INLINE_CONFIG_PREFIX))
-            .collect::<Vec<String>>()
+    pub fn config_lines(&self) -> impl Iterator<Item = String> + '_ {
+        self.docs.lines().map(remove_whitespaces).filter(|line| line.contains(INLINE_CONFIG_PREFIX))
     }
 }
 
@@ -89,6 +89,31 @@ fn contract_root_node<'a>(nodes: &'a [Node], contract_id: &'a str) -> Option<&'a
     None
 }
 
+// borrowed version of `Artifacts::into_artifacts`
+fn artifacts<'a, T, O: ArtifactOutput<Artifact = T>>(
+    a: &'a Artifacts<T>,
+) -> impl Iterator<Item = (ArtifactId, &'a T)> + 'a {
+    a.0.iter().flat_map(|(file, contract_artifacts)| {
+        contract_artifacts.iter().flat_map(move |(_contract_name, artifacts)| {
+            let source = PathBuf::from(file.clone());
+            artifacts.iter().filter_map(move |artifact| {
+                O::contract_name(&artifact.file).map(|name| {
+                    (
+                        ArtifactId {
+                            path: PathBuf::from(&artifact.file),
+                            name,
+                            source: source.clone(),
+                            version: artifact.version.clone(),
+                        }
+                        .with_slashed_paths(),
+                        &artifact.artifact,
+                    )
+                })
+            })
+        })
+    })
+}
+
 /// Implements a DFS over a compiler output node and its children.
 /// If a natspec is found it is added to `natspecs`
 fn apply(natspecs: &mut Vec<NatSpec>, contract: &str, node: &Node) {
@@ -137,7 +162,7 @@ fn get_fn_docs(fn_data: &BTreeMap<String, Value>) -> Option<(String, String)> {
                 let mut src_line = fn_docs
                     .get("src")
                     .map(|src| src.to_string())
-                    .unwrap_or(String::from("<no-src-line-available>"));
+                    .unwrap_or_else(|| String::from("<no-src-line-available>"));
 
                 src_line.retain(|c| c != '"');
                 return Some((comment.into(), src_line))
@@ -158,7 +183,7 @@ mod tests {
         let natspec = natspec();
         let config_lines = natspec.config_lines();
         assert_eq!(
-            config_lines,
+            config_lines.collect::<Vec<_>>(),
             vec![
                 "forge-config:default.fuzz.runs=600".to_string(),
                 "forge-config:ci.fuzz.runs=500".to_string(),
@@ -173,7 +198,7 @@ mod tests {
         let config_lines = natspec.current_profile_configs();
 
         assert_eq!(
-            config_lines,
+            config_lines.collect::<Vec<_>>(),
             vec![
                 "forge-config:default.fuzz.runs=600".to_string(),
                 "forge-config:default.invariant.runs=1".to_string()
@@ -186,9 +211,9 @@ mod tests {
         use super::INLINE_CONFIG_PREFIX;
         let natspec = natspec();
         let prefix = format!("{INLINE_CONFIG_PREFIX}:default");
-        let config_lines = natspec.config_lines_with_prefix(prefix);
+        let config_lines = natspec.config_lines_with_prefix(&prefix);
         assert_eq!(
-            config_lines,
+            config_lines.collect::<Vec<_>>(),
             vec![
                 "forge-config:default.fuzz.runs=600".to_string(),
                 "forge-config:default.invariant.runs=1".to_string()
diff --git a/crates/forge/bin/cmd/test/mod.rs b/crates/forge/bin/cmd/test/mod.rs
@@ -165,9 +165,8 @@ impl TestArgs {
         let test_options: TestOptions = TestOptionsBuilder::default()
             .fuzz(config.fuzz)
             .invariant(config.invariant)
-            .compile_output(&output)
             .profiles(profiles)
-            .build(project_root)?;
+            .build(&output, project_root)?;
 
         // Determine print verbosity and executor verbosity
         let verbosity = evm_opts.verbosity;
diff --git a/crates/forge/src/lib.rs b/crates/forge/src/lib.rs
@@ -47,6 +47,47 @@ pub struct TestOptions {
 }
 
 impl TestOptions {
+    /// Tries to create a new instance by detecting inline configurations from the project compile
+    /// output.
+    pub fn new(
+        output: &ProjectCompileOutput,
+        root: &Path,
+        profiles: Vec<String>,
+        base_fuzz: FuzzConfig,
+        base_invariant: InvariantConfig,
+    ) -> Result<Self, InlineConfigError> {
+        let natspecs: Vec<NatSpec> = NatSpec::parse(output, root);
+        let mut inline_invariant = InlineConfig::<InvariantConfig>::default();
+        let mut inline_fuzz = InlineConfig::<FuzzConfig>::default();
+
+        for natspec in natspecs {
+            // Perform general validation
+            validate_profiles(&natspec, &profiles)?;
+            FuzzConfig::validate_configs(&natspec)?;
+            InvariantConfig::validate_configs(&natspec)?;
+
+            // Apply in-line configurations for the current profile
+            let configs: Vec<String> = natspec.current_profile_configs().collect();
+            let c: &str = &natspec.contract;
+            let f: &str = &natspec.function;
+            let line: String = natspec.debug_context();
+
+            match base_fuzz.try_merge(&configs) {
+                Ok(Some(conf)) => inline_fuzz.insert(c, f, conf),
+                Ok(None) => { /* No inline config found, do nothing */ }
+                Err(e) => Err(InlineConfigError { line: line.clone(), source: e })?,
+            }
+
+            match base_invariant.try_merge(&configs) {
+                Ok(Some(conf)) => inline_invariant.insert(c, f, conf),
+                Ok(None) => { /* No inline config found, do nothing */ }
+                Err(e) => Err(InlineConfigError { line: line.clone(), source: e })?,
+            }
+        }
+
+        Ok(Self { fuzz: base_fuzz, invariant: base_invariant, inline_fuzz, inline_invariant })
+    }
+
     /// Returns a "fuzz" test runner instance. Parameters are used to select tight scoped fuzz
     /// configs that apply for a contract-function pair. A fallback configuration is applied
     /// if no specific setup is found for a given input.
@@ -127,83 +168,23 @@ impl TestOptions {
     }
 }
 
-impl<'a, P> TryFrom<(&'a ProjectCompileOutput, &'a P, Vec<String>, FuzzConfig, InvariantConfig)>
-    for TestOptions
-where
-    P: AsRef<Path>,
-{
-    type Error = InlineConfigError;
-
-    /// Tries to create an instance of `Self`, detecting inline configurations from the project
-    /// compile output.
-    ///
-    /// Param is a tuple, whose elements are:
-    /// 1. Solidity compiler output, essential to extract natspec test configs.
-    /// 2. Root path to express contract base dirs. This is essential to match inline configs at
-    /// runtime. 3. List of available configuration profiles
-    /// 4. Reference to a fuzz base configuration.
-    /// 5. Reference to an invariant base configuration.
-    fn try_from(
-        value: (&'a ProjectCompileOutput, &'a P, Vec<String>, FuzzConfig, InvariantConfig),
-    ) -> Result<Self, Self::Error> {
-        let output = value.0;
-        let root = value.1;
-        let profiles = &value.2;
-        let base_fuzz: FuzzConfig = value.3;
-        let base_invariant: InvariantConfig = value.4;
-
-        let natspecs: Vec<NatSpec> = NatSpec::parse(output, root);
-        let mut inline_invariant = InlineConfig::<InvariantConfig>::default();
-        let mut inline_fuzz = InlineConfig::<FuzzConfig>::default();
-
-        for natspec in natspecs {
-            // Perform general validation
-            validate_profiles(&natspec, profiles)?;
-            FuzzConfig::validate_configs(&natspec)?;
-            InvariantConfig::validate_configs(&natspec)?;
-
-            // Apply in-line configurations for the current profile
-            let configs: Vec<String> = natspec.current_profile_configs();
-            let c: &str = &natspec.contract;
-            let f: &str = &natspec.function;
-            let line: String = natspec.debug_context();
-
-            match base_fuzz.try_merge(&configs) {
-                Ok(Some(conf)) => inline_fuzz.insert(c, f, conf),
-                Err(e) => Err(InlineConfigError { line: line.clone(), source: e })?,
-                _ => { /* No inline config found, do nothing */ }
-            }
-
-            match base_invariant.try_merge(&configs) {
-                Ok(Some(conf)) => inline_invariant.insert(c, f, conf),
-                Err(e) => Err(InlineConfigError { line: line.clone(), source: e })?,
-                _ => { /* No inline config found, do nothing */ }
-            }
-        }
-
-        Ok(Self { fuzz: base_fuzz, invariant: base_invariant, inline_fuzz, inline_invariant })
-    }
-}
-
 /// Builder utility to create a [`TestOptions`] instance.
 #[derive(Default)]
+#[must_use = "builders do nothing unless you call `build` on them"]
 pub struct TestOptionsBuilder {
     fuzz: Option<FuzzConfig>,
     invariant: Option<InvariantConfig>,
     profiles: Option<Vec<String>>,
-    output: Option<ProjectCompileOutput>,
 }
 
 impl TestOptionsBuilder {
     /// Sets a [`FuzzConfig`] to be used as base "fuzz" configuration.
-    #[must_use = "A base 'fuzz' config must be provided"]
     pub fn fuzz(mut self, conf: FuzzConfig) -> Self {
         self.fuzz = Some(conf);
         self
     }
 
     /// Sets a [`InvariantConfig`] to be used as base "invariant" configuration.
-    #[must_use = "A base 'invariant' config must be provided"]
     pub fn invariant(mut self, conf: InvariantConfig) -> Self {
         self.invariant = Some(conf);
         self
@@ -216,40 +197,21 @@ impl TestOptionsBuilder {
         self
     }
 
-    /// Sets a project compiler output instance. This is used to extract
-    /// inline test configurations that override `self.fuzz` and `self.invariant`
-    /// specs when necessary.
-    pub fn compile_output(mut self, output: &ProjectCompileOutput) -> Self {
-        self.output = Some(output.clone());
-        self
-    }
-
     /// Creates an instance of [`TestOptions`]. This takes care of creating "fuzz" and
     /// "invariant" fallbacks, and extracting all inline test configs, if available.
     ///
     /// `root` is a reference to the user's project root dir. This is essential
     /// to determine the base path of generated contract identifiers. This is to provide correct
     /// matchers for inline test configs.
-    pub fn build(self, root: impl AsRef<Path>) -> Result<TestOptions, InlineConfigError> {
-        let default_profiles = vec![Config::selected_profile().into()];
-        let profiles: Vec<String> = self.profiles.unwrap_or(default_profiles);
+    pub fn build(
+        self,
+        output: &ProjectCompileOutput,
+        root: &Path,
+    ) -> Result<TestOptions, InlineConfigError> {
+        let profiles: Vec<String> =
+            self.profiles.unwrap_or_else(|| vec![Config::selected_profile().into()]);
         let base_fuzz = self.fuzz.unwrap_or_default();
         let base_invariant = self.invariant.unwrap_or_default();
-
-        match self.output {
-            Some(compile_output) => Ok(TestOptions::try_from((
-                &compile_output,
-                &root,
-                profiles,
-                base_fuzz,
-                base_invariant,
-            ))?),
-            None => Ok(TestOptions {
-                fuzz: base_fuzz,
-                invariant: base_invariant,
-                inline_fuzz: InlineConfig::default(),
-                inline_invariant: InlineConfig::default(),
-            }),
-        }
+        TestOptions::new(output, &root, profiles, base_fuzz, base_invariant)
     }
 }
diff --git a/crates/forge/tests/it/inline.rs b/crates/forge/tests/it/inline.rs
@@ -76,9 +76,8 @@ fn build_test_options() {
     let build_result = TestOptionsBuilder::default()
         .fuzz(FuzzConfig::default())
         .invariant(InvariantConfig::default())
-        .compile_output(&COMPILED)
         .profiles(profiles)
-        .build(root);
+        .build(&COMPILED, root);
 
     assert!(build_result.is_ok());
 }
@@ -90,9 +89,8 @@ fn build_test_options_just_one_valid_profile() {
     let build_result = TestOptionsBuilder::default()
         .fuzz(FuzzConfig::default())
         .invariant(InvariantConfig::default())
-        .compile_output(&COMPILED)
         .profiles(valid_profiles)
-        .build(root);
+        .build(&COMPILED, root);
 
     // We expect an error, since COMPILED contains in-line
     // per-test configs for "default" and "ci" profiles
@@ -104,7 +102,6 @@ fn test_options() -> TestOptions {
     TestOptionsBuilder::default()
         .fuzz(FuzzConfig::default())
         .invariant(InvariantConfig::default())
-        .compile_output(&COMPILED)
-        .build(root)
+        .build(&COMPILED, root)
         .expect("Config loaded")
 }
