fix(forge): consolidate flatten

Author: grandizzy

crates/cast/src/lib.rs

@@ -30,10 +30,10 @@ use foundry_common::{
     TransactionReceiptWithRevertReason,
     abi::{coerce_value, encode_function_args, get_event, get_func},
     compile::etherscan_project,
+    flatten,
     fmt::*,
     fs, get_pretty_tx_receipt_attr, shell,
 };
-use foundry_compilers::flatten::Flattener;
 use foundry_config::Chain;
 use foundry_evm_core::ic::decode_instructions;
 use futures::{FutureExt, StreamExt, future::Either};
@@ -2192,7 +2192,7 @@ impl SimpleCast {
         let project = etherscan_project(metadata, tmp.path())?;
         let target_path = project.find_contract_path(&metadata.contract_name)?;
 
-        let flattened = Flattener::new(project, &target_path)?.flatten();
+        let flattened = flatten(project, &target_path)?;
 
         if let Some(path) = output_path {
             fs::create_dir_all(path.parent().unwrap())?;

crates/common/src/utils.rs

@@ -1,9 +1,14 @@
 //! Uncategorised utilities.
 
 use alloy_primitives::{B256, Bytes, U256, hex, keccak256};
-use foundry_compilers::artifacts::BytecodeObject;
+use foundry_compilers::{
+    Project,
+    artifacts::{BytecodeObject, SolcLanguage},
+    error::SolcError,
+    flatten::{Flattener, FlattenerError},
+};
 use regex::Regex;
-use std::sync::LazyLock;
+use std::{path::Path, sync::LazyLock};
 
 static BYTECODE_PLACEHOLDER_RE: LazyLock<Regex> =
     LazyLock::new(|| Regex::new(r"__\$.{34}\$__").expect("invalid regex"));
@@ -84,3 +89,19 @@ pub fn strip_bytecode_placeholders(bytecode: &BytecodeObject) -> Option<Bytes> {
         }
     }
 }
+
+/// Flattens the given target of the project. Falls back to the old flattening implementation
+/// if the target cannot be compiled successfully. This would be the case if the target has invalid
+/// syntax. (e.g. Solang)
+pub fn flatten(project: Project, target_path: &Path) -> eyre::Result<String> {
+    let flattened = match Flattener::new(project.clone(), target_path) {
+        Ok(flattener) => Ok(flattener.flatten()),
+        Err(FlattenerError::Compilation(_)) => {
+            project.paths.with_language::<SolcLanguage>().flatten(target_path)
+        }
+        Err(FlattenerError::Other(err)) => Err(err),
+    }
+    .map_err(|err: SolcError| eyre::eyre!("Failed to flatten: {err}"))?;
+
+    Ok(flattened)
+}

crates/forge/src/cmd/flatten.rs

@@ -4,12 +4,7 @@ use foundry_cli::{
     opts::{BuildOpts, ProjectPathOpts},
     utils::LoadConfig,
 };
-use foundry_common::{compile::with_compilation_reporter, fs};
-use foundry_compilers::{
-    compilers::solc::SolcLanguage,
-    error::SolcError,
-    flatten::{Flattener, FlattenerError},
-};
+use foundry_common::{flatten, fs};
 use std::path::PathBuf;
 
 /// CLI arguments for `forge flatten`.
@@ -44,22 +39,7 @@ impl FlattenArgs {
         let project = config.ephemeral_project()?;
 
         let target_path = dunce::canonicalize(target_path)?;
-
-        let flattener = with_compilation_reporter(true, Some(project.root().to_path_buf()), || {
-            Flattener::new(project.clone(), &target_path)
-        });
-
-        let flattened = match flattener {
-            Ok(flattener) => Ok(flattener.flatten()),
-            Err(FlattenerError::Compilation(_)) => {
-                // Fallback to the old flattening implementation if we couldn't compile the target
-                // successfully. This would be the case if the target has invalid
-                // syntax. (e.g. Solang)
-                project.paths.with_language::<SolcLanguage>().flatten(&target_path)
-            }
-            Err(FlattenerError::Other(err)) => Err(err),
-        }
-        .map_err(|err: SolcError| eyre::eyre!("Failed to flatten: {err}"))?;
+        let flattened = flatten(project, &target_path)?;
 
         match output {
             Some(output) => {

crates/verify/src/etherscan/flatten.rs

@@ -1,7 +1,8 @@
 use super::{EtherscanSourceProvider, VerifyArgs};
 use crate::provider::VerificationContext;
-use eyre::{Context, Result};
+use eyre::Result;
 use foundry_block_explorers::verify::CodeFormat;
+use foundry_common::flatten;
 use foundry_compilers::{
     AggregatedCompilerOutput,
     artifacts::{BytecodeHash, Source, Sources},
@@ -32,18 +33,15 @@ impl EtherscanSourceProvider for EtherscanFlattenedSource {
             bch,
         );
 
-        let source = context
-            .project
-            .paths
-            .clone()
-            .with_language::<SolcLanguage>()
-            .flatten(&context.target_path)
-            .wrap_err("Failed to flatten contract")?;
-
+        let flattened_source = flatten(context.project.clone(), &context.target_path)?;
         if !args.force {
             // solc dry run of flattened code
-            self.check_flattened(source.clone(), &context.compiler_version, &context.target_path)
-                .map_err(|err| {
+            self.check_flattened(
+                flattened_source.clone(),
+                &context.compiler_version,
+                &context.target_path,
+            )
+            .map_err(|err| {
                 eyre::eyre!(
                     "Failed to compile the flattened code locally: `{}`\
             To skip this solc dry, have a look at the `--force` flag of this command.",
@@ -52,7 +50,7 @@ impl EtherscanSourceProvider for EtherscanFlattenedSource {
             })?;
         }
 
-        Ok((source, context.target_name.clone(), CodeFormat::SingleFile))
+        Ok((flattened_source, context.target_name.clone(), CodeFormat::SingleFile))
     }
 }