feat: interactive diagnostics take 1

src/Lean/Data/Lsp/Diagnostics.lean

@@ -28,7 +28,7 @@ instance : FromJson DiagnosticSeverity := ⟨fun j =>
   | Except.ok 2 => DiagnosticSeverity.warning
   | Except.ok 3 => DiagnosticSeverity.information
   | Except.ok 4 => DiagnosticSeverity.hint
-  | _      => throw "unknown DiagnosticSeverity"⟩
+  | _           => throw s!"unknown DiagnosticSeverity '{j}'"⟩
 
 instance : ToJson DiagnosticSeverity := ⟨fun
   | DiagnosticSeverity.error       => 1
@@ -44,7 +44,7 @@ inductive DiagnosticCode where
 instance : FromJson DiagnosticCode := ⟨fun
   | num (i : Int) => DiagnosticCode.int i
   | str s         => DiagnosticCode.string s
-  | _             => throw "the diagnostic code can only be a string or an integer"⟩
+  | j             => throw s!"expected string or integer diagnostic code, got '{j}'"⟩
 
 instance : ToJson DiagnosticCode := ⟨fun
   | DiagnosticCode.int i    => i
@@ -80,6 +80,9 @@ structure Diagnostic where
   message : String
   tags? : Option (Array DiagnosticTag) := none
   relatedInformation? : Option (Array DiagnosticRelatedInformation) := none
+  /-- Extension: interactive message widgets. We use untyped `Json` here due to dependency
+  ordering -- see `Lean.Widget.InteractiveDiagnostics`. -/
+  taggedMsg? : Option Json := none
   deriving Inhabited, BEq, ToJson, FromJson
 
 structure PublishDiagnosticsParams where
@@ -88,35 +91,5 @@ structure PublishDiagnosticsParams where
   diagnostics: Array Diagnostic
   deriving Inhabited, BEq, ToJson, FromJson
 
-/-- Transform a Lean Message concerning the given text into an LSP Diagnostic. -/
-def msgToDiagnostic (text : FileMap) (m : Message) : IO Diagnostic := do
-  let low : Lsp.Position := text.leanPosToLspPos m.pos
-  let fullHigh := text.leanPosToLspPos <| m.endPos.getD m.pos
-  let high : Lsp.Position := match m.endPos with
-    | some endPos =>
-      /-
-        Truncate messages that are more than one line long.
-        This is a workaround to avoid big blocks of "red squiggly lines" on VS Code.
-        TODO: should it be a parameter?
-      -/
-      let endPos := if endPos.line > m.pos.line then { line := m.pos.line + 1, column := 0 } else endPos
-      text.leanPosToLspPos endPos
-    | none        => low
-  let range : Range := ⟨low, high⟩
-  let fullRange : Range := ⟨low, fullHigh⟩
-  let severity := match m.severity with
-    | MessageSeverity.information => DiagnosticSeverity.information
-    | MessageSeverity.warning     => DiagnosticSeverity.warning
-    | MessageSeverity.error       => DiagnosticSeverity.error
-  let source := "Lean 4 server"
-  let message ← m.data.toString
-  pure {
-    range := range
-    fullRange := fullRange
-    severity? := severity
-    source? := source
-    message := message
-  }
-
 end Lsp
 end Lean

src/Lean/Server/FileWorker.lean

@@ -20,6 +20,8 @@ import Lean.Server.FileWorker.Utils
 import Lean.Server.FileWorker.RequestHandling
 import Lean.Server.Rpc.Basic
 
+import Lean.Widget.InteractiveDiagnostics
+
 /-!
 For general server architecture, see `README.md`. For details of IPC communication, see `Watchdog.lean`.
 This module implements per-file worker processes.
@@ -50,10 +52,21 @@ open Snapshots
 open Std (RBMap RBMap.empty)
 open JsonRpc
 
+def publishInteractiveMessages (m : DocumentMeta) (msgLog : MessageLog) (hOut : FS.Stream) (rpcSesh : RpcSession) : IO Unit := do
+  let diagnostics ← msgLog.msgs.mapM (Widget.msgToDiagnostic m.text · rpcSesh)
+  let diagParams : PublishDiagnosticsParams :=
+    { uri := m.uri
+      version? := some m.version
+      diagnostics := diagnostics.toArray }
+  hOut.writeLspNotification {
+    method := "textDocument/publishDiagnostics"
+    param := toJson diagParams
+  }
+
 /- Asynchronous snapshot elaboration. -/
 section Elab
   /-- Elaborates the next command after `parentSnap` and emits diagnostics into `hOut`. -/
-  private def nextCmdSnap (m : DocumentMeta) (parentSnap : Snapshot) (cancelTk : CancelToken) (hOut : FS.Stream)
+  private def nextCmdSnap (m : DocumentMeta) (parentSnap : Snapshot) (cancelTk : CancelToken) (hOut : FS.Stream) (rpcSesh : RpcSession)
       : ExceptT ElabTaskError IO Snapshot := do
     cancelTk.check
     publishProgressAtPos m parentSnap.endPos hOut
@@ -72,22 +85,22 @@ section Elab
         because we cannot guarantee that no further diagnostics are emitted after clearing
         them. -/
       if snap.msgLog.msgs.size > parentSnap.msgLog.msgs.size then
-        publishMessages m snap.msgLog hOut
+        publishInteractiveMessages m snap.msgLog hOut rpcSesh
       snap
     | Sum.inr msgLog =>
-      publishMessages m msgLog hOut
+      publishInteractiveMessages m msgLog hOut rpcSesh
       publishProgressDone m hOut
       throw ElabTaskError.eof
 
   /-- Elaborates all commands after `initSnap`, emitting the diagnostics into `hOut`. -/
-  def unfoldCmdSnaps (m : DocumentMeta) (initSnap : Snapshot) (cancelTk : CancelToken) (hOut : FS.Stream)
+  def unfoldCmdSnaps (m : DocumentMeta) (initSnap : Snapshot) (cancelTk : CancelToken) (hOut : FS.Stream) (rpcSesh : RpcSession)
     (initial : Bool) :
       IO (AsyncList ElabTaskError Snapshot) := do
     if initial && initSnap.msgLog.hasErrors then
       -- treat header processing errors as fatal so users aren't swamped with followup errors
       AsyncList.nil
     else
-      AsyncList.unfoldAsync (nextCmdSnap m . cancelTk hOut) initSnap
+      AsyncList.unfoldAsync (nextCmdSnap m . cancelTk hOut rpcSesh) initSnap
 end Elab
 
 -- Pending requests are tracked so they can be cancelled
@@ -167,7 +180,7 @@ section Initialization
     catch e =>  -- should be from `leanpkg print-paths`
       let msgs := MessageLog.empty.add { fileName := "<ignored>", pos := ⟨0, 0⟩, data := e.toString }
       pure (← mkEmptyEnvironment, msgs)
-    publishMessages m msgLog hOut
+    --publishMessages m msgLog hOut
     let cmdState := Elab.Command.mkState headerEnv msgLog opts
     let cmdState := { cmdState with infoState.enabled := true, scopes := [{ header := "", opts := opts }] }
     let headerSnap := {
@@ -187,7 +200,8 @@ section Initialization
       is a CR there. -/
     let (headerSnap, srcSearchPath) ← compileHeader meta o
     let cancelTk ← CancelToken.new
-    let cmdSnaps ← unfoldCmdSnaps meta headerSnap cancelTk o (initial := true)
+    let rpcSesh ← RpcSession.new false
+    let cmdSnaps ← unfoldCmdSnaps meta headerSnap cancelTk o (initial := true) rpcSesh
     let doc : EditableDocument := ⟨meta, headerSnap, cmdSnaps, cancelTk⟩
     return ({
       hIn                := i
@@ -198,7 +212,7 @@ section Initialization
     {
       doc             := doc
       pendingRequests := RBMap.empty
-      rpcSesh         := ← RpcSession.new false
+      rpcSesh
     })
 end Initialization
 
@@ -231,7 +245,7 @@ section Updates
       let mut validSnaps := cmdSnaps.finishedPrefix.takeWhile (fun s => s.endPos < changePos)
       if validSnaps.length = 0 then
         let cancelTk ← CancelToken.new
-        let newCmdSnaps ← unfoldCmdSnaps newMeta newHeaderSnap cancelTk ctx.hOut (initial := true)
+        let newCmdSnaps ← unfoldCmdSnaps newMeta newHeaderSnap cancelTk ctx.hOut (initial := true) (← get).rpcSesh
         modify fun st => { st with doc := ⟨newMeta, newHeaderSnap, newCmdSnaps, cancelTk⟩ }
       else
         /- When at least one valid non-header snap exists, it may happen that a change does not fall
@@ -249,7 +263,7 @@ section Updates
           validSnaps ← validSnaps.dropLast
           lastSnap ← preLastSnap
         let cancelTk ← CancelToken.new
-        let newSnaps ← unfoldCmdSnaps newMeta lastSnap cancelTk ctx.hOut (initial := false)
+        let newSnaps ← unfoldCmdSnaps newMeta lastSnap cancelTk ctx.hOut (initial := false) (← get).rpcSesh
         let newCmdSnaps := AsyncList.ofList validSnaps ++ newSnaps
         modify fun st => { st with doc := ⟨newMeta, newHeaderSnap, newCmdSnaps, cancelTk⟩ }
 end Updates

src/Lean/Server/Utils.lean

@@ -133,10 +133,6 @@ def publishDiagnostics (m : DocumentMeta) (diagnostics : Array Lsp.Diagnostic) (
     }
   }
 
-def publishMessages (m : DocumentMeta) (msgLog : MessageLog) (hOut : FS.Stream) : IO Unit := do
-  let diagnostics ← msgLog.msgs.mapM (msgToDiagnostic m.text)
-  publishDiagnostics m diagnostics.toArray hOut
-
 def publishProgress (m : DocumentMeta) (processing : Array LeanFileProgressProcessingInfo) (hOut : FS.Stream) : IO Unit :=
   hOut.writeLspNotification {
     method := "$/lean/fileProgress"

src/Lean/Widget.lean

@@ -6,4 +6,6 @@ Authors: Wojciech Nawrocki
 -/
 import Lean.Widget.ExprWithCtx
 import Lean.Widget.InteractiveGoals
+import Lean.Widget.InteractiveDiagnostics
 import Lean.Widget.TaggedText
+import Lean.Widget.ToHtmlFormat

src/Lean/Widget/InteractiveDiagnostics.lean

@@ -0,0 +1,127 @@
+import Lean.Data.Lsp
+import Lean.Message
+import Lean.Server.Rpc.Basic
+import Lean.Widget.TaggedText
+import Lean.Widget.ExprWithCtx
+
+namespace Lean.Widget
+open Lsp Server
+
+structure MsgEmbed where
+  -- TODO(WN): other variants, e.g. `ofGoal`
+  ef : ExprText
+  deriving Inhabited, RpcEncoding
+
+private def mkPPContext (nCtx : NamingContext) (ctx : MessageDataContext) : PPContext := {
+  env := ctx.env, mctx := ctx.mctx, lctx := ctx.lctx, opts := ctx.opts,
+  currNamespace := nCtx.currNamespace, openDecls := nCtx.openDecls
+}
+
+private inductive EmbedFmt
+  | ofExpr (e : ExprWithCtx)
+  deriving Inhabited
+
+private abbrev MsgFmtM := StateT (Array EmbedFmt) IO
+
+open MessageData in
+/-- We first build a `Nat`-tagged `Format` with the most shallow tag, if any,
+in every branch indexing into the array of embedded objects. -/
+private partial def msgToInteractiveAux (msgData : MessageData) : IO (Format × Array EmbedFmt) :=
+  go { currNamespace := Name.anonymous, openDecls := [] } none msgData #[]
+where
+  pushEmbedExpr (e : ExprWithCtx) : MsgFmtM Nat :=
+    modifyGet fun es => (es.size, es.push (EmbedFmt.ofExpr e))
+
+  go : NamingContext → Option MessageDataContext → MessageData → MsgFmtM Format
+  | _,    _,         ofFormat fmt             => pure fmt
+  | _,    _,         ofLevel u                => pure $ format u
+  | _,    _,         ofName n                 => pure $ format n
+  | nCtx, some ctx,  ofSyntax s               => ppTerm (mkPPContext nCtx ctx) s  -- HACK: might not be a term
+  | _,    none,      ofSyntax s               => pure $ s.formatStx
+  | _,    none,      ofExpr e                 => pure $ format (toString e)
+  | nCtx, some ctx,  ofExpr e                 => do
+    let f ← ppExpr (mkPPContext nCtx ctx) e
+    let t ← pushEmbedExpr {
+      expr := e
+      env := ctx.env
+      mctx := ctx.mctx
+      lctx := ctx.lctx
+      options := ctx.opts
+      currNamespace := nCtx.currNamespace
+      openDecls := nCtx.openDecls }
+    return Format.tag t f
+  | _,    none,      ofGoal mvarId            => pure $ "goal " ++ format (mkMVar mvarId)
+  | nCtx, some ctx,  ofGoal mvarId            => ppGoal (mkPPContext nCtx ctx) mvarId
+  | nCtx, _,         withContext ctx d        => go nCtx ctx d
+  | _,    ctx,       withNamingContext nCtx d => go nCtx ctx d
+  | nCtx, ctx,       tagged _ d               => go nCtx ctx d
+  | nCtx, ctx,       nest n d                 => Format.nest n <$> go nCtx ctx d
+  | nCtx, ctx,       compose d₁ d₂            => do let d₁ ← go nCtx ctx d₁; let d₂ ← go nCtx ctx d₂; pure $ d₁ ++ d₂
+  | nCtx, ctx,       group d                  => Format.group <$> go nCtx ctx d
+  | nCtx, ctx,       node ds                  => Format.nest 2 <$> ds.foldlM (fun r d => do let d ← go nCtx ctx d; pure $ r ++ Format.line ++ d) Format.nil
+
+private def msgToInteractive (msgData : MessageData) : IO (TaggedText MsgEmbed) := do
+  let (fmt, embeds) ← msgToInteractiveAux msgData
+  let tt := TaggedText.prettyTagged fmt
+  /- Here we rewrite a `TaggedText Nat` corresponding to a whole `MessageData` into one where
+  the tags are `TaggedText MsgEmbed`s corresponding to embedded objects with their subtree being
+  empty (`text ""`). In other words, we terminate the `MsgEmbed`-tree at embedded objects
+  and store the pretty-printed embed (which is itself a `TaggedText (WithRpcRef LazyExprWithCtx)`,
+  for example) in the tag. -/
+  tt.rewriteM fun n subTt =>
+    match embeds.get! n with
+    | EmbedFmt.ofExpr e =>
+      TaggedText.tag
+        { ef := subTt.map fun n => ⟨fun () => e.runMetaM (e.traverse n)⟩ }
+        (TaggedText.text "")
+
+/-- Remove tags, leaving just the pretty-printed string. -/
+partial def TaggedText.stripTags (tt : TaggedText α) : String :=
+  go "" [tt]
+where go (acc : String) : List (TaggedText α) → String
+  | []               => acc
+  | text s :: ts     => go (acc ++ s) ts
+  | append a b :: ts => go acc (a :: b :: ts)
+  | tag _ a :: ts    => go acc (a :: ts)
+
+partial def TaggedText.stripTags₂ (tt : TaggedText (MsgEmbed)) : String :=
+  go "" [tt]
+where go (acc : String) : List (TaggedText (MsgEmbed)) → String
+  | []               => acc
+  | text s :: ts     => go (acc ++ s) ts
+  | append a b :: ts => go acc (a :: b :: ts)
+  | tag ⟨et⟩ _ :: ts => go acc (text et.stripTags :: ts)
+
+/-- Transform a Lean Message concerning the given text into an LSP Diagnostic. -/
+def msgToDiagnostic (text : FileMap) (m : Message) : ReaderT RpcSession IO Diagnostic := do
+  let low : Lsp.Position := text.leanPosToLspPos m.pos
+  let fullHigh := text.leanPosToLspPos <| m.endPos.getD m.pos
+  let high : Lsp.Position := match m.endPos with
+    | some endPos =>
+      /-
+        Truncate messages that are more than one line long.
+        This is a workaround to avoid big blocks of "red squiggly lines" on VS Code.
+        TODO: should it be a parameter?
+      -/
+      let endPos := if endPos.line > m.pos.line then { line := m.pos.line + 1, column := 0 } else endPos
+      text.leanPosToLspPos endPos
+    | none        => low
+  let range : Range := ⟨low, high⟩
+  let fullRange : Range := ⟨low, fullHigh⟩
+  let severity := match m.severity with
+    | MessageSeverity.information => DiagnosticSeverity.information
+    | MessageSeverity.warning     => DiagnosticSeverity.warning
+    | MessageSeverity.error       => DiagnosticSeverity.error
+  let source := "Lean 4 server"
+  let tt ← msgToInteractive m.data
+  let ttJson ← toJson <$> rpcEncode tt
+  pure {
+    range := range
+    fullRange := fullRange
+    severity? := severity
+    source? := source
+    message := tt.stripTags₂
+    taggedMsg? := ttJson
+  }
+
+end Lean.Widget

src/Lean/Widget/TaggedText.lean

@@ -17,7 +17,7 @@ inductive TaggedText (α : Type u) where
   | text (s : String)
   | append (a b : TaggedText α)
   | tag (t : α) (a : TaggedText α)
-  deriving Inhabited, FromJson, ToJson
+  deriving Inhabited, BEq, FromJson, ToJson
 
 namespace TaggedText
 

src/Lean/Widget/ToHtmlFormat.lean

@@ -0,0 +1,6 @@
+/-- Types implementing this can displayed in HTML, for example as a SVG.
+This is non-interactive -- for interactivity, see widgets. -/
+-- should suffice for Dynkin
+-- https://leanprover.zulipchat.com/#narrow/stream/113488-general/topic/Dynkin.20diagrams
+class ToHtmlFormat (α : Type u) where
+  toHtmlFormat : α → String -- TODO HTML