fix: better performance for @[suggest_for] (#11598)

This PR fixes a performance issue resulting from misusing persistent
environment extensions that was introduced in #11554.

src/Init/Notation.lean

@@ -634,7 +634,7 @@ syntax (name := deprecated) "deprecated" (ppSpace ident)? (ppSpace str)?
     (" (" &"since" " := " str ")")? : attr
 
 /--
-The attribute `@[suggest_for]` on a declaration suggests likely ways in which
+The attribute `@[suggest_for ..]` on a declaration suggests likely ways in which
 someone might **incorrectly** refer to a definition.
 
 * `@[suggest_for String.endPos]` on the definition of `String.rawEndPos` suggests that `"str".endPos` might be correctable to `"str".rawEndPos`.
@@ -644,6 +644,10 @@ The namespace of the suggestions is always relative to the root namespace. In th
 adding an annotation `@[suggest_for Z.bar]` to `def Z.foo` will suggest `X.Y.Z.foo` only as a
 replacement for `Z.foo`. If your intent is to suggest `X.Y.Z.foo` as a replacement for
 `X.Y.Z.bar`, you must instead use the annotation `@[suggest_for X.Y.Z.bar]`.
+
+Suggestions can be defined for structure fields or inductive branches with the
+`attribute [suggest_for Exception] Except` syntax, and these attributes do not have to be added
+in the same module where the actual identifier was defined.
 -/
 syntax (name := suggest_for) "suggest_for" (ppSpace ident)+ : attr
 

src/Lean/IdentifierSuggestion.lean

@@ -17,36 +17,48 @@ namespace Lean
 
 set_option doc.verso true
 
-public structure IdentSuggestion where
-  existingToIncorrect : NameMap NameSet := {}
-  incorrectToExisting : NameMap NameSet := {}
-deriving Inhabited
+public abbrev NameMapExtension := PersistentEnvExtension (Name × Array Name) (Name × Array Name) (NameMap NameSet)
 
-def IdentSuggestion.add (table : IdentSuggestion) (trueName : Name) (altNames : Array Name) : IdentSuggestion :=
-  { existingToIncorrect :=
-      table.existingToIncorrect.alter trueName fun old =>
-        altNames.foldl (β := NameSet) (init := old.getD {}) fun accum altName =>
-          accum.insert altName
-    incorrectToExisting :=
-      altNames.foldl (init := table.incorrectToExisting) fun accum altName =>
-        accum.alter altName fun old =>
-          old.getD {} |>.insert trueName
-  }
+/--
+Create the extension mapping from existing identifiers to the incorrect alternatives for which we
+want to provide suggestions. This is mostly equivalent to a {name}`MapDeclarationExtension` or the
+extensions underlying {name}`ParametricAttribute` attributes, but it differs in allowing
+{name}`suggest_for` attributes to be assigned in files other than the ones where they were defined.
+-/
+def mkExistingToIncorrect : IO NameMapExtension := registerPersistentEnvExtension {
+  name := `Lean.identifierSuggestForAttr.existingToIncorrect
+  mkInitial := pure {},
+  addImportedFn := fun _ => pure {},
+  addEntryFn := fun table (trueName, altNames) =>
+    table.alter trueName fun old =>
+      altNames.foldl (β := NameSet) (init := old.getD {}) fun accum altName =>
+        accum.insert altName
+  exportEntriesFn table :=
+    table.toArray.map (fun (trueName, altNames) =>(trueName, altNames.toArray))
+      |>.qsort (lt := fun a b => Name.quickLt a.1 b.1)
+}
 
-builtin_initialize identifierSuggestionForAttr : PersistentEnvExtension
-    (Name × Array Name) /- Mapping from existing constant to potential incorrect alternative names -/
-    (Name × Array Name) /- Mapping from existing constant to potential incorrect alternative names -/
-    IdentSuggestion /- Two directional mapping between existing identifier and alternative incorrect mappings -/ ←
-  let ext ← registerPersistentEnvExtension {
-    mkInitial := pure {},
-    addImportedFn := fun modules => pure <|
-      (modules.foldl (init := {}) fun accum entries =>
-        entries.foldl (init := accum) fun accum (trueName, altNames) =>
-          accum.add trueName altNames),
-    addEntryFn := fun table (trueName, altNames) => table.add trueName altNames
-    exportEntriesFn table := table.existingToIncorrect.toArray.map fun (trueName, altNames) =>
-      (trueName, altNames.toArray)
-  }
+/--
+Create the extension mapping incorrect identifiers to the existing identifiers we want to suggest as
+replacements. This association is the opposite of the usual mapping for {name}`ParametricAttribute`
+attributes.
+-/
+def mkIncorrectToExisting : IO NameMapExtension := registerPersistentEnvExtension {
+  name := `Lean.identifierSuggestForAttr.incorrectToExisting
+  mkInitial := pure {},
+  addImportedFn := fun _ => pure {},
+  addEntryFn := fun table (trueName, altNames) =>
+    altNames.foldl (init := table) fun accum altName =>
+      accum.alter altName fun old =>
+        old.getD {} |>.insert trueName
+  exportEntriesFn table :=
+    table.toArray.map (fun (altName, trueNames) => (altName, trueNames.toArray))
+      |>.qsort (lt := fun a b => Name.quickLt a.1 b.1)
+}
+
+builtin_initialize identifierSuggestionsImpl : NameMapExtension × NameMapExtension ←
+  let existingToIncorrect ← mkExistingToIncorrect
+  let incorrectToExisting ← mkIncorrectToExisting
 
   registerBuiltinAttribute {
     name := `suggest_for,
@@ -61,36 +73,46 @@ builtin_initialize identifierSuggestionForAttr : PersistentEnvExtension
         | _ => throwError "Invalid `[suggest_for]` attribute syntax {repr stx}"
       if isPrivateName decl then
         throwError "Cannot make suggestions for private names"
-      modifyEnv (ext.addEntry · (decl, altSyntaxIds.map (·.getId)))
+      let altIds := altSyntaxIds.map (·.getId)
+      modifyEnv (existingToIncorrect.addEntry · (decl, altIds))
+      modifyEnv (incorrectToExisting.addEntry · (decl, altIds))
   }
 
-  return ext
+  return (existingToIncorrect, incorrectToExisting)
 
 /--
 Given a name, find all the stored correct, existing identifiers that mention that name in a
-{lit}`suggest_for` annotation.
+{name}`suggest_for` annotation.
 -/
 public def getSuggestions [Monad m] [MonadEnv m] (incorrectName : Name) : m NameSet := do
-  return identifierSuggestionForAttr.getState (← getEnv)
-    |>.incorrectToExisting
-    |>.find? incorrectName
-    |>.getD {}
+  let env ← getEnv
+  let { state, importedEntries } := identifierSuggestionsImpl.2.toEnvExtension.getState env
+  let localEntries := state.find? incorrectName |>.getD {}
+  return importedEntries.foldl (init := localEntries) fun results moduleEntry =>
+    match moduleEntry.binSearch (incorrectName, default) (fun a b => Name.quickLt a.1 b.1) with
+    | none => results
+    | some (_, extras) => extras.foldl (init := results) fun accum extra => accum.insert extra
 
 /--
 Given a (presumably existing) identifier, find all the {lit}`suggest_for` annotations that were given
 for that identifier.
 -/
 public def getStoredSuggestions [Monad m] [MonadEnv m] (trueName : Name) : m NameSet := do
-  return identifierSuggestionForAttr.getState (← getEnv)
-    |>.existingToIncorrect
-    |>.find? trueName
-    |>.getD {}
+  let env ← getEnv
+  let { state, importedEntries } := identifierSuggestionsImpl.1.toEnvExtension.getState env
+  let localEntries := state.find? trueName |>.getD {}
+  return importedEntries.foldl (init := localEntries) fun results moduleEntry =>
+    match moduleEntry.binSearch (trueName, default) (fun a b => Name.quickLt a.1 b.1) with
+    | none => results
+    | some (_, extras) => extras.foldl (init := results) fun accum extra => accum.insert extra
 
 /--
 Throw an unknown constant error message, potentially suggesting alternatives using
 {name}`suggest_for` attributes. (Like {name}`throwUnknownConstantAt` but with suggestions.)
 
-The "Unknown constant `<id>`" message will fully qualify the name, whereas the
+The replacement will mimic the path structure of the original as much as possible if they share a
+path prefix: if there is a suggestion for replacing `Foo.Bar.jazz` with `Foo.Bar.baz`, then
+`Bar.jazz` will be replaced by `Bar.baz` unless the resulting constant is ambiguous.
 -/
 public def throwUnknownConstantWithSuggestions (constName : Name) (ref? : Option Syntax := .none) : CoreM α := do
   let suggestions := (← getSuggestions constName).toArray

src/stdlib_flags.h

@@ -1,5 +1,6 @@
 #include "util/options.h"
 
+// this comment has been updated 1 time(s)
 namespace lean {
 options get_default_options() {
     options opts;

tests/lean/run/idSuggestPostHoc.lean

@@ -0,0 +1,24 @@
+-- test extending suggestions in a different module
+
+/--
+error: Invalid field `some`: The environment does not contain `String.some`, so it is not possible to project the field `some` from an expression
+  x
+of type `String`
+
+Hint: Perhaps you meant `String.contains` in place of `String.some`:
+  .s̵o̵m̵e̵c̲o̲n̲t̲a̲i̲n̲s̲
+-/
+#guard_msgs in example (x : String) := x.some fun _ => true
+
+attribute [suggest_for String.some] String.any
+
+/--
+error: Invalid field `some`: The environment does not contain `String.some`, so it is not possible to project the field `some` from an expression
+  x
+of type `String`
+
+Hint: Perhaps you meant one of these in place of `String.some`:
+  [apply] `String.contains`
+  [apply] `String.any`
+-/
+#guard_msgs in example (x : String) := x.some fun _ => true