Autodetect and recompile mocks

Author: grandizzy

crates/compilers/src/cache.rs

@@ -47,6 +47,7 @@ pub struct CompilerCache<S = Settings> {
     pub builds: BTreeSet<String>,
     pub profiles: BTreeMap<String, S>,
     pub preprocessed: bool,
+    pub mocks: HashSet<PathBuf>,
 }
 
 impl<S> CompilerCache<S> {
@@ -58,6 +59,7 @@ impl<S> CompilerCache<S> {
             builds: Default::default(),
             profiles: Default::default(),
             preprocessed,
+            mocks: Default::default(),
         }
     }
 }
@@ -381,6 +383,7 @@ impl<S> Default for CompilerCache<S> {
             paths: Default::default(),
             profiles: Default::default(),
             preprocessed: false,
+            mocks: Default::default(),
         }
     }
 }
@@ -792,7 +795,7 @@ impl<T: ArtifactOutput<CompilerContract = C::CompilerContract>, C: Compiler>
         self.missing_extra_files()
     }
 
-    // Walks over all cache entires, detects dirty files and removes them from cache.
+    // Walks over all cache entries, detects dirty files and removes them from cache.
     fn find_and_remove_dirty(&mut self) {
         fn populate_dirty_files<D>(
             file: &Path,
@@ -892,13 +895,21 @@ impl<T: ArtifactOutput<CompilerContract = C::CompilerContract>, C: Compiler>
                             self.dirty_sources.insert(file.clone());
                             break;
                         // For non-src files we mark them as dirty only if they import dirty
-                        // non-src file or src file for which
-                        // interface representation changed.
+                        // non-src file or src file for which interface representation changed.
+                        // For identified mock contracts (non-src contracts that extends contracts
+                        // from src file) we mark edges as dirty.
                         } else if !is_src
                             && self.dirty_sources.contains(import)
-                            && (!self.is_source_file(import) || self.is_dirty_impl(import, true))
+                            && (!self.is_source_file(import)
+                                || self.is_dirty_impl(import, true)
+                                || self.cache.mocks.contains(file))
                         {
-                            self.dirty_sources.insert(file.clone());
+                            if self.cache.mocks.contains(file) {
+                                // Mark all mock edges as dirty.
+                                populate_dirty_files(file, &mut self.dirty_sources, &edges);
+                            } else {
+                                self.dirty_sources.insert(file.clone());
+                            }
                         }
                     }
                 }
@@ -1122,6 +1133,16 @@ impl<'a, T: ArtifactOutput<CompilerContract = C::CompilerContract>, C: Compiler>
         }
     }
 
+    /// Adds the file's hashes to the set if not set yet
+    pub fn add_mocks(&mut self, mocks: Option<HashSet<PathBuf>>) {
+        if let Some(mocks) = mocks {
+            match self {
+                ArtifactsCache::Ephemeral(..) => {}
+                ArtifactsCache::Cached(cache) => cache.cache.mocks.extend(mocks),
+            }
+        }
+    }
+
     /// Filters out those sources that don't need to be compiled
     pub fn filter(&mut self, sources: &mut Sources, version: &Version, profile: &str) {
         match self {

crates/compilers/src/compile/project.rs

@@ -115,19 +115,25 @@ use crate::{
 use foundry_compilers_core::error::Result;
 use rayon::prelude::*;
 use semver::Version;
-use std::{collections::HashMap, fmt::Debug, path::PathBuf, time::Instant};
+use std::{
+    collections::{HashMap, HashSet},
+    fmt::Debug,
+    path::PathBuf,
+    time::Instant,
+};
 
 /// A set of different Solc installations with their version and the sources to be compiled
 pub(crate) type VersionedSources<'a, L, S> = HashMap<L, Vec<(Version, Sources, (&'a str, &'a S))>>;
 
 /// Invoked before the actual compiler invocation and can override the input.
+/// Returns preprocessed compiler input and identified mocks (if any) to be stored in cache.
 pub trait Preprocessor<C: Compiler>: Debug {
     fn preprocess(
         &self,
         compiler: &C,
         input: C::Input,
         paths: &ProjectPathsConfig<C::Language>,
-    ) -> Result<C::Input>;
+    ) -> Result<(C::Input, Option<HashSet<PathBuf>>)>;
 }
 
 #[derive(Debug)]
@@ -469,6 +475,7 @@ impl<L: Language, S: CompilerSettings> CompilerSources<'_, L, S> {
         include_paths.extend(graph.include_paths().clone());
 
         let mut jobs = Vec::new();
+        let mut mocks = None;
         for (language, versioned_sources) in self.sources {
             for (version, sources, (profile, opt_settings)) in versioned_sources {
                 let mut opt_settings = opt_settings.clone();
@@ -503,13 +510,16 @@ impl<L: Language, S: CompilerSettings> CompilerSources<'_, L, S> {
                 input.strip_prefix(project.paths.root.as_path());
 
                 if let Some(preprocessor) = preprocessor.as_ref() {
-                    input = preprocessor.preprocess(&project.compiler, input, &project.paths)?;
+                    (input, mocks) =
+                        preprocessor.preprocess(&project.compiler, input, &project.paths)?;
                 }
 
                 jobs.push((input, profile, actually_dirty));
             }
         }
 
+        cache.add_mocks(mocks);
+
         let results = if let Some(num_jobs) = jobs_cnt {
             compile_parallel(&project.compiler, jobs, num_jobs)
         } else {

crates/compilers/src/preprocessor/deps.rs

@@ -15,7 +15,7 @@ use solar_sema::{
 use std::{
     collections::{BTreeMap, BTreeSet, HashSet},
     ops::ControlFlow,
-    path::PathBuf,
+    path::{Path, PathBuf},
 };
 
 /// Holds data about referenced source contracts and bytecode dependencies.
@@ -24,12 +24,21 @@ pub(crate) struct PreprocessorDependencies {
     pub preprocessed_contracts: BTreeMap<ContractId, Vec<BytecodeDependency>>,
     // Referenced contract ids.
     pub referenced_contracts: HashSet<ContractId>,
+    // Mock contract paths (with a base contract from src dir).
+    pub mocks: HashSet<PathBuf>,
 }
 
 impl PreprocessorDependencies {
-    pub fn new(sess: &Session, hir: &Hir<'_>, paths: &[PathBuf], src_dir: &PathBuf) -> Self {
+    pub fn new(
+        sess: &Session,
+        hir: &Hir<'_>,
+        paths: &[PathBuf],
+        src_dir: &PathBuf,
+        root_dir: &Path,
+    ) -> Self {
         let mut preprocessed_contracts = BTreeMap::new();
         let mut referenced_contracts = HashSet::new();
+        let mut mocks = HashSet::new();
         for contract_id in hir.contract_ids() {
             let contract = hir.contract(contract_id);
             let source = hir.source(contract.source);
@@ -40,6 +49,24 @@ impl PreprocessorDependencies {
 
             // Collect dependencies only for tests and scripts.
             if !paths.contains(path) {
+                let path = path.display();
+                trace!("{path} is not test or script");
+                continue;
+            }
+
+            // Do not collect dependencies for mock contracts. Walk through base contracts and
+            // check if they're from src dir.
+            if contract.linearized_bases.iter().any(|base_contract_id| {
+                let base_contract = hir.contract(*base_contract_id);
+                let FileName::Real(path) = &hir.source(base_contract.source).file.name else {
+                    return false;
+                };
+                path.starts_with(src_dir)
+            }) {
+                // Record mock contracts to be evicted from preprocessed cache.
+                mocks.insert(root_dir.join(path));
+                let path = path.display();
+                trace!("found mock contract {path}");
                 continue;
             }
 
@@ -58,7 +85,7 @@ impl PreprocessorDependencies {
             // Record collected referenced contract ids.
             referenced_contracts.extend(deps_collector.referenced_contracts);
         }
-        Self { preprocessed_contracts, referenced_contracts }
+        Self { preprocessed_contracts, referenced_contracts, mocks }
     }
 }
 

crates/compilers/src/preprocessor/mod.rs

@@ -23,7 +23,10 @@ use solar_parse::{
     Parser,
 };
 use solar_sema::{hir::Arena, ParsingContext};
-use std::path::{Path, PathBuf};
+use std::{
+    collections::HashSet,
+    path::{Path, PathBuf},
+};
 
 mod data;
 mod deps;
@@ -60,62 +63,66 @@ impl Preprocessor<SolcCompiler> for TestOptimizerPreprocessor {
         _solc: &SolcCompiler,
         mut input: SolcVersionedInput,
         paths: &ProjectPathsConfig<SolcLanguage>,
-    ) -> Result<SolcVersionedInput> {
+    ) -> Result<(SolcVersionedInput, Option<HashSet<PathBuf>>)> {
         let sources = &mut input.input.sources;
         // Skip if we are not preprocessing any tests or scripts. Avoids unnecessary AST parsing.
         if sources.iter().all(|(path, _)| !is_test_or_script(path, paths)) {
             trace!("no tests or sources to preprocess");
-            return Ok(input);
+            return Ok((input, None));
         }
 
         let sess = Session::builder().with_buffer_emitter(Default::default()).build();
-        let _ = sess.enter_parallel(|| -> solar_parse::interface::Result<()> {
-            let hir_arena = Arena::new();
-            let mut parsing_context = ParsingContext::new(&sess);
-            // Set remappings into HIR parsing context.
-            for remapping in &paths.remappings {
-                parsing_context
-                    .file_resolver
-                    .add_import_map(PathBuf::from(&remapping.name), PathBuf::from(&remapping.path));
-            }
-            // Load and parse test and script contracts only (dependencies are automatically
-            // resolved).
-            let preprocessed_paths = sources
-                .into_iter()
-                .filter(|(path, source)| {
-                    is_test_or_script(path, paths) && !source.content.is_empty()
-                })
-                .map(|(path, _)| path.clone())
-                .collect_vec();
-            parsing_context.load_files(&preprocessed_paths)?;
+        let mocks = sess
+            .enter_parallel(|| -> solar_parse::interface::Result<Option<HashSet<PathBuf>>> {
+                let hir_arena = Arena::new();
+                let mut parsing_context = ParsingContext::new(&sess);
+                // Set remappings into HIR parsing context.
+                for remapping in &paths.remappings {
+                    parsing_context.file_resolver.add_import_map(
+                        PathBuf::from(&remapping.name),
+                        PathBuf::from(&remapping.path),
+                    );
+                }
+                // Load and parse test and script contracts only (dependencies are automatically
+                // resolved).
+                let preprocessed_paths = sources
+                    .into_iter()
+                    .filter(|(path, source)| {
+                        is_test_or_script(path, paths) && !source.content.is_empty()
+                    })
+                    .map(|(path, _)| path.clone())
+                    .collect_vec();
+                parsing_context.load_files(&preprocessed_paths)?;
 
-            let hir = &parsing_context.parse_and_lower_to_hir(&hir_arena)?;
-            // Collect tests and scripts dependencies.
-            let deps = PreprocessorDependencies::new(
-                &sess,
-                hir,
-                &preprocessed_paths,
-                &paths.paths_relative().sources,
-            );
-            // Collect data of source contracts referenced in tests and scripts.
-            let data = collect_preprocessor_data(&sess, hir, &deps.referenced_contracts);
+                let hir = &parsing_context.parse_and_lower_to_hir(&hir_arena)?;
+                // Collect tests and scripts dependencies and identify mock contracts.
+                let deps = PreprocessorDependencies::new(
+                    &sess,
+                    hir,
+                    &preprocessed_paths,
+                    &paths.paths_relative().sources,
+                    &paths.root,
+                );
+                // Collect data of source contracts referenced in tests and scripts.
+                let data = collect_preprocessor_data(&sess, hir, &deps.referenced_contracts);
 
-            // Extend existing sources with preprocessor deploy helper sources.
-            sources.extend(create_deploy_helpers(&data));
+                // Extend existing sources with preprocessor deploy helper sources.
+                sources.extend(create_deploy_helpers(&data));
 
-            // Generate and apply preprocessor source updates.
-            apply_updates(sources, remove_bytecode_dependencies(hir, &deps, &data));
+                // Generate and apply preprocessor source updates.
+                apply_updates(sources, remove_bytecode_dependencies(hir, &deps, &data));
 
-            Ok(())
-        });
+                Ok(Some(deps.mocks))
+            })
+            .unwrap_or_default();
 
         // Return if any diagnostics emitted during content parsing.
         if let Err(err) = sess.emitted_errors().unwrap() {
             trace!("failed preprocessing {err}");
             return Err(SolcError::Message(err.to_string()));
         }
 
-        Ok(input)
+        Ok((input, mocks))
     }
 }
 
@@ -125,18 +132,18 @@ impl Preprocessor<MultiCompiler> for TestOptimizerPreprocessor {
         compiler: &MultiCompiler,
         input: <MultiCompiler as Compiler>::Input,
         paths: &ProjectPathsConfig<MultiCompilerLanguage>,
-    ) -> Result<<MultiCompiler as Compiler>::Input> {
+    ) -> Result<(<MultiCompiler as Compiler>::Input, Option<HashSet<PathBuf>>)> {
         match input {
             MultiCompilerInput::Solc(input) => {
                 if let Some(solc) = &compiler.solc {
                     let paths = paths.clone().with_language::<SolcLanguage>();
-                    let input = self.preprocess(solc, input, &paths)?;
-                    Ok(MultiCompilerInput::Solc(input))
+                    let (input, mocks) = self.preprocess(solc, input, &paths)?;
+                    Ok((MultiCompilerInput::Solc(input), mocks))
                 } else {
-                    Ok(MultiCompilerInput::Solc(input))
+                    Ok((MultiCompilerInput::Solc(input), None))
                 }
             }
-            MultiCompilerInput::Vyper(input) => Ok(MultiCompilerInput::Vyper(input)),
+            MultiCompilerInput::Vyper(input) => Ok((MultiCompilerInput::Vyper(input), None)),
         }
     }
 }