feat: code action resolvers

src/Init/System/ST.lean

@@ -108,6 +108,11 @@ variable {σ : Type} {m : Type → Type} [Monad m] [MonadLiftT (ST σ) m]
 @[inline] def Ref.modify {α : Type} (r : Ref σ α) (f : α → α) : m Unit := liftM <| Prim.Ref.modify r f
 @[inline] def Ref.modifyGet {α : Type} {β : Type} (r : Ref σ α) (f : α → β × α) : m β := liftM <| Prim.Ref.modifyGet r f
 
+def Ref.toMonadStateOf (r : Ref σ α) : MonadStateOf α m where
+  get := r.get
+  set := r.set
+  modifyGet := r.modifyGet
+
 end
 
 end ST

src/Lean/Data/Json/Basic.lean

@@ -268,6 +268,15 @@ def setObjVal! : Json → String → Json → Json
   | obj kvs, k, v => obj <| kvs.insert compare k v
   | _      , _, _ => panic! "Json.setObjVal!: not an object: {j}"
 
+open Lean.RBNode in
+/-- Assuming both inputs `o₁, o₂` are json objects, will compute `{...o₁, ...o₂}`.
+If `o₁` is not a json object, `o₂` will be returned.
+-/
+def mergeObj : Json → Json → Json
+  | obj kvs₁, obj kvs₂ =>
+    obj <| fold (insert compare) kvs₁ kvs₂
+  | _, j₂ => j₂
+
 inductive Structured where
   | arr (elems : Array Json)
   | obj (kvPairs : RBNode String (fun _ => Json))

src/Lean/Server/CodeActions.lean

@@ -14,6 +14,70 @@ open Lsp
 open RequestM
 open Snapshots
 
+structure CodeActionData where
+  uri : DocumentUri
+  resolverId : UInt64
+  deriving ToJson, FromJson
+
+def CodeAction.getFileSource! (ca : CodeAction) : DocumentUri :=
+  let r : Except String DocumentUri := do
+    let some data := ca.data?
+      | throw s!"no data param on code action {ca.title}"
+    let data : CodeActionData ← fromJson? data
+    return data.uri
+  match r with
+  | Except.ok uri => uri
+  | Except.error e => panic! e
+
+instance : FileSource CodeAction where
+  fileSource x := CodeAction.getFileSource! x
+
+/-- A code action optionally supporting a lazy code action computation that is only run when the user clicks on
+the code action in the editor.
+
+If you want to use the lazy feature, make sure that the `edit?` field on the `eager` code action result is `none`.
+ -/
+structure LazyCodeAction where
+  /-- This is the initial code action that is sent to the server, to implement -/
+  eager : CodeAction
+  lazy? : Option (IO CodeAction) := none
+
+instance : Coe CodeAction LazyCodeAction where
+  coe c := { eager := c }
+
+abbrev CodeActionResolverState := RBMap UInt64 LazyCodeAction compare
+
+builtin_initialize codeActionResolver : IO.Ref CodeActionResolverState ← IO.mkRef ∅
+
+local instance : MonadStateOf CodeActionResolverState IO := codeActionResolver.toMonadStateOf
+
+def resetCodeActionResolverState : RequestM Unit := do
+  set (σ := CodeActionResolverState) <| ∅
+
+/-- Convert a LazyCodeAction to a CodeAction and register it. -/
+def LazyCodeAction.convert (uri : DocumentUri) (cs : LazyCodeAction) : RequestM CodeAction := do
+  if cs.lazy?.isNone || cs.eager.edit?.isSome then
+    return cs.eager
+  let s : CodeActionResolverState ← get
+  let resolverId := (Prod.fst <$> s.max).getD 0 |> (· + 1)
+  set <| s.insert resolverId cs
+  let data : CodeActionData := { resolverId, uri }
+  let data? := Option.merge Json.mergeObj cs.eager.data? (some <| toJson data)
+  return { cs.eager with data? }
+
+/-- Find the registered lazy resolver for the given code action and run it. -/
+def LazyCodeAction.resolve (ca : CodeAction) : RequestM CodeAction := do
+  let some data := ca.data?
+    | throw <| RequestError.internalError s!"No data field on {ca.title}"
+  let data : CodeActionData ← liftExcept <| Except.mapError RequestError.internalError <| fromJson? data
+  let s : CodeActionResolverState ← get
+  let some x := s.find? data.resolverId
+    | throw <| RequestError.internalError s!"Can't find resolver id {data.resolverId} for code action {ca.title}."
+  let some x := x.lazy?
+    | throw <| RequestError.internalError s!"No lazy resolver for {ca.title}."
+  let r ← liftM x
+  return r
+
 /-- A code action provider is a function for providing LSP code actions for an editor.
 You can register them with the `@[codeActionProvider]` attribute.
 
@@ -26,14 +90,10 @@ This interface can be used to implement the following applications:
 - Helper suggestions for tactics.
 
 When implementing your own `CodeActionProvider`, we assume that no long-running computations are allowed.
-If you need to create a code-action with a long-running computation, we suggest the following approach:
-
-Wrap the long-running computation in a tactic or macro which itself suggests a code action.
-This means that the long-running computation will be handled in the normal document processing thread.
-For example, suppose that you want to run some machine learning to suggest the next lemma to use.
-Rather than run the heavy ML computation in the CodeActionProvider, you could suggest a `by try_ml` tactic. Which once elaborated by Lean will suggest a new lemma.
+If you need to create a code-action with a long-running computation, you can use the `lazy?` field on `LazyCodeAction`
+to perform the computation after the user has clicked on the code action in their editor.
 -/
-def CodeActionProvider := CodeActionParams → RequestM (RequestTask (Array CodeAction))
+def CodeActionProvider := CodeActionParams → Snapshot → RequestM (Array LazyCodeAction)
 deriving instance Inhabited for CodeActionProvider
 
 builtin_initialize codeActionProviderExt : SimplePersistentEnvExtension Name NameSet ← registerSimplePersistentEnvExtension {
@@ -66,20 +126,30 @@ This is implemented by calling all of the registered `CodeActionProvider` functi
 def handleCodeAction (params : CodeActionParams) : RequestM (RequestTask (Array CodeAction)) := do
   let doc ← readDoc
   let pos := doc.meta.text.lspPosToUtf8Pos params.range.end
-  bindWaitFindSnap doc (fun s => s.endPos ≥ pos)
-    (notFoundX := return RequestTask.pure #[])
+  withWaitFindSnap doc (fun s => s.endPos ≥ pos)
+    (notFoundX := return #[])
     fun snap => do
       let caps ← RequestM.runCoreM snap do
         let env ← getEnv
         let names := codeActionProviderExt.getState env
         names.toArray.mapM evalCodeActionProvider
-      have : Monad Task := {bind := Task.bind, pure := Task.pure}
-      have : Monad RequestTask := show Monad (ExceptT _ Task) by infer_instance
-      let tasks : Array (RequestTask _) ← caps.mapM (· <| params)
-      let task := tasks.concatMapM id
-      return task
+      resetCodeActionResolverState
+      let cas ← caps.concatMapM (· params snap)
+      let cas ← cas.mapM <| LazyCodeAction.convert doc.meta.uri
+      return cas
 
 builtin_initialize
   registerLspRequestHandler "textDocument/codeAction" CodeActionParams (Array CodeAction) handleCodeAction
 
+/-- Handler for `"codeAction/resolve"`.
+
+[reference](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#codeAction_resolve)
+-/
+def handleCodeActionResolve (param : CodeAction) : RequestM (RequestTask CodeAction) := do
+  let r ← LazyCodeAction.resolve param
+  return RequestTask.pure r
+
+builtin_initialize
+  registerLspRequestHandler "codeAction/resolve" CodeAction CodeAction handleCodeActionResolve
+
 end Lean.Server

tests/lean/interactive/codeaction.lean

@@ -3,16 +3,28 @@ import Lean
 open Lean Server Lsp
 
 @[codeActionProvider]
-def helloProvider : CodeActionProvider := fun (params : CodeActionParams) => do
-  let uri := params.textDocument.uri
-  return RequestTask.pure #[{
+def helloProvider : CodeActionProvider := fun params _snap => do
+  let td := params.textDocument
+  let edit : TextEdit := {
+      range := params.range,
+      newText := "hello!!!"
+    }
+  let ca : CodeAction := {
     title := "hello world",
     kind? := "quickfix",
-    edit? := WorkspaceEdit.ofTextEdit uri {
-      range := params.range,
-      newText := "hello!!!",
+    edit? := WorkspaceEdit.ofTextEdit td.uri edit
+  }
+  let longRunner : CodeAction := {
+    title := "a long-running action",
+    kind? := "refactor",
+  }
+  let lazyResult : IO CodeAction := do
+    let v? ← IO.getEnv "PWD"
+    let v := v?.getD "none"
+    return { longRunner with
+      edit? := WorkspaceEdit.ofTextEdit td.uri { range := params.range, newText := v}
     }
-  }]
+  return #[ca, {eager := longRunner, lazy? := lazyResult}]
 
 theorem asdf : (x : Nat) → x = x := by
   intro x

tests/lean/interactive/codeaction.lean.expected.out

@@ -5,8 +5,11 @@
   [{"textDocument": {"version": 0, "uri": "file://codeaction.lean"},
     "edits":
     [{"range":
-      {"start": {"line": 17, "character": 4},
-       "end": {"line": 17, "character": 4}},
+      {"start": {"line": 29, "character": 4},
+       "end": {"line": 29, "character": 4}},
       "newText": "hello!!!"}]}],
   "changes": {},
   "changeAnnotations": {}}}
+{"title": "a long-running action",
+ "kind": "refactor",
+ "data": {"uri": "file://codeaction.lean", "resolverId": "1"}}