refactor: unify sources and filtered sources (#162)

Allows reusing resources and just doing `retain` instead of
`.into_iter().collect()` on 3 nested maps

Author: DaniPopes

crates/artifacts/solc/src/lib.rs

@@ -106,8 +106,8 @@ impl SolcInput {
     /// Builds one or two inputs from given sources set. Returns two inputs in cases when there are
     /// both Solidity and Yul sources.
     pub fn resolve_and_build(sources: Sources, settings: Settings) -> Vec<Self> {
-        let mut solidity_sources = BTreeMap::new();
-        let mut yul_sources = BTreeMap::new();
+        let mut solidity_sources = Sources::new();
+        let mut yul_sources = Sources::new();
 
         for (file, source) in sources {
             if file.extension().map_or(false, |e| e == "yul") {

crates/artifacts/solc/src/sources.rs

@@ -9,8 +9,97 @@ use std::{
     sync::Arc,
 };
 
-/// An ordered list of files and their source
-pub type Sources = BTreeMap<PathBuf, Source>;
+type SourcesInner = BTreeMap<PathBuf, Source>;
+
+/// An ordered list of files and their source.
+#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
+pub struct Sources(pub SourcesInner);
+
+impl Sources {
+    /// Returns a new instance of [Sources].
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Returns `true` if no sources should have optimized output selection.
+    pub fn all_dirty(&self) -> bool {
+        self.0.values().all(|s| s.is_dirty())
+    }
+
+    /// Returns all entries that should not be optimized.
+    pub fn dirty(&self) -> impl Iterator<Item = (&PathBuf, &Source)> + '_ {
+        self.0.iter().filter(|(_, s)| s.is_dirty())
+    }
+
+    /// Returns all entries that should be optimized.
+    pub fn clean(&self) -> impl Iterator<Item = (&PathBuf, &Source)> + '_ {
+        self.0.iter().filter(|(_, s)| !s.is_dirty())
+    }
+
+    /// Returns all files that should not be optimized.
+    pub fn dirty_files(&self) -> impl Iterator<Item = &PathBuf> + '_ {
+        self.dirty().map(|(k, _)| k)
+    }
+}
+
+impl std::ops::Deref for Sources {
+    type Target = SourcesInner;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl std::ops::DerefMut for Sources {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl<I> From<I> for Sources
+where
+    SourcesInner: From<I>,
+{
+    fn from(value: I) -> Self {
+        Self(From::from(value))
+    }
+}
+
+impl<I> FromIterator<I> for Sources
+where
+    SourcesInner: FromIterator<I>,
+{
+    fn from_iter<T: IntoIterator<Item = I>>(iter: T) -> Self {
+        Self(FromIterator::from_iter(iter))
+    }
+}
+
+impl IntoIterator for Sources {
+    type Item = <SourcesInner as IntoIterator>::Item;
+    type IntoIter = <SourcesInner as IntoIterator>::IntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter()
+    }
+}
+
+impl<'a> IntoIterator for &'a Sources {
+    type Item = <&'a SourcesInner as IntoIterator>::Item;
+    type IntoIter = <&'a SourcesInner as IntoIterator>::IntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.iter()
+    }
+}
+
+impl<'a> IntoIterator for &'a mut Sources {
+    type Item = <&'a mut SourcesInner as IntoIterator>::Item;
+    type IntoIter = <&'a mut SourcesInner as IntoIterator>::IntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.iter_mut()
+    }
+}
 
 /// Content of a solidity file
 ///
@@ -23,12 +112,14 @@ pub struct Source {
     /// conflicting versions then the same [Source] may be required by conflicting versions and
     /// needs to be duplicated.
     pub content: Arc<String>,
+    #[serde(skip, default)]
+    pub kind: SourceCompilationKind,
 }
 
 impl Source {
     /// Creates a new instance of [Source] with the given content.
     pub fn new(content: impl Into<String>) -> Self {
-        Self { content: Arc::new(content.into()) }
+        Self { content: Arc::new(content.into()), kind: SourceCompilationKind::Complete }
     }
 
     /// Reads the file's content
@@ -45,6 +136,11 @@ impl Source {
         Ok(Self::new(content))
     }
 
+    /// Returns `true` if the source should be compiled with full output selection.
+    pub fn is_dirty(&self) -> bool {
+        self.kind.is_dirty()
+    }
+
     /// Recursively finds all source files under the given dir path and reads them all
     pub fn read_all_from(dir: &Path, extensions: &[&str]) -> Result<Sources, SolcIoError> {
         Self::read_all_files(utils::source_files(dir, extensions))
@@ -91,7 +187,8 @@ impl Source {
             .par_bridge()
             .map(Into::into)
             .map(|file| Self::read(&file).map(|source| (file, source)))
-            .collect()
+            .collect::<Result<BTreeMap<_, _>, _>>()
+            .map(Sources)
     }
 
     /// Generate a non-cryptographically secure checksum of the file's content
@@ -156,3 +253,21 @@ impl AsRef<[u8]> for Source {
         self.content.as_bytes()
     }
 }
+
+/// Represents the state of a filtered [`Source`].
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub enum SourceCompilationKind {
+    /// We need a complete compilation output for the source.
+    #[default]
+    Complete,
+    /// A source for which we don't need a complete output and want to optimize its compilation by
+    /// reducing output selection.
+    Optimized,
+}
+
+impl SourceCompilationKind {
+    /// Whether this file should be compiled with full output selection
+    pub fn is_dirty(&self) -> bool {
+        matches!(self, Self::Complete)
+    }
+}

crates/compilers/src/buildinfo.rs

@@ -128,18 +128,16 @@ impl<L: Language> RawBuildInfo<L> {
 
 #[cfg(test)]
 mod tests {
-    use foundry_compilers_artifacts::{sources::Source, Error, SolcLanguage};
-
-    use crate::compilers::solc::SolcVersionedInput;
-
     use super::*;
-    use std::{collections::BTreeMap, path::PathBuf};
+    use crate::compilers::solc::SolcVersionedInput;
+    use foundry_compilers_artifacts::{sources::Source, Error, SolcLanguage, Sources};
+    use std::path::PathBuf;
 
     #[test]
     fn build_info_serde() {
         let v: Version = "0.8.4+commit.c7e474f2".parse().unwrap();
         let input = SolcVersionedInput::build(
-            BTreeMap::from([(PathBuf::from("input.sol"), Source::new(""))]),
+            Sources::from([(PathBuf::from("input.sol"), Source::new(""))]),
             Default::default(),
             SolcLanguage::Solidity,
             v,

crates/compilers/src/cache.rs

@@ -5,8 +5,8 @@ use crate::{
     compilers::{Compiler, CompilerSettings, Language},
     output::Builds,
     resolver::GraphEdges,
-    ArtifactFile, ArtifactOutput, Artifacts, ArtifactsMap, FilteredSources, Graph, OutputContext,
-    Project, ProjectPaths, ProjectPathsConfig, SourceCompilationKind,
+    ArtifactFile, ArtifactOutput, Artifacts, ArtifactsMap, Graph, OutputContext, Project,
+    ProjectPaths, ProjectPathsConfig, SourceCompilationKind,
 };
 use foundry_compilers_artifacts::{
     sources::{Source, Sources},
@@ -648,10 +648,10 @@ impl<'a, T: ArtifactOutput, C: Compiler> ArtifactsCacheInner<'a, T, C> {
     /// 2. [SourceCompilationKind::Optimized] - the file is not dirty, but is imported by a dirty
     ///    file and thus will be processed by solc. For such files we don't need full data, so we
     ///    are marking them as clean to optimize output selection later.
-    fn filter(&mut self, sources: Sources, version: &Version) -> FilteredSources {
+    fn filter(&mut self, sources: &mut Sources, version: &Version) {
         // sources that should be passed to compiler.
-        let mut compile_complete = BTreeSet::new();
-        let mut compile_optimized = BTreeSet::new();
+        let mut compile_complete = HashSet::new();
+        let mut compile_optimized = HashSet::new();
 
         for (file, source) in sources.iter() {
             self.sources_in_scope.insert(file.clone(), version.clone());
@@ -677,20 +677,16 @@ impl<'a, T: ArtifactOutput, C: Compiler> ArtifactsCacheInner<'a, T, C> {
             }
         }
 
-        let filtered = sources
-            .into_iter()
-            .filter_map(|(file, source)| {
-                if compile_complete.contains(&file) {
-                    Some((file, SourceCompilationKind::Complete(source)))
-                } else if compile_optimized.contains(&file) {
-                    Some((file, SourceCompilationKind::Optimized(source)))
-                } else {
-                    None
-                }
-            })
-            .collect();
-
-        FilteredSources(filtered)
+        sources.retain(|file, source| {
+            source.kind = if compile_complete.contains(file) {
+                SourceCompilationKind::Complete
+            } else if compile_optimized.contains(file) {
+                SourceCompilationKind::Optimized
+            } else {
+                return false;
+            };
+            true
+        });
     }
 
     /// Returns whether we are missing artifacts for the given file and version.
@@ -746,7 +742,7 @@ impl<'a, T: ArtifactOutput, C: Compiler> ArtifactsCacheInner<'a, T, C> {
         // Iterate over existing cache entries.
         let files = self.cache.files.keys().cloned().collect::<HashSet<_>>();
 
-        let mut sources = BTreeMap::new();
+        let mut sources = Sources::new();
 
         // Read all sources, marking entries as dirty on I/O errors.
         for file in &files {
@@ -939,14 +935,14 @@ impl<'a, T: ArtifactOutput, C: Compiler> ArtifactsCache<'a, T, C> {
     // only useful for debugging for debugging purposes
     pub fn as_cached(&self) -> Option<&ArtifactsCacheInner<'a, T, C>> {
         match self {
-            ArtifactsCache::Ephemeral(_, _) => None,
+            ArtifactsCache::Ephemeral(..) => None,
             ArtifactsCache::Cached(cached) => Some(cached),
         }
     }
 
     pub fn output_ctx(&self) -> OutputContext<'_> {
         match self {
-            ArtifactsCache::Ephemeral(_, _) => Default::default(),
+            ArtifactsCache::Ephemeral(..) => Default::default(),
             ArtifactsCache::Cached(inner) => OutputContext::new(&inner.cache),
         }
     }
@@ -961,15 +957,15 @@ impl<'a, T: ArtifactOutput, C: Compiler> ArtifactsCache<'a, T, C> {
     /// Adds the file's hashes to the set if not set yet
     pub fn remove_dirty_sources(&mut self) {
         match self {
-            ArtifactsCache::Ephemeral(_, _) => {}
+            ArtifactsCache::Ephemeral(..) => {}
             ArtifactsCache::Cached(cache) => cache.find_and_remove_dirty(),
         }
     }
 
     /// Filters out those sources that don't need to be compiled
-    pub fn filter(&mut self, sources: Sources, version: &Version) -> FilteredSources {
+    pub fn filter(&mut self, sources: &mut Sources, version: &Version) {
         match self {
-            ArtifactsCache::Ephemeral(_, _) => sources.into(),
+            ArtifactsCache::Ephemeral(..) => {}
             ArtifactsCache::Cached(cache) => cache.filter(sources, version),
         }
     }