Small flattener features (#75)

This PR implements two small flattener features mentioned in
makerdao/spells-mainnet#391 (comment)

1. Combine pragmas from different sources into one pragma so that
version restrictions are kept the same before and after flattening
2. Comments specifying original source file of the flattened code (`//
src/File.sol`)

Cargo.toml

@@ -57,6 +57,7 @@ home = "0.5"
 svm = { package = "svm-rs", version = "0.3", default-features = false, optional = true }
 svm-builds = { package = "svm-rs-builds", version = "0.3", default-features = false, optional = true }
 sha2 = { version = "0.10", default-features = false, optional = true }
+itertools = "0.12"
 
 [dev-dependencies]
 alloy-primitives = { version = "0.6", features = ["serde", "rand"] }

src/config.rs

@@ -2,7 +2,7 @@ use crate::{
     artifacts::{output_selection::ContractOutputSelection, Settings},
     cache::SOLIDITY_FILES_CACHE_FILENAME,
     error::{Result, SolcError, SolcIoError},
-    flatten::collect_ordered_deps,
+    flatten::{collect_ordered_deps, combine_version_pragmas},
     remappings::Remapping,
     resolver::{Graph, SolImportAlias},
     utils, Source, Sources,
@@ -412,6 +412,8 @@ impl ProjectPathsConfig {
         let ordered_deps = collect_ordered_deps(&flatten_target, self, &graph)?;
 
         let mut sources = Vec::new();
+        let mut experimental_pragma = None;
+        let mut version_pragmas = Vec::new();
 
         let mut result = String::new();
 
@@ -420,7 +422,7 @@ impl ProjectPathsConfig {
                 SolcError::msg(format!("cannot resolve file at {}", path.display()))
             })?;
             let node = graph.node(*node_id);
-            let content = node.content().to_owned();
+            let content = node.content();
 
             // Firstly we strip all licesnses, verson pragmas
             // We keep target file pragma and license placing them in the beginning of the result.
@@ -434,17 +436,14 @@ impl ProjectPathsConfig {
                 }
             }
             if let Some(version) = node.version() {
+                let content = &content[version.loc()];
                 ranges_to_remove.push(version.loc());
-                if *path == flatten_target {
-                    result.push_str(&content[version.loc()]);
-                    result.push('\n');
-                }
+                version_pragmas.push(content);
             }
             if let Some(experimental) = node.experimental() {
                 ranges_to_remove.push(experimental.loc());
-                if *path == flatten_target {
-                    result.push_str(&content[experimental.loc()]);
-                    result.push('\n');
+                if experimental_pragma.is_none() {
+                    experimental_pragma = Some(content[experimental.loc()].to_owned());
                 }
             }
             for import in node.imports() {
@@ -489,12 +488,27 @@ impl ProjectPathsConfig {
                 }
             }
 
+            let content = format!(
+                "// {}\n{}",
+                path.strip_prefix(&self.root).unwrap_or(path).display(),
+                content
+            );
+
             sources.push(content);
         }
 
+        if let Some(version) = combine_version_pragmas(version_pragmas) {
+            result.push_str(&version);
+            result.push('\n');
+        }
+        if let Some(experimental) = experimental_pragma {
+            result.push_str(&experimental);
+            result.push('\n');
+        }
+
         for source in sources {
-            result.push_str(&source);
             result.push_str("\n\n");
+            result.push_str(&source);
         }
 
         Ok(format!("{}\n", utils::RE_THREE_OR_MORE_NEWLINES.replace_all(&result, "\n\n").trim()))

src/flatten.rs

@@ -4,6 +4,8 @@ use std::{
     path::{Path, PathBuf},
 };
 
+use itertools::Itertools;
+
 use crate::{
     artifacts::{
         ast::SourceLocation,
@@ -12,7 +14,7 @@ use crate::{
         MemberAccess, Source, SourceUnit, SourceUnitPart, Sources,
     },
     error::SolcError,
-    utils, Graph, Project, ProjectCompileOutput, ProjectPathsConfig, Result,
+    utils, Graph, Project, ProjectCompileOutput, ProjectPathsConfig, Result, Solc,
 };
 
 /// Alternative of `SourceLocation` which includes path of the file.
@@ -95,11 +97,11 @@ impl Visitor for ReferencesCollector {
 /// source_path -> (start, end, new_value)
 type Updates = HashMap<PathBuf, HashSet<(usize, usize, String)>>;
 
-struct FlatteningResult<'a> {
+pub struct FlatteningResult<'a> {
     /// Updated source in the order they shoud be written to the output file.
     sources: Vec<String>,
     /// Pragmas that should be present in the target file.
-    pragmas: Vec<&'a str>,
+    pragmas: Vec<String>,
     /// License identifier that should be present in the target file.
     license: Option<&'a str>,
 }
@@ -108,7 +110,7 @@ impl<'a> FlatteningResult<'a> {
     fn new(
         flattener: &Flattener,
         mut updates: Updates,
-        pragmas: Vec<&'a str>,
+        pragmas: Vec<String>,
         license: Option<&'a str>,
     ) -> Self {
         let mut sources = Vec::new();
@@ -127,7 +129,12 @@ impl<'a> FlatteningResult<'a> {
                     offset += new_value.len() as isize - (end - start) as isize;
                 }
             }
-            sources.push(String::from_utf8(content).unwrap());
+            let content = format!(
+                "// {}\n{}",
+                path.strip_prefix(&flattener.project_root).unwrap_or(path).display(),
+                String::from_utf8(content).unwrap()
+            );
+            sources.push(content);
         }
 
         Self { sources, pragmas, license }
@@ -143,7 +150,7 @@ impl<'a> FlatteningResult<'a> {
             result.push_str(&format!("{}\n", pragma));
         }
         for source in &self.sources {
-            result.push_str(&format!("{}\n\n", source));
+            result.push_str(&format!("\n\n{}", source));
         }
 
         format!("{}\n", utils::RE_THREE_OR_MORE_NEWLINES.replace_all(&result, "\n\n").trim())
@@ -160,6 +167,8 @@ pub struct Flattener {
     asts: Vec<(PathBuf, SourceUnit)>,
     /// Sources in the order they should be written to the output file.
     ordered_sources: Vec<PathBuf>,
+    /// Project root directory.
+    project_root: PathBuf,
 }
 
 impl Flattener {
@@ -177,9 +186,9 @@ impl Flattener {
         let sources = Source::read_all_files(input_files)?;
         let graph = Graph::resolve_sources(&project.paths, sources)?;
 
-        let ordered_deps = collect_ordered_deps(&target.to_path_buf(), &project.paths, &graph)?;
+        let ordered_sources = collect_ordered_deps(&target.to_path_buf(), &project.paths, &graph)?;
 
-        let sources = Source::read_all(&ordered_deps)?;
+        let sources = Source::read_all(&ordered_sources)?;
 
         // Convert all ASTs from artifacts to strongly typed ASTs
         let mut asts: Vec<(PathBuf, SourceUnit)> = Vec::new();
@@ -197,7 +206,13 @@ impl Flattener {
             asts.push((PathBuf::from(path), serde_json::from_str(&serde_json::to_string(ast)?)?));
         }
 
-        Ok(Flattener { target: target.into(), sources, asts, ordered_sources: ordered_deps })
+        Ok(Flattener {
+            target: target.into(),
+            sources,
+            asts,
+            ordered_sources,
+            project_root: project.root().clone(),
+        })
     }
 
     /// Flattens target file and returns the result as a string
@@ -227,7 +242,7 @@ impl Flattener {
     fn flatten_result<'a>(
         &'a self,
         updates: Updates,
-        target_pragmas: Vec<&'a str>,
+        target_pragmas: Vec<String>,
         target_license: Option<&'a str>,
     ) -> FlatteningResult<'_> {
         FlatteningResult::new(self, updates, target_pragmas, target_license)
@@ -629,25 +644,22 @@ impl Flattener {
             .collect()
     }
 
-    /// Removes all pragma directives from all sources. Returns Vec of pragmas that were found in
-    /// target file.
-    fn process_pragmas(&self, updates: &mut Updates) -> Vec<&str> {
-        // Pragmas that will be used in the resulted file
-        let mut target_pragmas = Vec::new();
+    /// Removes all pragma directives from all sources. Returns Vec with experimental and combined
+    /// version pragmas (if present).
+    fn process_pragmas(&self, updates: &mut Updates) -> Vec<String> {
+        let mut experimental = None;
 
         let pragmas = self.collect_pragmas();
-
-        let mut seen_experimental = false;
+        let mut version_pragmas = Vec::new();
 
         for loc in &pragmas {
             let pragma_content = self.read_location(loc);
             if pragma_content.contains("experimental") {
-                if !seen_experimental {
-                    seen_experimental = true;
-                    target_pragmas.push(loc);
+                if experimental.is_none() {
+                    experimental = Some(self.read_location(loc).to_string());
                 }
-            } else if loc.path == self.target {
-                target_pragmas.push(loc);
+            } else if pragma_content.contains("solidity") {
+                version_pragmas.push(pragma_content);
             }
 
             updates.entry(loc.path.clone()).or_default().insert((
@@ -657,8 +669,17 @@ impl Flattener {
             ));
         }
 
-        target_pragmas.sort_by_key(|loc| loc.start);
-        target_pragmas.iter().map(|loc| self.read_location(loc)).collect::<Vec<_>>()
+        let mut pragmas = Vec::new();
+
+        if let Some(version_pragma) = combine_version_pragmas(version_pragmas) {
+            pragmas.push(version_pragma);
+        }
+
+        if let Some(pragma) = experimental {
+            pragmas.push(pragma);
+        }
+
+        pragmas
     }
 
     // Collects all pragma directives locations.
@@ -785,3 +806,27 @@ pub fn collect_ordered_deps(
 
     Ok(ordered_deps)
 }
+
+pub fn combine_version_pragmas(pragmas: Vec<&str>) -> Option<String> {
+    let mut versions = pragmas
+        .into_iter()
+        .filter_map(|p| {
+            Solc::version_req(
+                p.replace("pragma", "").replace("solidity", "").replace(';', "").trim(),
+            )
+            .ok()
+        })
+        .flat_map(|req| req.comparators)
+        .collect::<HashSet<_>>()
+        .into_iter()
+        .map(|comp| comp.to_string())
+        .collect::<Vec<_>>();
+
+    versions.sort();
+
+    if !versions.is_empty() {
+        return Some(format!("pragma solidity {};", versions.iter().format(" ")));
+    }
+
+    None
+}

src/resolver/mod.rs

@@ -151,7 +151,7 @@ impl GraphEdges {
     /// Returns all files that import the given file
     pub fn importers(&self, file: impl AsRef<Path>) -> HashSet<&PathBuf> {
         if let Some(start) = self.indices.get(file.as_ref()).copied() {
-            self.rev_edges[start].iter().map(move |idx| &self.rev_indices[&idx]).collect()
+            self.rev_edges[start].iter().map(move |idx| &self.rev_indices[idx]).collect()
         } else {
             HashSet::new()
         }