rustc_metadata: Preprocess search paths for better performance

Over in Zed we've noticed that loading crates for a large-ish workspace can take almost 200ms. We've pinned it down to how rustc searches for paths, as it performs a linear search over the list of candidate paths. In our case the candidate list had about 20k entries which we had to iterate over for each dependency being loaded.

This commit introduces a simple FilesIndex that's just a BTreeMap under the hood. Since crates are looked up by both prefix and suffix, we perform a range search on said BTree (which constraints the search space based on prefix) and follow up with a linear scan of entries with matching suffixes.

Overall, this commit brings down build time for us in dev scenarios by about 6%.
100ms might not seem like much, but this is a constant cost that each of our workspace crates has to pay, even when said crate is miniscule.

Author: osiewicz

compiler/rustc_metadata/src/locator.rs

@@ -247,7 +247,7 @@ pub(crate) struct CrateLocator<'a> {
     cfg_version: &'static str,
 
     // Immutable per-search configuration.
-    crate_name: Symbol,
+    pub crate_name: Symbol,
     exact_paths: Vec<CanonicalizedPath>,
     pub hash: Option<Svh>,
     extra_filename: Option<&'a str>,
@@ -409,45 +409,40 @@ impl<'a> CrateLocator<'a> {
         // handled safely in `extract_one`.
         for search_path in self.filesearch.search_paths() {
             debug!("searching {}", search_path.dir.display());
-            for spf in search_path.files.iter() {
-                debug!("testing {}", spf.path.display());
-
-                let f = &spf.file_name_str;
-                let (hash, kind) = if let Some(f) = f.strip_prefix(rlib_prefix)
-                    && let Some(f) = f.strip_suffix(rlib_suffix)
-                {
-                    (f, CrateFlavor::Rlib)
-                } else if let Some(f) = f.strip_prefix(rmeta_prefix)
-                    && let Some(f) = f.strip_suffix(rmeta_suffix)
-                {
-                    (f, CrateFlavor::Rmeta)
-                } else if let Some(f) = f.strip_prefix(dylib_prefix)
-                    && let Some(f) = f.strip_suffix(dylib_suffix.as_ref())
-                {
-                    (f, CrateFlavor::Dylib)
-                } else {
-                    if f.starts_with(staticlib_prefix) && f.ends_with(staticlib_suffix.as_ref()) {
-                        self.crate_rejections.via_kind.push(CrateMismatch {
-                            path: spf.path.clone(),
-                            got: "static".to_string(),
-                        });
-                    }
-                    continue;
-                };
+            let spf = &search_path.files;
 
-                info!("lib candidate: {}", spf.path.display());
+            for (prefix, suffix, kind) in [
+                (rlib_prefix.as_str(), rlib_suffix, CrateFlavor::Rlib),
+                (rmeta_prefix.as_str(), rmeta_suffix, CrateFlavor::Rmeta),
+                (dylib_prefix, dylib_suffix, CrateFlavor::Dylib),
+            ] {
+                if let Some(matches) = spf.query(prefix, suffix) {
+                    for (hash, spf) in matches {
+                        info!("lib candidate: {}", spf.path.display());
 
-                let (rlibs, rmetas, dylibs) = candidates.entry(hash.to_string()).or_default();
-                let path = try_canonicalize(&spf.path).unwrap_or_else(|_| spf.path.clone());
-                if seen_paths.contains(&path) {
-                    continue;
-                };
-                seen_paths.insert(path.clone());
-                match kind {
-                    CrateFlavor::Rlib => rlibs.insert(path, search_path.kind),
-                    CrateFlavor::Rmeta => rmetas.insert(path, search_path.kind),
-                    CrateFlavor::Dylib => dylibs.insert(path, search_path.kind),
-                };
+                        let (rlibs, rmetas, dylibs) =
+                            candidates.entry(hash.to_string()).or_default();
+                        let path =
+                            try_canonicalize(&spf.path).unwrap_or_else(|_| spf.path.to_path_buf());
+                        if seen_paths.contains(&path) {
+                            continue;
+                        };
+                        seen_paths.insert(path.clone());
+                        match kind {
+                            CrateFlavor::Rlib => rlibs.insert(path, search_path.kind),
+                            CrateFlavor::Rmeta => rmetas.insert(path, search_path.kind),
+                            CrateFlavor::Dylib => dylibs.insert(path, search_path.kind),
+                        };
+                    }
+                }
+            }
+            if let Some(static_matches) = spf.query(staticlib_prefix, staticlib_suffix) {
+                for (_, spf) in static_matches {
+                    self.crate_rejections.via_kind.push(CrateMismatch {
+                        path: spf.path.to_path_buf(),
+                        got: "static".to_string(),
+                    });
+                }
             }
         }
 

compiler/rustc_session/src/search_paths.rs

@@ -1,4 +1,7 @@
+use std::collections::BTreeMap;
+use std::ops::Bound;
 use std::path::{Path, PathBuf};
+use std::sync::Arc;
 
 use rustc_macros::{Decodable, Encodable, HashStable_Generic};
 use rustc_target::spec::TargetTuple;
@@ -10,9 +13,40 @@ use crate::filesearch::make_target_lib_path;
 pub struct SearchPath {
     pub kind: PathKind,
     pub dir: PathBuf,
-    pub files: Vec<SearchPathFile>,
+    pub files: FilesIndex,
 }
 
+/// [FilesIndex] contains paths that can be efficiently looked up with (prefix, suffix) pairs.
+#[derive(Clone, Debug)]
+pub struct FilesIndex(BTreeMap<Arc<str>, Arc<SearchPathFile>>);
+
+impl FilesIndex {
+    /// Look up [SearchPathFile] by (prefix, suffix) pair.
+    pub fn query<'this, 'prefix, 'suffix>(
+        &'this self,
+        prefix: &'prefix str,
+        suffix: &'suffix str,
+    ) -> Option<impl Iterator<Item = (String, Arc<SearchPathFile>)> + use<'this, 'prefix, 'suffix>>
+    {
+        let start = prefix;
+        let end = format!("{prefix}\u{FFFF}");
+        let mut ret = self
+            .0
+            .range::<str, _>((Bound::Included(start), Bound::Excluded(end.as_str())))
+            .filter_map(move |(k, v)| {
+                k.ends_with(suffix).then(|| {
+                    (
+                        String::from(
+                            &v.file_name_str[prefix.len()..v.file_name_str.len() - suffix.len()],
+                        ),
+                        v.clone(),
+                    )
+                })
+            })
+            .peekable();
+        ret.peek().is_some().then(|| ret)
+    }
+}
 /// The obvious implementation of `SearchPath::files` is a `Vec<PathBuf>`. But
 /// it is searched repeatedly by `find_library_crate`, and the searches involve
 /// checking the prefix and suffix of the filename of each `PathBuf`. This is
@@ -26,8 +60,8 @@ pub struct SearchPath {
 /// UTF-8, and so a non-UTF-8 filename couldn't be one we're looking for.)
 #[derive(Clone, Debug)]
 pub struct SearchPathFile {
-    pub path: PathBuf,
-    pub file_name_str: String,
+    pub path: Arc<Path>,
+    pub file_name_str: Arc<str>,
 }
 
 #[derive(PartialEq, Clone, Copy, Debug, Hash, Eq, Encodable, Decodable, HashStable_Generic)]
@@ -102,16 +136,20 @@ impl SearchPath {
             Ok(files) => files
                 .filter_map(|e| {
                     e.ok().and_then(|e| {
-                        e.file_name().to_str().map(|s| SearchPathFile {
-                            path: e.path(),
-                            file_name_str: s.to_string(),
+                        e.file_name().to_str().map(|s| {
+                            let file_name_str: Arc<str> = s.into();
+                            (
+                                file_name_str.clone(),
+                                Arc::new(SearchPathFile { path: e.path().into(), file_name_str }),
+                            )
                         })
                     })
                 })
-                .collect::<Vec<_>>(),
-            Err(..) => vec![],
-        };
+                .collect::<BTreeMap<_, _>>(),
 
+            Err(..) => Default::default(),
+        };
+        let files = FilesIndex(files);
         SearchPath { kind, dir, files }
     }
 }