perf: parallelize Remapping::get_many

Author: DaniPopes

crates/artifacts/solc/src/remappings/find.rs

@@ -1,8 +1,10 @@
 use super::Remapping;
 use foundry_compilers_core::utils;
+use rayon::prelude::*;
 use std::{
     collections::{btree_map::Entry, BTreeMap, HashSet},
     path::{Path, PathBuf},
+    sync::Mutex,
 };
 
 const DAPPTOOLS_CONTRACTS_DIR: &str = "src";
@@ -52,6 +54,7 @@ impl Remapping {
     /// which would be multiple rededications according to our rules ("governance", "protocol-v2"),
     /// are unified into `@aave` by looking at their common ancestor, the root of this subdirectory
     /// (`@aave`)
+    #[instrument(level = "trace", name = "Remapping::find_many")]
     pub fn find_many(dir: &Path) -> Vec<Self> {
         /// prioritize
         ///   - ("a", "1/2") over ("a", "1/2/3")
@@ -76,41 +79,31 @@ impl Remapping {
             }
         }
 
-        // all combined remappings from all subdirs
-        let mut all_remappings = BTreeMap::new();
-
         let is_inside_node_modules = dir.ends_with("node_modules");
+        let visited_symlink_dirs = Mutex::new(HashSet::new());
 
-        let mut visited_symlink_dirs = HashSet::new();
         // iterate over all dirs that are children of the root
-        for dir in walkdir::WalkDir::new(dir)
-            .sort_by_file_name()
-            .follow_links(true)
-            .min_depth(1)
-            .max_depth(1)
-            .into_iter()
-            .filter_entry(|e| !is_hidden(e))
-            .filter_map(Result::ok)
+        let candidates = read_dir(dir)
             .filter(|e| e.file_type().is_dir())
-        {
-            let depth1_dir = dir.path();
-            // check all remappings in this depth 1 folder
-            let candidates = find_remapping_candidates(
-                depth1_dir,
-                depth1_dir,
-                0,
-                is_inside_node_modules,
-                &mut visited_symlink_dirs,
-            );
-
-            for candidate in candidates {
-                if let Some(name) = candidate.window_start.file_name().and_then(|s| s.to_str()) {
-                    insert_prioritized(
-                        &mut all_remappings,
-                        format!("{name}/"),
-                        candidate.source_dir,
-                    );
-                }
+            .collect::<Vec<_>>()
+            .par_iter()
+            .flat_map_iter(|entry| {
+                let dir = entry.path();
+                find_remapping_candidates(
+                    dir,
+                    dir,
+                    0,
+                    is_inside_node_modules,
+                    &visited_symlink_dirs,
+                )
+            })
+            .collect::<Vec<_>>();
+
+        // all combined remappings from all subdirs
+        let mut all_remappings = BTreeMap::new();
+        for candidate in candidates {
+            if let Some(name) = candidate.window_start.file_name().and_then(|s| s.to_str()) {
+                insert_prioritized(&mut all_remappings, format!("{name}/"), candidate.source_dir);
             }
         }
 
@@ -298,32 +291,23 @@ fn is_hidden(entry: &walkdir::DirEntry) -> bool {
 
 /// Finds all remappings in the directory recursively
 ///
-/// Note: this supports symlinks and will short-circuit if a symlink dir has already been visited, this can occur in pnpm setups: <https://github.com/foundry-rs/foundry/issues/7820>
+/// Note: this supports symlinks and will short-circuit if a symlink dir has already been visited,
+/// this can occur in pnpm setups: <https://github.com/foundry-rs/foundry/issues/7820>
 fn find_remapping_candidates(
     current_dir: &Path,
     open: &Path,
     current_level: usize,
     is_inside_node_modules: bool,
-    visited_symlink_dirs: &mut HashSet<PathBuf>,
+    visited_symlink_dirs: &Mutex<HashSet<PathBuf>>,
 ) -> Vec<Candidate> {
+    trace!("find_remapping_candidates({})", current_dir.display());
+
     // this is a marker if the current root is a candidate for a remapping
     let mut is_candidate = false;
 
-    // all found candidates
-    let mut candidates = Vec::new();
-
     // scan all entries in the current dir
-    for entry in walkdir::WalkDir::new(current_dir)
-        .sort_by_file_name()
-        .follow_links(true)
-        .min_depth(1)
-        .max_depth(1)
-        .into_iter()
-        .filter_entry(|e| !is_hidden(e))
-        .filter_map(Result::ok)
-    {
-        let entry: walkdir::DirEntry = entry;
-
+    let mut search = Vec::new();
+    for entry in read_dir(current_dir) {
         // found a solidity file directly the current dir
         if !is_candidate
             && entry.file_type().is_file()
@@ -341,7 +325,7 @@ fn find_remapping_candidates(
             // ```
             if entry.path_is_symlink() {
                 if let Ok(target) = utils::canonicalize(entry.path()) {
-                    if !visited_symlink_dirs.insert(target.clone()) {
+                    if !visited_symlink_dirs.lock().unwrap().insert(target.clone()) {
                         // short-circuiting if we've already visited the symlink
                         return Vec::new();
                     }
@@ -357,38 +341,45 @@ fn find_remapping_candidates(
 
             let subdir = entry.path();
             // we skip commonly used subdirs that should not be searched for recursively
-            if !(subdir.ends_with("tests") || subdir.ends_with("test") || subdir.ends_with("demo"))
-            {
-                // scan the subdirectory for remappings, but we need a way to identify nested
-                // dependencies like `ds-token/lib/ds-stop/lib/ds-note/src/contract.sol`, or
-                // `oz/{tokens,auth}/{contracts, interfaces}/contract.sol` to assign
-                // the remappings to their root, we use a window that lies between two barriers. If
-                // we find a solidity file within a window, it belongs to the dir that opened the
-                // window.
-
-                // check if the subdir is a lib barrier, in which case we open a new window
-                if is_lib_dir(subdir) {
-                    candidates.extend(find_remapping_candidates(
-                        subdir,
-                        subdir,
-                        current_level + 1,
-                        is_inside_node_modules,
-                        visited_symlink_dirs,
-                    ));
-                } else {
-                    // continue scanning with the current window
-                    candidates.extend(find_remapping_candidates(
-                        subdir,
-                        open,
-                        current_level,
-                        is_inside_node_modules,
-                        visited_symlink_dirs,
-                    ));
-                }
+            if !no_recurse(subdir) {
+                search.push(subdir.to_path_buf());
             }
         }
     }
 
+    // all found candidates
+    let mut candidates = search
+        .par_iter()
+        .flat_map_iter(|subdir| {
+            // scan the subdirectory for remappings, but we need a way to identify nested
+            // dependencies like `ds-token/lib/ds-stop/lib/ds-note/src/contract.sol`, or
+            // `oz/{tokens,auth}/{contracts, interfaces}/contract.sol` to assign
+            // the remappings to their root, we use a window that lies between two barriers. If
+            // we find a solidity file within a window, it belongs to the dir that opened the
+            // window.
+
+            // check if the subdir is a lib barrier, in which case we open a new window
+            if is_lib_dir(subdir) {
+                find_remapping_candidates(
+                    subdir,
+                    subdir,
+                    current_level + 1,
+                    is_inside_node_modules,
+                    visited_symlink_dirs,
+                )
+            } else {
+                // continue scanning with the current window
+                find_remapping_candidates(
+                    subdir,
+                    open,
+                    current_level,
+                    is_inside_node_modules,
+                    visited_symlink_dirs,
+                )
+            }
+        })
+        .collect::<Vec<_>>();
+
     // need to find the actual next window in the event `open` is a lib dir
     let window_start = next_nested_window(open, current_dir);
     // finally, we need to merge, adjust candidates from the same level and open window
@@ -425,6 +416,21 @@ fn find_remapping_candidates(
     candidates
 }
 
+fn read_dir(dir: &Path) -> impl Iterator<Item = walkdir::DirEntry> {
+    walkdir::WalkDir::new(dir)
+        .sort_by_file_name()
+        .follow_links(true)
+        .min_depth(1)
+        .max_depth(1)
+        .into_iter()
+        .filter_entry(|e| !is_hidden(e))
+        .filter_map(Result::ok)
+}
+
+fn no_recurse(dir: &Path) -> bool {
+    dir.ends_with("tests") || dir.ends_with("test") || dir.ends_with("demo")
+}
+
 /// Counts the number of components between `root` and `current`
 /// `dir_distance("root/a", "root/a/b/c") == 2`
 fn dir_distance(root: &Path, current: &Path) -> usize {