feat: server-side for trace search (#10365)

This PR implements the server-side for a new trace search mechanism in
the InfoView.

Demo:
![Search
demo](https://github.com/user-attachments/assets/f8f1cdfd-a4f2-4258-8cb8-360f64ea06e9)

src/Lean/Data/Lsp/Capabilities.lean

@@ -89,6 +89,7 @@ def ClientCapabilities.silentDiagnosticSupport (c : ClientCapabilities) : Bool :
 
 structure LeanServerCapabilities where
   moduleHierarchyProvider? : Option ModuleHierarchyOptions
+  rpcProvider? : Option RpcOptions
   deriving FromJson, ToJson
 
 -- TODO largely unimplemented

src/Lean/Data/Lsp/Diagnostics.lean

@@ -117,7 +117,7 @@ structure DiagnosticRelatedInformation where
 /-- Represents a diagnostic, such as a compiler error or warning. Diagnostic objects are only valid in the scope of a resource.
 
 LSP accepts a `Diagnostic := DiagnosticWith String`.
-The infoview also accepts `InteractiveDiagnostic := DiagnosticWith (TaggedText MsgEmbed)`.
+The infoview also accepts `InteractiveDiagnostic := DiagnosticWith InteractiveMessage`.
 
 [reference](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#diagnostic) -/
 structure DiagnosticWith (α : Type) where

src/Lean/Data/Lsp/Extra.lean

@@ -128,6 +128,13 @@ structure PlainTermGoal where
 structure ModuleHierarchyOptions where
   deriving FromJson, ToJson
 
+structure HighlightMatchesOptions where
+  deriving FromJson, ToJson
+
+structure RpcOptions where
+  highlightMatchesProvider? : Option HighlightMatchesOptions := none
+  deriving FromJson, ToJson
+
 structure LeanModule where
   name  : String
   uri   : DocumentUri

src/Lean/Server/FileWorker/WidgetRequests.lean

@@ -32,7 +32,7 @@ builtin_initialize
   registerBuiltinRpcProcedure
     `Lean.Widget.InteractiveDiagnostics.msgToInteractive
     MsgToInteractive
-    (TaggedText MsgEmbed)
+    InteractiveMessage
     fun ⟨m, i⟩ => RequestM.pureTask do msgToInteractive m.val i (hasWidgets := true)
 
 /-- The information that the infoview uses to render a popup
@@ -152,12 +152,315 @@ builtin_initialize
         return ls.map (·.toLocationLink)
 
 def lazyTraceChildrenToInteractive (children : WithRpcRef LazyTraceChildren) :
-    RequestM (RequestTask (Array (TaggedText MsgEmbed))) :=
+    RequestM (RequestTask (Array InteractiveMessage)) :=
   RequestM.pureTask do
     let ⟨indent, children⟩ := children.val
-    children.mapM fun child =>
+    children.mapM fun child => do
       msgToInteractive child.val (hasWidgets := true) (indent := indent)
 
 builtin_initialize registerBuiltinRpcProcedure ``lazyTraceChildrenToInteractive _ _ lazyTraceChildrenToInteractive
 
+private def kmpSearch (query text : String) : Array String.Pos := Id.run do
+  if query.isEmpty then
+    return #[]
+  let query := query.toUTF8
+  let text := text.toUTF8
+  let table := buildKMPTable query
+  let mut r := #[]
+  let mut qi : Int := 0
+  for h:ti in *...text.size do
+    while qi >= 0 && text[ti] != query[qi.toNat]! do
+      qi := table[qi.toNat]!
+    qi := qi + 1
+    if qi == query.size then
+      r := r.push <| (ti + 1) - query.size
+      qi := table[qi.toNat]!
+  return r.map (⟨·⟩)
+where
+  buildKMPTable (w : ByteArray) : Array Int := Id.run do
+    let mut t := Array.emptyWithCapacity w.size
+    let mut n := -1
+    t := t.push n
+    for h:i in *...w.size do
+      while n >= 0 && w[n.toNat]! != w[i] do
+        n := t[n.toNat]!
+      n := n + 1
+      t := t.push n
+    return t
+
+private def contains (query text : String) : Bool :=
+  ! (kmpSearch query text).isEmpty
+
+private def matchEndPos (query : String) (startPos : String.Pos) : String.Pos :=
+  startPos + ⟨query.utf8ByteSize⟩
+
+@[specialize]
+private def hightlightStringMatches? (query text : String) (matchPositions : Array String.Pos)
+    (highlight : α) (offset : String.Pos := ⟨0⟩) (mapIdx : Nat → Nat := id) :
+    Option (TaggedText α) := Id.run do
+  if query.isEmpty || text.isEmpty || matchPositions.isEmpty then
+    return none
+  let mut anyMatch : Bool := false
+  let mut r : Array (TaggedText α) := #[]
+  let mut p : String.Pos := ⟨0⟩
+  for i in 0...matchPositions.size do
+    if p >= text.endPos then
+      break
+    let i := mapIdx i
+    let globalMatchPos := matchPositions[i]!
+    let matchPos := globalMatchPos - offset
+    if matchPos >= text.endPos then
+      break
+    if let some nonMatch := nonMatch? p matchPos then
+      r := r.push nonMatch
+    let globalMatchEndPos := matchEndPos query globalMatchPos
+    let matchEndPos := globalMatchEndPos - offset
+    let «match» := text.extract matchPos matchEndPos
+    r := r.push <| .tag highlight (.text «match»)
+    p := matchEndPos
+    anyMatch := true
+  if let some nonMatch := nonMatch? p text.endPos then
+    r := r.push nonMatch
+  if ! anyMatch then
+    return none
+  if h : r.size = 1 then
+    return some r[0]
+  return some (.append r)
+where
+  nonMatch? (p matchPosition : String.Pos) : Option (TaggedText α) := do
+    guard <| p < matchPosition
+    let nonMatch := text.extract p matchPosition
+    return .text nonMatch
+
+private def findTaggedTextMatches (query : String) (tt : TaggedText α) (toText : α → String) :
+    Array String.Pos :=
+  let tt : TaggedText Empty := tt.rewrite fun t _ => .text (toText t)
+  kmpSearch query tt.stripTags
+
+private structure TaggedTextHighlightState where
+  query        : String
+  ms           : Array String.Pos
+  p            : String.Pos
+  anyHighlight : Bool
+
+private def advanceTaggedTextHighlightState (text : String) (highlighted : α) :
+    StateM TaggedTextHighlightState (TaggedText α) := do
+  let query := (← get).query
+  let p := (← get).p
+  let ms := (← get).ms
+  let highlighted? := hightlightStringMatches? query text ms highlighted (offset := p)
+    (mapIdx := fun i => ms.size - i - 1)
+  updateState text highlighted?.isSome
+  return highlighted?.getD (.text text)
+where
+  updateState (text : String) (isHighlighted : Bool) : StateM TaggedTextHighlightState Unit :=
+    modify fun s =>
+      let p : String.Pos := s.p + ⟨text.utf8ByteSize⟩
+      let ms := updateMatches s.query s.ms p
+      let anyHighlight := s.anyHighlight || isHighlighted
+      { s with p, ms, anyHighlight }
+  updateMatches (query : String) (ms : Array String.Pos) (p : String.Pos) : Array String.Pos := Id.run do
+    let n := ms.size
+    let mut ms := ms
+    for i in 0...n do
+      if p >= matchEndPos query ms[n - i - 1]! then
+        ms := ms.pop
+    return ms
+
+@[specialize]
+private partial def highlightTaggedText (query : String) (tt : TaggedText α) (highlighted : β)
+    (toText : (t : α) → String)
+    (highlightTag :
+      (t : α) →
+      (highlightTaggedText :
+        TaggedText α →
+        StateM TaggedTextHighlightState (TaggedText β)) →
+      StateM TaggedTextHighlightState β) :
+    TaggedText β × Bool :=
+  let ms := findTaggedTextMatches query tt toText
+  let (tt, s) := go tt |>.run { query, p := ⟨0⟩, ms := ms.reverse, anyHighlight := false }
+  (tt, s.anyHighlight)
+where
+  go (tt : TaggedText α) : StateM TaggedTextHighlightState (TaggedText β) := do
+    match tt with
+    | .text s =>
+      advanceTaggedTextHighlightState s highlighted
+    | .append a =>
+      return .append (← a.mapM go)
+    | .tag t a =>
+      let t ← highlightTag t go
+      let a ← go a
+      return .tag t a
+
+inductive HighlightedSubexprInfo where
+  | subexpr (info : SubexprInfo)
+  | highlighted
+
+instance : RpcEncodable HighlightedSubexprInfo where
+  rpcEncode
+    | .subexpr info => rpcEncode info
+    | .highlighted => pure "highlighted"
+  rpcDecode
+    | .str "highlighted" => pure .highlighted
+    | j => do
+      return .subexpr (← rpcDecode j)
+
+abbrev HighlightedCodeWithInfos := TaggedText HighlightedSubexprInfo
+
+private def HighlightedCodeWithInfos.ofCodeWithInfos (c : CodeWithInfos) :
+    HighlightedCodeWithInfos :=
+  c.map .subexpr
+
+private partial def highlightCodeWithInfos (query : String) (c : CodeWithInfos) :
+    HighlightedCodeWithInfos × Bool :=
+  highlightTaggedText query c .highlighted (fun _ => "") fun i _ => pure <| .subexpr i
+
+inductive HighlightedMsgEmbed where
+  | expr : HighlightedCodeWithInfos → HighlightedMsgEmbed
+  | goal : InteractiveGoal → HighlightedMsgEmbed
+  | widget (wi : Widget.WidgetInstance) (alt : TaggedText HighlightedMsgEmbed)
+  | trace (indent : Nat) (cls : Name) (msg : TaggedText HighlightedMsgEmbed) (collapsed : Bool)
+      (children : StrictOrLazy
+        (Array (TaggedText HighlightedMsgEmbed))
+        (WithRpcRef LazyTraceChildren))
+  | highlighted
+  deriving Inhabited, RpcEncodable
+
+mutual
+
+private partial def HighlightedMsgEmbed.traceChildrenOfMsgEmbed :
+    StrictOrLazy
+      (Array (TaggedText MsgEmbed))
+      (WithRpcRef LazyTraceChildren) →
+    StrictOrLazy
+      (Array (TaggedText HighlightedMsgEmbed))
+      (WithRpcRef LazyTraceChildren)
+  | .strict cs => .strict <| cs.map fun c => c.map (.ofMsgEmbed ·)
+  | .lazy cs => .lazy cs
+
+private partial def HighlightedMsgEmbed.ofMsgEmbed : MsgEmbed → HighlightedMsgEmbed
+  | .expr c => .expr (.ofCodeWithInfos c)
+  | .goal g => .goal g
+  | .widget wi alt => .widget wi <| alt.map (.ofMsgEmbed ·)
+  | .trace indent cls msg collapsed children =>
+    let msg := msg.map (.ofMsgEmbed ·)
+    let children := traceChildrenOfMsgEmbed children
+    .trace indent cls msg collapsed children
+
+end
+
+abbrev HighlightedInteractiveMessage := TaggedText HighlightedMsgEmbed
+
+private def HighlightedInteractiveMessage.ofInteractiveMessage (msg : InteractiveMessage) :
+    HighlightedInteractiveMessage :=
+  msg.map (.ofMsgEmbed ·)
+
+private def highlightInteractiveMessageWithExprs (query : String) (msg : InteractiveMessage) :
+    HighlightedInteractiveMessage × Bool :=
+  highlightTaggedText query msg .highlighted toText go
+where
+  toText : MsgEmbed → String
+    | .expr tt => tt.stripTags
+    | _        => ""
+  go e _ := do
+    match e with
+    | .expr tt =>
+      let tt ← highlightTaggedText.go (tt := tt) .highlighted fun i _ => pure <| .subexpr i
+      return .expr tt
+    | e =>
+      pure <| .ofMsgEmbed e
+
+variable (query : String) (indent : Nat) in
+private partial def unfoldMessageDataTracesContaining (msg : MessageData) (nctx : NamingContext) (ctx? : Option MessageDataContext) : BaseIO (MessageData × Bool) := do
+  -- We currently only support `.trace` trees.
+  match msg with
+  | .trace data traceMsg children =>
+    let unfolded ← children.mapM (unfoldMessageDataTracesContaining · nctx ctx?)
+    let children := unfolded.map (·.1)
+    let anyUnfolded := unfolded.any (·.2)
+    if anyUnfolded then
+      return (.trace { data with collapsed := false } traceMsg children, true)
+    let fmt ← traceMsg.formatAux nctx ctx?
+    let s := fmt.pretty (indent := indent)
+    if contains query s then
+      return (.trace { data with collapsed := false } traceMsg children, true)
+    return (.trace data traceMsg children, false)
+  | .withContext ctx msg =>
+    let (msg, unfolded) ← unfoldMessageDataTracesContaining msg nctx ctx
+    return (.withContext ctx msg, unfolded)
+  | .withNamingContext nctx msg =>
+    let (msg, unfolded) ← unfoldMessageDataTracesContaining msg nctx ctx?
+    return (.withNamingContext nctx msg, unfolded)
+  | .tagged tag msg =>
+    let (msg, unfolded) ← unfoldMessageDataTracesContaining msg nctx ctx?
+    return (.tagged tag msg, unfolded)
+  | _ =>
+    return (msg, false)
+
+private structure Highlighted where
+  msg : HighlightedInteractiveMessage
+  isHighlighted : Bool
+  isTrace : Bool
+
+variable (query : String) in
+private partial def highlightMatchesAux (msg : InteractiveMessage) :
+    IO Highlighted := do
+  match msg with
+  | .tag (.trace indent cls msg collapsed children) a => do
+    let a : HighlightedInteractiveMessage := .ofInteractiveMessage a
+    let (msg, isMsgHighlighted) := highlightInteractiveMessageWithExprs query msg
+    let unchangedChildren := Thunk.mk fun _ =>
+      let children := HighlightedMsgEmbed.traceChildrenOfMsgEmbed children
+      let trace := .trace indent cls msg collapsed children
+      { msg := .tag trace a, isHighlighted := isMsgHighlighted, isTrace := true }
+    let expandedChildren children :=
+      let trace := .trace indent cls msg false (.strict children)
+      { msg := .tag trace a, isHighlighted := true, isTrace := true }
+    match children with
+    | .strict children =>
+      let highlighted ← children.mapM (highlightMatchesAux ·)
+      let anyHighlighted := highlighted.any (·.2)
+      let children := highlighted.map (·.1)
+      if ! anyHighlighted then
+        return unchangedChildren.get
+      return expandedChildren children
+    | .lazy c => do
+      let indent := c.val.indent
+      let unfolded ← c.val.children.mapM fun childRef =>
+        let child := childRef.val
+        let nctx := { currNamespace := Name.anonymous, openDecls := [] }
+        unfoldMessageDataTracesContaining query indent child nctx none
+      let anyUnfolded := unfolded.any (·.2)
+      let children := unfolded.map (·.1)
+      if ! anyUnfolded then
+        return unchangedChildren.get
+      let children ← children.mapM (msgToInteractive · true indent)
+      let highlighted ← children.mapM highlightMatchesAux
+      let children := highlighted.map (·.1)
+      return expandedChildren children
+  | .append as =>
+    let highlighted ← as.mapM highlightMatchesAux
+    let anyHighlighted := highlighted.any (·.isHighlighted)
+    let anyTrace := highlighted.any (·.isTrace)
+    let as := highlighted.map (·.msg)
+    if ! anyTrace then
+      let (msg, isHighlighted) := highlightInteractiveMessageWithExprs query msg
+      return { msg, isHighlighted, isTrace := false }
+    return { msg := .append as, isHighlighted := anyHighlighted, isTrace := true }
+  | msg =>
+    let (msg, isHighlighted) := highlightInteractiveMessageWithExprs query msg
+    return { msg, isHighlighted, isTrace := false }
+
+structure HighlightMatchesParams where
+  query : String
+  msg : InteractiveMessage
+  deriving RpcEncodable
+
+def highlightMatches (params : HighlightMatchesParams) : RequestM (RequestTask HighlightedInteractiveMessage) :=
+  RequestM.pureTask do
+    let r ← highlightMatchesAux params.query params.msg
+    return r.msg
+
+builtin_initialize registerBuiltinRpcProcedure ``highlightMatches _ _ highlightMatches
+
 end Lean.Widget

src/Lean/Server/InfoUtils.lean

@@ -388,11 +388,14 @@ where
         return (some { eFmt with fmt := f!"```lean\n{eFmt.fmt}\n```" }, ← fmtModule? c)
       let eFmt ← Meta.ppExpr e
       -- Try not to show too scary internals
-      let showTerm := if let .fvar _ := e then
-        if let some ldecl := (← getLCtx).findFVar? e then
-          !ldecl.userName.hasMacroScopes
-        else false
-      else isAtomicFormat eFmt
+      let showTerm :=
+        if let .fvar _ := e then
+          if let some ldecl := (← getLCtx).findFVar? e then
+            !ldecl.userName.hasMacroScopes
+          else
+            false
+        else
+          isAtomicFormat eFmt
       let fmt := if showTerm then f!"{eFmt} : {tpFmt}" else tpFmt
       return (some f!"```lean\n{fmt}\n```", none)
     | Info.ofFieldInfo fi =>

src/Lean/Server/Watchdog.lean

@@ -1599,6 +1599,9 @@ def mkLeanServerCapabilities : ServerCapabilities := {
   }
   experimental? := some {
     moduleHierarchyProvider? := some {}
+    rpcProvider? := some {
+      highlightMatchesProvider? := some {}
+    }
   }
 }
 

src/Lean/Widget/InteractiveDiagnostic.lean

@@ -46,6 +46,10 @@ inductive MsgEmbed where
       (children : StrictOrLazy (Array (TaggedText MsgEmbed)) (WithRpcRef LazyTraceChildren))
   deriving Inhabited, RpcEncodable
 
+abbrev InteractiveMessage := TaggedText MsgEmbed
+
+instance : TypeName InteractiveMessage := unsafe (TypeName.mk _ ``InteractiveMessage)
+
 /-- The `message` field is the text of a message
 possibly containing interactive *embeds* of type `MsgEmbed`.
 We maintain the invariant that embeds are stored in `.tag`s with empty `.text` subtrees,
@@ -53,7 +57,7 @@ i.e., `.tag embed (.text "")`.
 
 Client-side display algorithms render tags in a custom way,
 ignoring the nested text. -/
-abbrev InteractiveDiagnostic := Lsp.DiagnosticWith (TaggedText MsgEmbed)
+abbrev InteractiveDiagnostic := Lsp.DiagnosticWith InteractiveMessage
 
 deriving instance RpcEncodable for Lsp.DiagnosticWith
 
@@ -63,8 +67,8 @@ open MsgEmbed
 partial def toDiagnostic (diag : InteractiveDiagnostic) : Lsp.Diagnostic :=
   { diag with message := prettyTt diag.message }
 where
-  prettyTt (tt : TaggedText MsgEmbed) : String :=
-    let tt : TaggedText MsgEmbed := tt.rewrite fun
+  prettyTt (tt : InteractiveMessage) : String :=
+    let tt : InteractiveMessage := tt.rewrite fun
       | .expr tt,      _ => .text tt.stripTags
       | .goal g,       _ => .text (toString g.pretty)
       | .widget _ alt, _ => .text $ prettyTt alt
@@ -78,7 +82,7 @@ private def mkPPContext (nCtx : NamingContext) (ctx : MessageDataContext) : PPCo
   currNamespace := nCtx.currNamespace, openDecls := nCtx.openDecls
 }
 
-/-! The `msgToInteractive` algorithm turns a `MessageData` into `TaggedText MsgEmbed` in two stages.
+/-! The `msgToInteractive` algorithm turns a `MessageData` into `InteractiveMessage` in two stages.
 
 First, in `msgToInteractiveAux` we produce a `Format` object whose `.tag` nodes refer to `EmbedFmt`
 objects stored in an auxiliary array. Only the most shallow `.tag` in every branch through the
@@ -86,11 +90,11 @@ objects stored in an auxiliary array. Only the most shallow `.tag` in every bran
 object (possibly including further `.tag`s), is processed. For example, if the output is
 `.tag (.expr ctx infos) fmt` then tags in the nested `fmt` object refer to elements of `infos`.
 
-In the second stage, we recursively transform such a `Format` into `TaggedText MsgEmbed` according
+In the second stage, we recursively transform such a `Format` into `InteractiveMessage` according
 to the rule above by first pretty-printing it and then grabbing data referenced by the tags from
 all the nested arrays (such as the `infos` array in the example above).
 
-We cannot easily do the translation in a single `MessageData → TaggedText MsgEmbed` step because
+We cannot easily do the translation in a single `MessageData → InteractiveMessage` step because
 that would effectively require reimplementing the (stateful, to keep track of indentation)
 `Format.prettyM` algorithm.
 -/
@@ -191,11 +195,11 @@ where
           f!"{children.size - blockSize} more entries..." more
     else children
 
-partial def msgToInteractive (msgData : MessageData) (hasWidgets : Bool) (indent : Nat := 0) : IO (TaggedText MsgEmbed) := do
+partial def msgToInteractive (msgData : MessageData) (hasWidgets : Bool) (indent : Nat := 0) : IO InteractiveMessage := do
   if !hasWidgets then
     return (TaggedText.prettyTagged (← msgData.format)).rewrite fun _ tt => .text tt.stripTags
   let (fmt, embeds) ← msgToInteractiveAux msgData
-  let rec fmtToTT (fmt : Format) (indent : Nat) : IO (TaggedText MsgEmbed) :=
+  let rec fmtToTT (fmt : Format) (indent : Nat) : IO InteractiveMessage :=
     (TaggedText.prettyTagged fmt indent).rewriteM fun (n, col) tt =>
       match embeds[n]! with
         | .code ctx infos => do