Use Vec as the underlying store for FilesIndex

Author: osiewicz

compiler/rustc_session/src/search_paths.rs

@@ -1,5 +1,3 @@
-use std::collections::BTreeMap;
-use std::ops::Bound;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
@@ -18,36 +16,45 @@ pub struct SearchPath {
 
 /// [FilesIndex] contains paths that can be efficiently looked up with (prefix, suffix) pairs.
 #[derive(Clone, Debug)]
-pub struct FilesIndex(BTreeMap<Arc<str>, Arc<SearchPathFile>>);
+pub struct FilesIndex(Vec<(Arc<str>, 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>>
+    ) -> Option<impl Iterator<Item = (String, &'this 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()],
-                        ),
-                        Arc::clone(v),
-                    )
-                })
+        let start = self.0.partition_point(|(k, _)| {
+            let k: &str = &k;
+            k < prefix
+        });
+        if start == self.0.len() {
+            return None;
+        }
+        let prefixed_items;
+        let end = self.0[start..].partition_point(|(k, _)| k.starts_with(prefix));
+        if end + start == self.0.len() {
+            prefixed_items = &self.0[start..];
+        } else {
+            prefixed_items = &self.0[start..start + end];
+        }
+
+        let ret = prefixed_items.into_iter().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,
+                )
             })
-            .peekable();
-        ret.peek().is_some().then(|| ret)
+        });
+        Some(ret)
     }
     pub fn retain(&mut self, prefixes: &[&str]) {
-        self.0.retain(|k, _| prefixes.iter().any(|prefix| k.starts_with(prefix)));
+        self.0.retain(|(k, _)| prefixes.iter().any(|prefix| k.starts_with(prefix)));
     }
 }
 /// The obvious implementation of `SearchPath::files` is a `Vec<PathBuf>`. But
@@ -135,23 +142,24 @@ impl SearchPath {
 
     pub fn new(kind: PathKind, dir: PathBuf) -> Self {
         // Get the files within the directory.
-        let files = match std::fs::read_dir(&dir) {
+        let mut files = match std::fs::read_dir(&dir) {
             Ok(files) => files
                 .filter_map(|e| {
                     e.ok().and_then(|e| {
                         e.file_name().to_str().map(|s| {
                             let file_name_str: Arc<str> = s.into();
-                            (
-                                Arc::clone(&file_name_str),
-                                Arc::new(SearchPathFile { path: e.path().into(), file_name_str }),
-                            )
+                            (Arc::clone(&file_name_str), SearchPathFile {
+                                path: e.path().into(),
+                                file_name_str,
+                            })
                         })
                     })
                 })
-                .collect::<BTreeMap<_, _>>(),
+                .collect::<Vec<_>>(),
 
             Err(..) => Default::default(),
         };
+        files.sort_by(|(lhs, _), (rhs, _)| lhs.cmp(rhs));
         let files = FilesIndex(files);
         SearchPath { kind, dir, files }
     }