refactor(flatten): move compilation logic into Flattener

crates/artifacts/solc/src/output_selection.rs

@@ -130,6 +130,20 @@ impl OutputSelection {
         Default::default()
     }
 
+    /// Returns output selection which requests only AST for all sources.
+    pub fn ast_output_selection() -> Self {
+        BTreeMap::from([(
+            "*".to_string(),
+            BTreeMap::from([
+                // Do not request any output for separate contracts
+                ("*".to_string(), vec![]),
+                // Request AST for all sources.
+                ("".to_string(), vec!["ast".to_string()]),
+            ]),
+        )])
+        .into()
+    }
+
     /// Returns true if this output selection is a subset of the other output selection.
     /// TODO: correctly process wildcard keys to reduce false negatives
     pub fn is_subset_of(&self, other: &Self) -> bool {

crates/project/src/flatten.rs

@@ -2,10 +2,11 @@ use crate::{
     compilers::{Compiler, ParsedSource},
     filter::MaybeSolData,
     resolver::parse::SolData,
-    Graph, Project, ProjectCompileOutput, ProjectPathsConfig,
+    CompilerSettings, Graph, Project, ProjectPathsConfig,
 };
 use foundry_compilers_artifacts::{
     ast::{visitor::Visitor, *},
+    output_selection::OutputSelection,
     solc::ExternalInlineAssemblyReference,
     sources::{Source, Sources},
     ContractDefinitionPart, SourceUnit, SourceUnitPart,
@@ -162,6 +163,20 @@ impl<'a> FlatteningResult<'a> {
     }
 }
 
+#[derive(thiserror::Error, Debug)]
+pub enum FlattenerError {
+    #[error("Failed to compile {0}")]
+    Compilation(SolcError),
+    #[error(transparent)]
+    Other(SolcError),
+}
+
+impl<T: Into<SolcError>> From<T> for FlattenerError {
+    fn from(err: T) -> Self {
+        Self::Other(err.into())
+    }
+}
+
 /// Context for flattening. Stores all sources and ASTs that are in scope of the flattening target.
 pub struct Flattener {
     /// Target file to flatten.
@@ -177,25 +192,25 @@ pub struct Flattener {
 }
 
 impl Flattener {
-    /// Compilation output is expected to contain all artifacts for all sources.
-    /// Flattener caller is expected to resolve all imports of target file, compile them and pass
-    /// into this function.
+    /// Compiles the target file and prepares AST and analysis data for flattening.
     pub fn new<C: Compiler>(
-        project: &Project<C>,
-        output: &ProjectCompileOutput<C>,
+        mut project: Project<C>,
         target: &Path,
-    ) -> Result<Self>
+    ) -> std::result::Result<Self, FlattenerError>
     where
         C::ParsedSource: MaybeSolData,
     {
-        let input_files = output
-            .artifacts_with_files()
-            .map(|(file, _, _)| PathBuf::from(file))
-            .collect::<HashSet<_>>()
-            .into_iter()
-            .collect::<Vec<_>>();
-
-        let sources = Source::read_all_files(input_files)?;
+        // Configure project to compile the target file and only request AST for target file.
+        project.cached = false;
+        project.no_artifacts = true;
+        project.settings.update_output_selection(|selection| {
+            *selection = OutputSelection::ast_output_selection();
+        });
+
+        let output =
+            project.compile_file(target).map_err(FlattenerError::Compilation)?.compiler_output;
+
+        let sources = Source::read_all_files(vec![target.to_path_buf()])?;
         let graph = Graph::<C::ParsedSource>::resolve_sources(&project.paths, sources)?;
 
         let ordered_sources = collect_ordered_deps(&target.to_path_buf(), &project.paths, &graph)?;
@@ -214,9 +229,9 @@ impl Flattener {
 
         // Convert all ASTs from artifacts to strongly typed ASTs
         let mut asts: Vec<(PathBuf, SourceUnit)> = Vec::new();
-        for (path, ast) in output.artifacts_with_files().filter_map(|(path, _, artifact)| {
-            if let Some(ast) = artifact.ast.as_ref() {
-                if sources.contains_key(Path::new(path)) {
+        for (path, ast) in output.sources.0.iter().filter_map(|(path, files)| {
+            if let Some(ast) = files.first().and_then(|source| source.source_file.ast.as_ref()) {
+                if sources.contains_key(path) {
                     return Some((path, ast));
                 }
             }

crates/project/tests/project.rs

@@ -544,8 +544,7 @@ fn test_flatteners(project: &TempProject, target: &Path, additional_checks: fn(&
     let target = canonicalize(target).unwrap();
     let result =
         project.project().paths.clone().with_language::<SolcLanguage>().flatten(&target).unwrap();
-    let solc_result =
-        Flattener::new(project.project(), &project.compile().unwrap(), &target).unwrap().flatten();
+    let solc_result = Flattener::new(project.project().clone(), &target).unwrap().flatten();
 
     assert_eq!(result, solc_result);
 
@@ -1148,8 +1147,7 @@ contract Bar is Foo {}
         )
         .unwrap();
 
-    let result =
-        Flattener::new(project.project(), &project.compile().unwrap(), &target).unwrap().flatten();
+    let result = Flattener::new(project.project().clone(), &target).unwrap().flatten();
     assert_eq!(
         result,
         r"pragma solidity ^0.8.10;
@@ -1254,8 +1252,7 @@ contract D is C_File.B_File.A_File.A {
 }
 "#,).unwrap();
 
-    let result =
-        Flattener::new(project.project(), &project.compile().unwrap(), &target).unwrap().flatten();
+    let result = Flattener::new(project.project().clone(), &target).unwrap().flatten();
     assert_eq!(
         result,
         r"pragma solidity ^0.8.10;
@@ -1350,8 +1347,7 @@ contract B {
         )
         .unwrap();
 
-    let result =
-        Flattener::new(project.project(), &project.compile().unwrap(), &target).unwrap().flatten();
+    let result = Flattener::new(project.project().clone(), &target).unwrap().flatten();
     assert_eq!(
         result,
         r#"pragma solidity ^0.8.10;
@@ -1410,8 +1406,7 @@ contract B is A {}
         )
         .unwrap();
 
-    let result =
-        Flattener::new(project.project(), &project.compile().unwrap(), &target).unwrap().flatten();
+    let result = Flattener::new(project.project().clone(), &target).unwrap().flatten();
     assert_eq!(
         result,
         r"pragma solidity =0.6.12;
@@ -1529,8 +1524,7 @@ contract B is A {
         )
         .unwrap();
 
-    let result =
-        Flattener::new(project.project(), &project.compile().unwrap(), &target).unwrap().flatten();
+    let result = Flattener::new(project.project().clone(), &target).unwrap().flatten();
     assert_eq!(
         result,
         r"pragma solidity ^0.8.10;
@@ -1586,8 +1580,7 @@ contract B is Alias {
         )
         .unwrap();
 
-    let result =
-        Flattener::new(project.project(), &project.compile().unwrap(), &target).unwrap().flatten();
+    let result = Flattener::new(project.project().clone(), &target).unwrap().flatten();
     assert_eq!(
         result,
         r"pragma solidity ^0.8.10;
@@ -1664,8 +1657,7 @@ contract Foo {
         )
         .unwrap();
 
-    let result =
-        Flattener::new(project.project(), &project.compile().unwrap(), &target).unwrap().flatten();
+    let result = Flattener::new(project.project().clone(), &target).unwrap().flatten();
     assert_eq!(
         result,
         r"pragma solidity ^0.8.10;
@@ -1746,8 +1738,7 @@ contract Foo {
         )
         .unwrap();
 
-    let result =
-        Flattener::new(project.project(), &project.compile().unwrap(), &target).unwrap().flatten();
+    let result = Flattener::new(project.project().clone(), &target).unwrap().flatten();
     assert_eq!(
         result,
         r"pragma solidity ^0.8.10;
@@ -1814,8 +1805,7 @@ contract Foo {
         )
         .unwrap();
 
-    let result =
-        Flattener::new(project.project(), &project.compile().unwrap(), &target).unwrap().flatten();
+    let result = Flattener::new(project.project().clone(), &target).unwrap().flatten();
     assert_eq!(
         result,
         r"pragma solidity ^0.8.10;