feat: support code actions in linters

src/Lean/Elab/Command.lean

@@ -335,15 +335,17 @@ instance : MonadLog CommandElabM where
     let msg := { msg with data := MessageData.withNamingContext { currNamespace := currNamespace, openDecls := openDecls } msg.data }
     modify fun s => { s with messages := s.messages.add msg }
 
-def runLinters (stx : Syntax) : CommandElabM Unit := do
+def runLinters (stx : Syntax) (promise : Option (IO.Promise InfoTree) := .none) : CommandElabM Unit := do
   profileitM Exception "linting" (← getOptions) do
     withTraceNode `Elab.lint (fun _ => return m!"running linters") do
       let linters ← lintersRef.get
+      let producedInfoTrees ← IO.mkRef ({} : PersistentArray InfoTree)
       unless linters.isEmpty do
         for linter in linters do
           withTraceNode `Elab.lint (fun _ => return m!"running linter: {.ofConstName linter.name}")
               (tag := linter.name.toString) do
             let savedState ← get
+            let originalSize := savedState.infoState.trees.size
             try
               linter.run stx
             catch ex =>
@@ -355,7 +357,15 @@ def runLinters (stx : Syntax) : CommandElabM Unit := do
             finally
               -- TODO: it would be good to preserve even more state (#4363) but preserving info
               -- trees currently breaks from linters adding context-less info nodes
+              let newInfoState ← getInfoState
+              if newInfoState.enabled then
+                producedInfoTrees.modify fun old =>
+                  old.append (newInfoState.trees.foldl (·.push ·) {} (start := originalSize))
               modify fun s => { savedState with messages := s.messages, traceState := s.traceState }
+      if let some promise := promise then
+        if (← getInfoState).enabled then
+          promise.resolve <|
+            mkLinterInfoGroupNode (← producedInfoTrees.get)
 
 def runModuleLinters (cmds : Array Syntax) : CommandElabM Unit := do
   profileitM Exception "module linting" (← getOptions) do
@@ -377,16 +387,19 @@ def runModuleLinters (cmds : Array Syntax) : CommandElabM Unit := do
             finally
               modify fun s => { savedState with messages := s.messages, traceState := s.traceState }
 
-def runStatefulLinters (stx : Syntax) (prev : Array LinterState) : CommandElabM (Array LinterState) := do
+def runStatefulLinters (stx : Syntax) (prev : Array LinterState)
+    (promise : Option (IO.Promise InfoTree) := .none) : CommandElabM (Array LinterState) := do
   profileitM Exception "stateful linting" (← getOptions) do
     withTraceNode `Elab.lint (fun _ => return m!"running stateful linters") do
       let linters ← statefulLintersRef.get
+      let producedInfoTrees ← IO.mkRef ({} : PersistentArray InfoTree)
       let run {α : Type} (phase : String) (idx : Nat) (onError : CommandElabM α)
           (act : CommandElabM α) : CommandElabM α :=
         withTraceNode `Elab.lint
             (fun _ => return m!"running stateful linter #{idx} ({phase})")
             (tag := toString idx) do
           let savedState ← get
+          let originalSize := savedState.infoState.trees.size
           try
             act
           catch ex =>
@@ -396,6 +409,10 @@ def runStatefulLinters (stx : Syntax) (prev : Array LinterState) : CommandElabM
             | .internal _ _ => logException ex
             onError
           finally
+            let newInfoState ← getInfoState
+            if newInfoState.enabled then
+              producedInfoTrees.modify fun old =>
+                old.append (newInfoState.trees.foldl (·.push ·) {} (start := originalSize))
             modify fun s => { savedState with messages := s.messages, traceState := s.traceState }
       let mut preSt : Array (Option LinterState) := .emptyWithCapacity linters.size
       let mut i := 0
@@ -407,6 +424,10 @@ def runStatefulLinters (stx : Syntax) (prev : Array LinterState) : CommandElabM
       for l in linters do
         postSt := postSt.push (← run "post" i (pure prev[i]!) (l.post stx prev preSt))
         i := i + 1
+      if let some promise := promise then
+        if (← getInfoState).enabled then
+          promise.resolve <|
+            mkLinterInfoGroupNode (← producedInfoTrees.get)
       return postSt
 
 def initialLinterStates : BaseIO (Array LinterState) := do
@@ -493,6 +514,9 @@ def runLintersAsync (stx : Syntax) (cmds : Array Syntax) : CommandElabM Unit :=
         runModuleLinters cmds
     return
 
+  -- We create a promise for the info trees produced by the linters
+  let lintersInfoPromise ← IO.Promise.new (α := InfoTree)
+
   -- linters should have access to the complete info tree and message log
   let mut snaps := (← get).snapshotTasks
   if let some elabSnap := (← read).snap? then
@@ -509,15 +533,24 @@ def runLintersAsync (stx : Syntax) (cmds : Array Syntax) : CommandElabM Unit :=
     let messages := messages.markAllReported
     modify fun st => { st with messages := st.messages ++ messages }
     modifyInfoState fun _ => infoSt
-    runLinters stx
+    runLinters stx lintersInfoPromise
     if Parser.isTerminalCommand stx then
+        -- TODO: support code actions in module linters
+        -- Currently, code actions provided by terminal command are ignored
         runModuleLinters cmds
 
   let task ← BaseIO.bindTask (sync := true) (t := (← getInfoState).substituteLazy) fun infoSt =>
     BaseIO.mapTask (t := treeTask) fun _ =>
       lintAct infoSt
   logSnapshotTask { stx? := none, task, cancelTk? := cancelTk }
 
+  let infoHole ← liftCoreM mkFreshMVarId
+  modifyInfoState fun s => { s with
+    trees          := s.trees.modify 0 (pushInfoChild · (.hole infoHole))
+    lazyAssignment := s.lazyAssignment
+      |>.insert infoHole (lintersInfoPromise.resultD default)
+  }
+
 open Language in
 def runStatefulLintersAsync (stx : Syntax) : CommandElabM Unit := do
   if (← statefulLintersRef.get).isEmpty then return
@@ -539,13 +572,16 @@ def runStatefulLintersAsync (stx : Syntax) : CommandElabM Unit := do
   let inits ← initialLinterStates
   modify fun s => { s with prevLinterStates := some (statePromise.resultD inits) }
 
+  -- We create a promise for the info trees produced by the linters
+  let lintersInfoPromise ← IO.Promise.new (α := InfoTree)
+
   let cancelTk ← IO.CancelToken.new
   let lintAct ← wrapAsyncAsSnapshot (cancelTk? := cancelTk) fun (prev, infoSt) => do
     let messages := tree.getAll.map (·.diagnostics.msgLog) |>.foldl (· ++ ·) .empty
     let messages := messages.markAllReported
     modify fun st => { st with messages := st.messages ++ messages }
     modifyInfoState fun _ => infoSt
-    let postSt ← runStatefulLinters stx prev
+    let postSt ← runStatefulLinters stx prev lintersInfoPromise
     statePromise.resolve postSt
 
   let task ← BaseIO.bindTask (sync := true) (t := (← getInfoState).substituteLazy) fun infoSt =>
@@ -554,6 +590,13 @@ def runStatefulLintersAsync (stx : Syntax) : CommandElabM Unit := do
         lintAct (prev, infoSt)
   logSnapshotTask { stx? := none, task, cancelTk? := cancelTk }
 
+  let infoHole ← liftCoreM mkFreshMVarId
+  modifyInfoState fun s => { s with
+    trees          := s.trees.modify 0 (pushInfoChild · (.hole infoHole))
+    lazyAssignment := s.lazyAssignment
+      |>.insert infoHole (lintersInfoPromise.resultD default)
+  }
+
 /--
 Registers a command elaborator for the given syntax node kind.
 

src/Lean/Elab/InfoTree/Basic.lean

@@ -91,6 +91,16 @@ def InfoState.substituteLazy (s : InfoState) : Task InfoState :=
     lazyAssignment := {}
   }
 
+structure LinterInfoGroup where
+  deriving TypeName
+
+def mkLinterInfoGroupNode (trees : PersistentArray InfoTree) : InfoTree :=
+  InfoTree.node (Info.ofCustomInfo { stx := .missing, value := Dynamic.mk (⟨⟩ : LinterInfoGroup) }) trees
+
+def pushInfoChild : InfoTree → InfoTree → InfoTree
+  | .context ctx (.node i cs), child => .context ctx (.node i (cs.push child))
+  | t, _ => t
+
 def Info.toElabInfo? : Info → Option ElabInfo
   | ofTacticInfo i         => some i.toElabInfo
   | ofTermInfo i           => some i.toElabInfo

tests/elab/binopInfoTree.lean.out.expected

@@ -2,6 +2,7 @@
   • [Command] @ ⟨3, 0⟩-⟨3, 31⟩ @ Lean.Elab.Command.elabSetOption
     • [Completion] (Command.set_option "set_option" `trace.Elab.info []) @ ⟨3, 0⟩-⟨3, 26⟩
     • [Option] trace.Elab.info @ ⟨3, 11⟩-⟨3, 26⟩
+    • [CustomInfo(Lean.Elab.LinterInfoGroup)]
 1 + 2 + 3 : Nat
 [Elab.info] 
   • [Command] @ ⟨5, 0⟩-⟨5, 16⟩ @ Lean.Elab.Command.elabCheck
@@ -29,6 +30,7 @@
                   • [Term] 1 : Nat @ ⟨5, 7⟩-⟨5, 8⟩ @ Lean.Elab.Term.elabNumLit
                   • [Term] 2 : Nat @ ⟨5, 11⟩-⟨5, 12⟩ @ Lean.Elab.Term.elabNumLit
             • [Term] 3 : Nat @ ⟨5, 15⟩-⟨5, 16⟩ @ Lean.Elab.Term.elabNumLit
+    • [CustomInfo(Lean.Elab.LinterInfoGroup)]
 fun n m l => ↑n + (↑m + ↑l) : Nat → Nat → Nat → Int
 [Elab.info] 
   • [Command] @ ⟨7, 0⟩-⟨7, 48⟩ @ Lean.Elab.Command.elabCheck
@@ -81,5 +83,7 @@ fun n m l => ↑n + (↑m + ↑l) : Nat → Nat → Nat → Int
                           • [Term] l : Nat @ ⟨7, 39⟩-⟨7, 40⟩ @ Lean.Elab.Term.elabIdent
                             • [Completion-Id] l : none @ ⟨7, 39⟩-⟨7, 40⟩
                             • [Term] l : Nat @ ⟨7, 39⟩-⟨7, 40⟩
+    • [CustomInfo(Lean.Elab.LinterInfoGroup)]
 [Elab.info] 
   • [Command] @ ⟨8, 0⟩-⟨8, 0⟩ @ Lean.Elab.Command.elabEoi
+    • [CustomInfo(Lean.Elab.LinterInfoGroup)]

tests/elab/kernelErrorFollowup.lean

@@ -13,7 +13,7 @@ theorem zero_eq_one : 0 = 1 := bad.elim
 
 /--
 info: def bad : Empty :=
-?_uniq.4
+?_uniq.5
 -/
 #guard_msgs in
 set_option pp.raw true in

tests/pkg/stateful_linter/LinterTest/Suggest.lean

@@ -0,0 +1,27 @@
+import Lean
+
+/-! A stateful linter that attaches a `Try this` suggestion to the `#suggestMark` command in each of
+its `pre` and `post` passes. Checks that stateful linters can emit code-action suggestions and that
+both passes contribute. The suggestions are observed here as their `Try this` messages; the info-tree
+splicing that feeds the code-action provider is shared with ordinary linters and exercised by the
+interactive tests. -/
+
+namespace LinterTest.Suggest
+open Lean Elab Command Lean.Meta.Tactic.TryThis
+
+/-- Marker command the suggestion linter attaches `Try this` code actions to. -/
+syntax (name := suggestMark) "#suggestMark" : command
+
+elab_rules : command
+  | `(#suggestMark) => pure ()
+
+initialize _suggestLinter : StatefulLinter Unit Unit ←
+  registerStatefulLinter ()
+    (pre := fun stx _ _ => do
+      if stx.isOfKind ``suggestMark then
+        liftCoreM <| addSuggestion stx { suggestion := "#eval 0" }
+      pure none)
+    (post := fun stx _ _ _ _ => do
+      if stx.isOfKind ``suggestMark then
+        liftCoreM <| addSuggestion stx { suggestion := "#eval 1" }
+      pure ())

tests/pkg/stateful_linter/LinterTest/TestSuggest.lean

@@ -0,0 +1,14 @@
+import LinterTest.Suggest
+
+/- The `#suggestMark` linter attaches a `Try this` suggestion in each of its `pre` and `post` passes,
+so both surface as messages on the command. -/
+
+/--
+info: Try this:
+  [apply] #eval 0
+---
+info: Try this:
+  [apply] #eval 1
+-/
+#guard_msgs (ordering := sorted) in
+#suggestMark

tests/server_interactive/533.lean.out.expected

@@ -3,18 +3,18 @@
 {"items":
  [{"label": "Foo",
    "kind": 6,
-   "data": ["file:///533.lean", 4, 10, 0, "f_uniq.4"]},
+   "data": ["file:///533.lean", 4, 10, 0, "f_uniq.5"]},
   {"label": "F",
    "kind": 6,
-   "data": ["file:///533.lean", 4, 10, 0, "f_uniq.8"]}],
+   "data": ["file:///533.lean", 4, 10, 0, "f_uniq.9"]}],
  "isIncomplete": false}
 Resolution of Foo:
 {"label": "Foo",
  "kind": 6,
  "detail": "Sort ?u",
- "data": ["file:///533.lean", 4, 10, 0, "f_uniq.4"]}
+ "data": ["file:///533.lean", 4, 10, 0, "f_uniq.5"]}
 Resolution of F:
 {"label": "F",
  "kind": 6,
  "detail": "Sort ?u",
- "data": ["file:///533.lean", 4, 10, 0, "f_uniq.8"]}
+ "data": ["file:///533.lean", 4, 10, 0, "f_uniq.9"]}

tests/server_interactive/compHeader.lean.out.expected

@@ -3,7 +3,7 @@
 {"items":
  [{"label": "veryLongNam",
    "kind": 6,
-   "data": ["file:///compHeader.lean", 2, 22, 0, "f_uniq.11"]},
+   "data": ["file:///compHeader.lean", 2, 22, 0, "f_uniq.12"]},
   {"label": "veryLongNameForCompletion",
    "kind": 21,
    "data":
@@ -13,7 +13,7 @@ Resolution of veryLongNam:
 {"label": "veryLongNam",
  "kind": 6,
  "detail": "Sort u_1",
- "data": ["file:///compHeader.lean", 2, 22, 0, "f_uniq.11"]}
+ "data": ["file:///compHeader.lean", 2, 22, 0, "f_uniq.12"]}
 Resolution of veryLongNameForCompletion:
 {"label": "veryLongNameForCompletion",
  "kind": 21,

tests/server_interactive/completionPrefixIssue.lean.out.expected

@@ -3,7 +3,7 @@
 {"items":
  [{"label": "veryLongDefinitionName",
    "kind": 6,
-   "data": ["file:///completionPrefixIssue.lean", 1, 64, 0, "f_uniq.51"]},
+   "data": ["file:///completionPrefixIssue.lean", 1, 64, 0, "f_uniq.52"]},
   {"label": "veryLongDefinitionNameVeryLongDefinitionName",
    "kind": 21,
    "data":
@@ -17,7 +17,7 @@ Resolution of veryLongDefinitionName:
 {"label": "veryLongDefinitionName",
  "kind": 6,
  "detail": "Nat",
- "data": ["file:///completionPrefixIssue.lean", 1, 64, 0, "f_uniq.51"]}
+ "data": ["file:///completionPrefixIssue.lean", 1, 64, 0, "f_uniq.52"]}
 Resolution of veryLongDefinitionNameVeryLongDefinitionName:
 {"label": "veryLongDefinitionNameVeryLongDefinitionName",
  "kind": 21,

tests/server_interactive/linterCodeAction.lean

@@ -0,0 +1,24 @@
+import Lean
+
+/-!
+Checks that `Try this` suggestions emitted by a linter surface as `textDocument/codeAction` code
+actions.
+-/
+
+set_option Elab.async true
+
+open Lean Lean.Elab.Command Lean.Meta.Tactic.TryThis
+
+syntax (name := normalMark) "#normalMark" : command
+
+@[command_elab normalMark] def elabNormalMark : CommandElab := fun _ => pure ()
+
+def normalLinter : Linter where
+  run stx := do
+    if stx.isOfKind ``normalMark then
+      liftCoreM <| addSuggestion stx { suggestion := "#eval 0" }
+
+run_cmd liftCoreM <| addLinter normalLinter
+
+#normalMark
+--^ codeAction

tests/server_interactive/linterCodeAction.lean.out.expected

@@ -0,0 +1,44 @@
+{"textDocument": {"uri": "file:///linterCodeAction.lean"},
+ "range":
+ {"start": {"line": 22, "character": 2}, "end": {"line": 22, "character": 2}},
+ "context": {"diagnostics": []}}
+[{"title": "Try this: #eval 0",
+  "kind": "quickfix",
+  "edit":
+  {"documentChanges":
+   [{"textDocument": {"version": 1, "uri": "file:///linterCodeAction.lean"},
+     "edits":
+     [{"range":
+       {"start": {"line": 22, "character": 0},
+        "end": {"line": 22, "character": 11}},
+       "newText": "#eval 0"}]}]},
+  "data":
+  {"providerResultIndex": 0,
+   "providerName":
+   "_private.Lean.Meta.Tactic.TryThis.0.Lean.Meta.Tactic.TryThis.tryThisProvider",
+   "params":
+   {"textDocument": {"uri": "file:///linterCodeAction.lean"},
+    "range":
+    {"start": {"line": 22, "character": 2},
+     "end": {"line": 22, "character": 2}},
+    "context": {"diagnostics": []}}}}]
+Resolution of Try this: #eval 0:
+{"title": "Try this: #eval 0",
+ "kind": "quickfix",
+ "edit":
+ {"documentChanges":
+  [{"textDocument": {"version": 1, "uri": "file:///linterCodeAction.lean"},
+    "edits":
+    [{"range":
+      {"start": {"line": 22, "character": 0},
+       "end": {"line": 22, "character": 11}},
+      "newText": "#eval 0"}]}]},
+ "data":
+ {"providerResultIndex": 0,
+  "providerName":
+  "_private.Lean.Meta.Tactic.TryThis.0.Lean.Meta.Tactic.TryThis.tryThisProvider",
+  "params":
+  {"textDocument": {"uri": "file:///linterCodeAction.lean"},
+   "range":
+   {"start": {"line": 22, "character": 2}, "end": {"line": 22, "character": 2}},
+   "context": {"diagnostics": []}}}}