defer suggested stdlib imports to diagnostic display & use global cache

Summary:
strsim::levenshtein was using 12% of our CPU time from profiling initial indexing on langchain/dify

when we fail to find a module, we check the distance between the module name & all the stdlib modules, and compute a list of suggestions (did you mean...?) for the error message.

this computation is very expensive and uncached, and mostly unnecessary since we don't emit diagnostics for most files during indexing.

this diff makes it so that we only compute the suggestions when we display an error, and use a global cache for the suggestions.

Reviewed By: rchen152

Differential Revision: D91750918

fbshipit-source-id: 1e88754c540ad79fcfc4f8bf84f5c13c8078cd8f

Author: yangdanny97

pyrefly/lib/module/finder.rs

@@ -8,16 +8,23 @@
 use std::iter;
 use std::path::Path;
 use std::path::PathBuf;
+use std::sync::LazyLock;
 
 use pyrefly_python::COMPILED_FILE_SUFFIXES;
 use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::module_path::ModulePath;
 use pyrefly_python::module_path::ModuleStyle;
+use pyrefly_util::locked_map::LockedMap;
 use pyrefly_util::suggest::best_suggestion;
 use ruff_python_ast::name::Name;
 use starlark_map::small_map::SmallMap;
 use vec1::Vec1;
 
+/// Global cache for stdlib import suggestions.
+/// Keyed by the missing module name, returns the suggested module name (if any).
+static STDLIB_SUGGESTION_CACHE: LazyLock<LockedMap<ModuleName, Option<ModuleName>>> =
+    LazyLock::new(LockedMap::new);
+
 use crate::config::config::ConfigFile;
 use crate::module::bundled::BundledStub;
 use crate::module::third_party::get_bundled_third_party;
@@ -641,7 +648,6 @@ pub fn find_import_filtered(
             config.structured_import_lookup_path(origin),
             module,
             &config.source,
-            suggest_stdlib_import(module),
         ))
     }
 }
@@ -698,7 +704,14 @@ fn recommended_stubs_package(module: ModuleName) -> Option<ModuleName> {
 
 /// Suggest a similar stdlib module name for a mistyped import.
 /// Uses Levenshtein distance to find the closest match from typeshed's stdlib modules.
-fn suggest_stdlib_import(missing: ModuleName) -> Option<ModuleName> {
+/// Results are cached globally since typeshed doesn't change during a session.
+pub fn suggest_stdlib_import(missing: ModuleName) -> Option<ModuleName> {
+    *STDLIB_SUGGESTION_CACHE
+        .ensure(&missing, || suggest_stdlib_import_uncached(missing))
+        .0
+}
+
+fn suggest_stdlib_import_uncached(missing: ModuleName) -> Option<ModuleName> {
     let ts = typeshed().ok()?;
     let missing_str = missing.as_str();
 
@@ -1054,7 +1067,6 @@ mod tests {
                 config.structured_import_lookup_path(None),
                 ModuleName::from_str("spp_priority.d"),
                 &config.source,
-                None,
             )),
         );
     }

pyrefly/lib/state/loader.rs

@@ -24,6 +24,7 @@ use crate::config::config::ImportLookupPathPart;
 use crate::error::context::ErrorContext;
 use crate::module::finder::find_import;
 use crate::module::finder::find_import_filtered;
+use crate::module::finder::suggest_stdlib_import;
 
 #[derive(Debug, Clone, Dupe, PartialEq, Eq)]
 pub enum FindError {
@@ -51,7 +52,6 @@ impl FindError {
         path: Vec<ImportLookupPathPart>,
         module: ModuleName,
         config_source: &ConfigSource,
-        suggestion: Option<ModuleName>,
     ) -> FindError {
         let config_suffix = match config_source {
             ConfigSource::File(p) => format!(" (from config in `{}`)", p.display()),
@@ -79,18 +79,22 @@ impl FindError {
             format!("Looked in these locations{config_suffix}:")
         }];
         explanation.extend(nonempty_paths);
-        if let Some(suggested) = suggestion {
-            explanation.insert(0, format!("Did you mean `{suggested}`?"));
-        }
         FindError::NotFound(module, Arc::new(explanation))
     }
 
     pub fn display(&self) -> (Option<Box<dyn Fn() -> ErrorContext + '_>>, Vec1<String>) {
         match self {
-            Self::NotFound(module, err) => (
-                Some(Box::new(|| ErrorContext::ImportNotFound(*module))),
-                (**err).clone(),
-            ),
+            Self::NotFound(module, err) => {
+                let mut lines = (**err).clone();
+                // Compute suggestion lazily at display time, using global cache
+                if let Some(suggested) = suggest_stdlib_import(*module) {
+                    lines.insert(0, format!("Did you mean `{suggested}`?"));
+                }
+                (
+                    Some(Box::new(|| ErrorContext::ImportNotFound(*module))),
+                    lines,
+                )
+            }
             Self::Ignored => (None, vec1!["Ignored import".to_owned()]),
             Self::NoSource(module) => (
                 None,