fix: modify @[suggest_for] to work with the Prelude (#11529)

This PR fixes a syntax-pattern-matching issue from #11367 that prevented
the addition of suggestions in Init prior to Lean.Parser being
introduced, which was a significant shortcoming. It preserves the
ability to have multiple suggestions for one annotation later in the
process.

Additionally, tweaks a (not-yet-user-visible) error message and modifies
the attribute declaration to store a wrongIdentifier ->
correctIdentifier mapping instead of a correctIdentifier ->
wrongIdentifier mapping.

Author: robsimmons

src/Lean/Elab/App.lean

@@ -1420,15 +1420,26 @@ where
     -- Possible alternatives provided with `@[suggest_for]` annotations
     let suggestions := (← Lean.getSuggestions fullName).filter (·.getPrefix = fullName.getPrefix)
     let suggestForHint ←
-      if suggestions.size = 0 then
+      if h : suggestions.size = 0 then
         pure .nil
+      else if suggestions.size = 1 then
+        MessageData.hint (ref? := ref)
+          m!"Perhaps you meant `{.ofConstName suggestions[0]}` in place of `{fullName}`:"
+          (suggestions.map fun suggestion => {
+            preInfo? := .some s!"{e}.",
+            suggestion := suggestion.getString!,
+            toCodeActionTitle? := .some (s!"Suggested replacement: {e}.{·}"),
+            diffGranularity := .all,
+          })
       else
-        m!"Perhaps you meant one of these in place of `{fullName}`:".hint (suggestions.map fun suggestion => {
-          suggestion := suggestion.getString!,
-          toCodeActionTitle? := .some (s!"Suggested replacement: {e}.{·}"),
-          diffGranularity := .all,
-          messageData? := .some m!"`{.ofConstName suggestion}`: {e}.{suggestion.getString!}",
-        }) ref
+        MessageData.hint (ref? := ref)
+          m!"Perhaps you meant one of these in place of `{fullName}`:"
+          (suggestions.map fun suggestion => {
+            suggestion := suggestion.getString!,
+            toCodeActionTitle? := .some (s!"Suggested replacement: {e}.{·}"),
+            diffGranularity := .all,
+            messageData? := .some m!"`{.ofConstName suggestion}`: {e}.{suggestion.getString!}",
+          })
 
     -- By using `mkUnknownIdentifierMessage`, the tag `Lean.unknownIdentifierMessageTag` is
     -- incorporated within the message, as required for the "import unknown identifier" code action.

src/Lean/IdentifierSuggestion.lean

@@ -17,35 +17,46 @@ namespace Lean
 
 set_option doc.verso true
 
-builtin_initialize identifierSuggestionForAttr : ParametricAttribute (Name × Array Name) ←
-  registerParametricAttribute {
+
+builtin_initialize identifierSuggestionForAttr : PersistentEnvExtension
+    (Name × Array Name) /- *Store* mappings from incorrect names to corrections (identifiers that exist) -/
+    (Name × Array Name) /- *Add* mappings from existing functions to possible incorrect namings -/
+    (NameMap NameSet) /- *Use* mapping from incorrect names to corrections (identifiers that exist) -/ ←
+  let ext ← registerPersistentEnvExtension {
+    mkInitial := pure {},
+    addImportedFn := fun modules => pure <|
+      modules.foldl (init := {}) fun accum entries =>
+        entries.foldl (init := accum) fun accum (altName, trueNames) =>
+          accum.alter altName (fun old => trueNames.foldl (β := NameSet) (init := old.getD ∅) fun accum trueName =>
+            accum.insert trueName),
+    addEntryFn := fun table (trueName, altNames) =>
+      altNames.foldl (init := table) fun accum alt =>
+        accum.alter alt (·.getD {} |>.insert trueName)
+    exportEntriesFn table := table.toArray.map fun (altName, trueNames) =>
+      (altName, trueNames.toArray)
+  }
+
+  registerBuiltinAttribute {
     name := `suggest_for,
     descr := "suggest other (incorrect, not-existing) identifiers that someone might use when they actually want this definition",
-    getParam := fun trueDeclName stx => do
-      let `(attr| suggest_for $altNames:ident*) := stx
-        | throwError "Invalid `[suggest_for]` attribute syntax"
-      let ns := trueDeclName.getPrefix
-      return (trueDeclName, altNames.map (·.getId))
-    }
+    add (decl : Name) (stx : Syntax) (kind : AttributeKind) := do
+      unless kind == AttributeKind.global do throwAttrMustBeGlobal `suggest_for kind
+      let altSyntaxIds : Array Syntax ← match stx with
+        | -- Attributes parsed _after_ the suggest_for notation is added
+          .node _ ``suggest_for #[.atom _ "suggest_for", .node _ `null ids] => pure ids
+        | -- Attributes parsed _before the suggest_for notation is added
+          .node _ ``Parser.Attr.simple #[.ident _ _ `suggest_for [], .node _ `null #[id]] => pure #[id]
+        | _ => throwError "Invalid `[suggest_for]` attribute syntax {repr stx}"
+      modifyEnv (ext.addEntry · (decl, altSyntaxIds.map (·.getId)))
+  }
+
+  return ext
 
 public def getSuggestions [Monad m] [MonadEnv m] (fullName : Name) : m (Array Name) := do
-  let mut possibleReplacements := #[]
-  let (_, allSuggestions) := identifierSuggestionForAttr.ext.getState (← getEnv)
-  for (_, trueName, suggestions) in allSuggestions do
-    for suggestion in suggestions do
-      if fullName = suggestion then
-        possibleReplacements := possibleReplacements.push trueName
-  return possibleReplacements.qsort (lt := lt)
-where
-  -- Ensure the result of getSuggestions is stable (if arbitrary)
-  lt : Name -> Name -> Bool
-    | .anonymous, _ => false
-    | .str _ _, .anonymous => true
-    | .num _ _, .anonymous => true
-    | .str _ _, .num _ _ => true
-    | .num _ _, .str _ _ => false
-    | .num a n, .num b m => n < m || n == m && lt a b
-    | .str a s1, .str b s2 => s1 < s2 || s1 == s2 && lt a b
+  return identifierSuggestionForAttr.getState (← getEnv)
+    |>.find? fullName
+    |>.getD {}
+    |>.toArray
 
 /--
 Throw an unknown constant error message, potentially suggesting alternatives using

tests/lean/run/identifierSuggestions.lean

@@ -22,8 +22,8 @@ error: Invalid field `test0`: The environment does not contain `String.test0`, s
   "abc"
 of type `String`
 
-Hint: Perhaps you meant one of these in place of `String.test0`:
-  [apply] `String.foo`: "abc".foo
+Hint: Perhaps you meant `String.foo` in place of `String.test0`:
+  "abc".t̵e̵s̵t̵0̵f̲o̲o̲
 -/
 #guard_msgs in
 #check "abc".test0
@@ -44,8 +44,8 @@ error: Invalid field `test1`: The environment does not contain `String.test1`, s
 of type `String`
 
 Hint: Perhaps you meant one of these in place of `String.test1`:
-  [apply] `String.baz`: "abc".baz
   [apply] `String.foo`: "abc".foo
+  [apply] `String.baz`: "abc".baz
 -/
 #guard_msgs in
 #check "abc".test1
@@ -54,8 +54,8 @@ Hint: Perhaps you meant one of these in place of `String.test1`:
 error: Unknown constant `String.test1`
 
 Hint: Perhaps you meant one of these in place of `String.test1`:
-  • S̵t̵r̵i̵n̵g̵.̵t̵e̵s̵t̵1̵S̲t̲r̲i̲n̲g̲.̲b̲a̲z̲
   • S̵t̵r̵i̵n̵g̵.̵t̵e̵s̵t̵1̵S̲t̲r̲i̲n̲g̲.̲f̲o̲o̲
+  • S̵t̵r̵i̵n̵g̵.̵t̵e̵s̵t̵1̵S̲t̲r̲i̲n̲g̲.̲b̲a̲z̲
 -/
 #guard_msgs in
 #check String.test1
@@ -67,9 +67,9 @@ error: Invalid field `test2`: The environment does not contain `String.test2`, s
 of type `String`
 
 Hint: Perhaps you meant one of these in place of `String.test2`:
-  [apply] `String.bar`: "abc".bar
-  [apply] `String.baz`: "abc".baz
   [apply] `String.foo`: "abc".foo
+  [apply] `String.baz`: "abc".baz
+  [apply] `String.bar`: "abc".bar
 -/
 #guard_msgs in
 #check "abc".test2
@@ -79,11 +79,11 @@ Hint: Perhaps you meant one of these in place of `String.test2`:
 error: Unknown constant `String.test2`
 
 Hint: Perhaps you meant one of these in place of `String.test2`:
-  • S̵t̵r̵i̵n̵g̵.̵t̵e̵s̵t̵2̵S̲t̲r̲i̲n̲g̲.̲b̲a̲r̲
-  • S̵t̵r̵i̵n̵g̵.̵t̵e̵s̵t̵2̵S̲t̲r̲i̲n̲g̲.̲b̲a̲z̲
-  • S̵t̵r̵i̵n̵g̵.̵t̵e̵s̵t̵2̵S̲t̲r̲i̲n̲g̲.̲f̲o̲o̲
   • S̵t̵r̵i̵n̵g̵.̵t̵e̵s̵t̵2̵o̲t̲h̲e̲r̲B̲a̲z̲
+  • S̵t̵r̵i̵n̵g̵.̵t̵e̵s̵t̵2̵S̲t̲r̲i̲n̲g̲.̲f̲o̲o̲
   • S̵t̵r̵i̵n̵g̵.̵t̵e̵s̵t̵2̵o̲t̲h̲e̲r̲F̲o̲o̲
+  • S̵t̵r̵i̵n̵g̵.̵t̵e̵s̵t̵2̵S̲t̲r̲i̲n̲g̲.̲b̲a̲z̲
+  • S̵t̵r̵i̵n̵g̵.̵t̵e̵s̵t̵2̵S̲t̲r̲i̲n̲g̲.̲b̲a̲r̲
 -/
 #guard_msgs in
 #check String.test2
@@ -118,8 +118,8 @@ error: Invalid field `toNum`: The environment does not contain `Foo.Bar.toNum`,
   Foo.Bar.three
 of type `Foo.Bar`
 
-Hint: Perhaps you meant one of these in place of `Foo.Bar.toNum`:
-  [apply] `Foo.Bar.toNat`: Foo.Bar.three.toNat
+Hint: Perhaps you meant `Foo.Bar.toNat` in place of `Foo.Bar.toNum`:
+  Foo.Bar.three.t̵o̵N̵u̵m̵t̲o̲N̲a̲t̲
 -/
 #guard_msgs in
 #eval Foo.Bar.three.toNum
@@ -129,8 +129,8 @@ error: Invalid field `toStr`: The environment does not contain `Foo.Bar.toStr`,
   Foo.Bar.two
 of type `Foo.Bar`
 
-Hint: Perhaps you meant one of these in place of `Foo.Bar.toStr`:
-  [apply] `Foo.Bar.toString`: Foo.Bar.two.toString
+Hint: Perhaps you meant `Foo.Bar.toString` in place of `Foo.Bar.toStr`:
+  Foo.Bar.two.t̵o̵S̵t̵r̵t̲o̲S̲t̲r̲i̲n̲g̲
 -/
 #guard_msgs in
 #eval Foo.Bar.two.toStr
@@ -273,8 +273,8 @@ error: Invalid field `not`: The environment does not contain `MyBool.not`, so it
   MyBool.tt
 of type `MyBool`
 
-Hint: Perhaps you meant one of these in place of `MyBool.not`:
-  [apply] `MyBool.swap`: MyBool.tt.swap
+Hint: Perhaps you meant `MyBool.swap` in place of `MyBool.not`:
+  MyBool.tt.n̵o̵t̵s̲w̲a̲p̲
 -/
 #guard_msgs in
 example := MyBool.tt.not