From e2d32789141e62432f3f4afd40f56e34f1011224 Mon Sep 17 00:00:00 2001
From: Arsenii Kulikov <klkvrr@gmail.com>
Date: Fri, 2 Feb 2024 15:24:53 +0300
Subject: [PATCH] feat(`forge`): new `flatten` implementation (#6936)

* Update Flatten impl

* Bump foundry-compilers

* fix error handling

* Update crates/forge/bin/cmd/flatten.rs

Co-authored-by: DaniPopes <57450786+DaniPopes@users.noreply.github.com>

* error handling

* Bump compilers and block-explorers

* Simplify compilation

* fmt

* fix doc

* use patch

* bump compilers

---------

Co-authored-by: DaniPopes <57450786+DaniPopes@users.noreply.github.com>
---
 Cargo.lock                      | 12 ++++--------
 Cargo.toml                      |  6 +++---
 crates/forge/bin/cmd/flatten.rs | 23 +++++++++++++++++++----
 3 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 72c7688b2f1a..c65fd0e82b4d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2951,9 +2951,9 @@ dependencies = [
 
 [[package]]
 name = "foundry-block-explorers"
-version = "0.2.2"
+version = "0.2.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ebeafdc703baf5bb879b276653735cddb2e0244e5e59c24f3aeab5686d801006"
+checksum = "5a056d4aa33a639c0aa1e9e473c25b9b191be30cbea94b31445fac5c272418ae"
 dependencies = [
  "alloy-chains",
  "alloy-json-abi",
@@ -3104,23 +3104,20 @@ dependencies = [
 
 [[package]]
 name = "foundry-compilers"
-version = "0.2.4"
+version = "0.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5106d26aa9a9955c852bed6f9d8994d7229a177969ae5a813008022b022f31ae"
+checksum = "7fd693c0870b3817b215417ff0f2f0df02cddf296dc5524c330f99186a42cf29"
 dependencies = [
  "alloy-json-abi",
  "alloy-primitives",
  "cfg-if",
- "const-hex",
  "dirs 5.0.1",
  "dunce",
  "fs_extra",
  "futures-util",
- "glob",
  "home",
  "md-5 0.10.6",
  "memmap2 0.9.4",
- "num_cpus",
  "once_cell",
  "path-slash",
  "rand 0.8.5",
@@ -3135,7 +3132,6 @@ dependencies = [
  "svm-rs-builds 0.3.5",
  "tempfile",
  "thiserror",
- "tiny-keccak",
  "tokio",
  "tracing",
  "walkdir",
diff --git a/Cargo.toml b/Cargo.toml
index 88ecea80484a..1d5f7db36c06 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -126,8 +126,8 @@ foundry-macros = { path = "crates/macros" }
 foundry-test-utils = { path = "crates/test-utils" }
 
 # solc & compilation utilities
-foundry-block-explorers = { version = "0.2.0", default-features = false }
-foundry-compilers = { version = "0.2.4", default-features = false }
+foundry-block-explorers = { version = "0.2.3", default-features = false }
+foundry-compilers = { version = "0.3.1", default-features = false }
 
 ## revm
 # no default features to avoid c-kzg
@@ -226,4 +226,4 @@ revm-primitives = { git = "https://github.com/bluealloy/revm", branch = "reth_fr
 revm-interpreter = { git = "https://github.com/bluealloy/revm", branch = "reth_freeze" }
 revm-precompile = { git = "https://github.com/bluealloy/revm", branch = "reth_freeze" }
 
-revm-inspectors = { git = "https://github.com/paradigmxyz/evm-inspectors" }
+revm-inspectors = { git = "https://github.com/paradigmxyz/evm-inspectors" }
\ No newline at end of file
diff --git a/crates/forge/bin/cmd/flatten.rs b/crates/forge/bin/cmd/flatten.rs
index 0716af47267a..9831e17cecce 100644
--- a/crates/forge/bin/cmd/flatten.rs
+++ b/crates/forge/bin/cmd/flatten.rs
@@ -4,7 +4,8 @@ use foundry_cli::{
     opts::{CoreBuildArgs, ProjectPathsArgs},
     utils::LoadConfig,
 };
-use foundry_common::fs;
+use foundry_common::{compile::ProjectCompiler, fs};
+use foundry_compilers::{error::SolcError, flatten::Flattener};
 use std::path::PathBuf;
 
 /// CLI arguments for `forge flatten`.
@@ -38,10 +39,24 @@ impl FlattenArgs {
 
         let config = build_args.try_load_config_emit_warnings()?;
 
-        let paths = config.project_paths();
         let target_path = dunce::canonicalize(target_path)?;
-        let flattened =
-            paths.flatten(&target_path).map_err(|err| eyre::eyre!("Failed to flatten: {err}"))?;
+
+        let project = config.ephemeral_no_artifacts_project()?;
+
+        let compiler_output = ProjectCompiler::new().files([target_path.clone()]).compile(&project);
+
+        let flattened = match compiler_output {
+            Ok(compiler_output) => {
+                Flattener::new(&project, &compiler_output, &target_path).map(|f| f.flatten())
+            }
+            Err(_) => {
+                // 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.flatten(&target_path)
+            }
+        }
+        .map_err(|err: SolcError| eyre::eyre!("Failed to flatten: {err}"))?;
 
         match output {
             Some(output) => {