works

src/highlighting/SubVerso/Highlighting/Code.lean

@@ -21,6 +21,12 @@ initialize registerTraceClass `SubVerso.Highlighting.Code
 
 namespace SubVerso.Highlighting
 
+structure Context where
+  ids : HashMap Lsp.RefIdent Lsp.RefIdent
+  definitionsPossible : Bool
+
+def Context.noDefinitions (ctxt : Context) : Context := {ctxt with definitionsPossible := false}
+
 partial def Token.Kind.priority : Token.Kind → Nat
   | .var .. => 2
   | .str .. => 3
@@ -148,13 +154,13 @@ where
       go more
 
 def exprKind [Monad m] [MonadLiftT IO m] [MonadMCtx m] [MonadEnv m]
-    (ids : HashMap Lsp.RefIdent Lsp.RefIdent) (ci : ContextInfo) (lctx : LocalContext) (expr : Expr)
+    (ctxt : Context) (ci : ContextInfo) (lctx : LocalContext) (expr : Expr)
     : m (Option Token.Kind) := do
   let runMeta {α} (act : MetaM α) : m α := ci.runMetaM lctx act
   match ← instantiateMVars expr with
   | Expr.fvar id =>
     let seen ←
-      if let some y := ids[(← Compat.mkRefIdentFVar id)]? then
+      if let some y := ctxt.ids[(← Compat.mkRefIdentFVar id)]? then
         Compat.refIdentCase y
           (onFVar := fun x => do
             let ty ← instantiateMVars (← runMeta <| Meta.inferType expr)
@@ -183,6 +189,13 @@ def exprKind [Monad m] [MonadLiftT IO m] [MonadMCtx m] [MonadEnv m]
 
 /-- Checks whether an occurrence of a name is in fact the definition of the name -/
 def isDefinition [Monad m] [MonadEnv m] [MonadLiftT IO m] [MonadFileMap m] (name : Name) (stx : Syntax) : m Bool := do
+  -- This gets called a lot, so it's important to bail early if it's not likely to be a global
+  -- definition. Right now, all global definitions in Lean use `declId` to enable explicit universe
+  -- polymorphism, except for constructors. This does mean that `example` and `instance` won't work
+  -- yet, but they're a more marginal use case - there's no name to hyperlink to them in rendered
+  -- HTML.
+  if !((← getEnv).isConstructor name || (← getEnv).isProjectionFn name || stx.getKind == ``Parser.Command.declId) then return false
+  if stx.getHeadInfo == .none then return false
   let ranges :=
     if let some r := (← findDeclarationRangesCore? name) then
       some r
@@ -193,12 +206,13 @@ def isDefinition [Monad m] [MonadEnv m] [MonadLiftT IO m] [MonadFileMap m] (name
   return false
 
 def termInfoKind [Monad m] [MonadLiftT IO m] [MonadMCtx m] [MonadEnv m] [MonadFileMap m]
-    (ids : HashMap Lsp.RefIdent Lsp.RefIdent) (ci : ContextInfo) (termInfo : TermInfo)
+    (ctxt : Context) (ci : ContextInfo) (termInfo : TermInfo)
     : m (Option Token.Kind) := do
-  let k ← exprKind ids ci termInfo.lctx termInfo.expr
-  if let some (.const name sig docs isDef) := k then
-    dbg_trace termInfo.stx
-    return some (.const name sig docs isDef) --(← isDefinition name termInfo.stx))
+  let k ← exprKind ctxt ci termInfo.lctx termInfo.expr
+  if ctxt.definitionsPossible then
+    if let some (.const name sig docs _isDef) := k then
+      (some ∘ .const name sig docs) <$> (isDefinition name termInfo.stx)
+    else return k
   else return k
 
 def fieldInfoKind [Monad m] [MonadMCtx m] [MonadLiftT IO m] [MonadEnv m]
@@ -211,10 +225,10 @@ def fieldInfoKind [Monad m] [MonadMCtx m] [MonadLiftT IO m] [MonadEnv m]
   return .const fieldInfo.projName tyStr docs false
 
 def infoKind [Monad m] [MonadLiftT IO m] [MonadMCtx m] [MonadEnv m] [MonadFileMap m]
-     (ids : HashMap Lsp.RefIdent Lsp.RefIdent) (ci : ContextInfo) (info : Info)
+     (ctxt : Context) (ci : ContextInfo) (info : Info)
     : m (Option Token.Kind) := do
   match info with
-    | .ofTermInfo termInfo => termInfoKind ids ci termInfo
+    | .ofTermInfo termInfo => termInfoKind ctxt ci termInfo
     | .ofFieldInfo fieldInfo => some <$> fieldInfoKind ci fieldInfo
     | .ofOptionInfo oi =>
       let doc := (← getOptionDecls).find? oi.optionName |>.map (·.descr)
@@ -226,24 +240,24 @@ def infoKind [Monad m] [MonadLiftT IO m] [MonadMCtx m] [MonadEnv m] [MonadFileMa
       pure none
 
 def identKind [Monad m] [MonadLiftT IO m] [MonadFileMap m] [MonadEnv m] [MonadMCtx m]
-    (ids : HashMap Lsp.RefIdent Lsp.RefIdent) (trees : PersistentArray InfoTree) (stx : TSyntax `ident)
+    (ctxt : Context) (trees : PersistentArray InfoTree) (stx : TSyntax `ident)
     : m Token.Kind := do
   let mut kind : Token.Kind := .unknown
   for t in trees do
     for (ci, info) in infoForSyntax t stx do
-      if let some seen ← infoKind ids ci info then
+      if let some seen ← infoKind ctxt ci info then
         if seen.priority > kind.priority then kind := seen
       else continue
 
   pure kind
 
 def anonCtorKind [Monad m] [MonadLiftT IO m] [MonadFileMap m] [MonadEnv m] [MonadMCtx m]
-    (ids : HashMap Lsp.RefIdent Lsp.RefIdent) (trees : PersistentArray InfoTree) (stx : Syntax)
+    (ctxt : Context) (trees : PersistentArray InfoTree) (stx : Syntax)
     : m (Option Token.Kind) := do
   let mut kind : Token.Kind := .unknown
   for t in trees do
     for (ci, info) in infoForSyntax t stx do
-      if let some seen ← infoKind ids ci info then
+      if let some seen ← infoKind ctxt ci info then
         if seen.priority > kind.priority then kind := seen
       else continue
   return match kind with
@@ -540,7 +554,7 @@ partial def childHasTactics (stx : Syntax) (trees : PersistentArray Lean.Elab.In
 
 
 partial def renderTagged [Monad m] [MonadLiftT IO m] [MonadMCtx m] [MonadEnv m] [MonadFileMap m]
-    (ids : HashMap Lsp.RefIdent Lsp.RefIdent) (outer : Option Token.Kind) (doc : CodeWithInfos)
+    (ctxt : Context) (outer : Option Token.Kind) (doc : CodeWithInfos)
     : m Highlighted := do
   match doc with
   | .text txt => do
@@ -578,13 +592,13 @@ partial def renderTagged [Monad m] [MonadLiftT IO m] [MonadMCtx m] [MonadEnv m]
   | .tag t doc' =>
     let ⟨{ctx, info, children := _}⟩ := t.info
     if let .text tok := doc' then
-      if let some k ← infoKind ids ctx info then
+      if let some k ← infoKind ctxt ctx info then
         pure <| .token ⟨k, tok⟩
       else pure <| .text tok
     else
-      let k? ← infoKind ids ctx info
-      renderTagged ids k? doc'
-  | .append xs => .seq <$> xs.mapM (renderTagged ids outer)
+      let k? ← infoKind ctxt ctx info
+      renderTagged ctxt k? doc'
+  | .append xs => .seq <$> xs.mapM (renderTagged ctxt outer)
 where
   tokenEnder str := str.isEmpty || !(str.get 0 |>.isAlphanum)
 
@@ -602,10 +616,11 @@ partial def _root_.Lean.Widget.TaggedText.indent (doc : TaggedText α) : TaggedT
   .append #[.text "  ", indent' doc]
 
 def highlightGoals
-    (ids : HashMap Lsp.RefIdent Lsp.RefIdent)
+    (ctxt : Context)
     (ci : ContextInfo)
     (goals : List MVarId)
     : HighlightM (Array (Highlighted.Goal Highlighted)) := do
+  let ctxt := ctxt.noDefinitions
   let mut goalView := #[]
   for g in goals do
     let mut hyps := #[]
@@ -625,23 +640,23 @@ def highlightGoals
       if decl.isAuxDecl || decl.isImplementationDetail then continue
       match decl with
       | .cdecl _index fvar name type _ _ =>
-        let nk ← exprKind ids ci lctx (.fvar fvar)
-        let tyStr ← renderTagged ids none (← runMeta (ppExprTagged =<< instantiateMVars type))
+        let nk ← exprKind ctxt ci lctx (.fvar fvar)
+        let tyStr ← renderTagged ctxt none (← runMeta (ppExprTagged =<< instantiateMVars type))
         hyps := hyps.push (name, nk.getD .unknown, tyStr)
       | .ldecl _index fvar name type val _ _ =>
-        let nk ← exprKind ids ci lctx (.fvar fvar)
+        let nk ← exprKind ctxt ci lctx (.fvar fvar)
         let tyDoc ← runMeta (ppExprTagged =<< instantiateMVars type)
         let valDoc ← runMeta (ppExprTagged =<< instantiateMVars val)
-        let tyValStr ← renderTagged ids none <| .append <| #[tyDoc].append <|
+        let tyValStr ← renderTagged ctxt none <| .append <| #[tyDoc].append <|
           if tyDoc.oneLine && valDoc.oneLine then #[.text " := ", valDoc]
           else #[.text " := \n", valDoc.indent]
         hyps := hyps.push (name, nk.getD .unknown, tyValStr)
-    let concl ← renderTagged ids none (← runMeta <| ppExprTagged =<< instantiateMVars mvDecl.type)
+    let concl ← renderTagged ctxt none (← runMeta <| ppExprTagged =<< instantiateMVars mvDecl.type)
     goalView := goalView.push ⟨name, Meta.getGoalPrefix mvDecl, hyps, concl⟩
   pure goalView
 
 partial def findTactics'
-    (ids : HashMap Lsp.RefIdent Lsp.RefIdent)
+    (ctxt : Context)
     (trees : PersistentArray Lean.Elab.InfoTree)
     (stx : Syntax)
     (startPos endPos : String.Pos)
@@ -660,14 +675,14 @@ partial def findTactics'
 
         let goals := if before then tacticInfo.goalsBefore else tacticInfo.goalsAfter
         if goals.isEmpty then continue
-        let goalView ← highlightGoals ids ci goals
+        let goalView ← highlightGoals ctxt ci goals
 
         if !Output.inTacticState (← get).output goalView then
           return some (goalView, startPos.byteIdx, endPos.byteIdx, endPosition)
   return none
 
 partial def findTactics
-    (ids : HashMap Lsp.RefIdent Lsp.RefIdent)
+    (ctxt : Context)
     (trees : PersistentArray Lean.Elab.InfoTree)
     (stx : Syntax)
     : HighlightM (Option (Array (Highlighted.Goal Highlighted) × Nat × Nat × Position)) :=
@@ -693,7 +708,7 @@ partial def findTactics
           | `(Lean.Parser.Term.byTactic| by%$tk $tactics)
           | `(Lean.Parser.Term.byTactic'| by%$tk $tactics) =>
             if tk == stx then
-              return (← findTactics' ids trees tactics startPos endPos endPosition (before := true))
+              return (← findTactics' ctxt trees tactics startPos endPos endPosition (before := true))
           | _ => continue
 
   -- Special handling for =>: show the _before state_
@@ -708,7 +723,7 @@ partial def findTactics
           | `(tactic| case $_ $_* =>%$tk $rhs )
           | `(tactic| case' $_ $_* =>%$tk $rhs ) =>
             if tk == stx then
-              return (← findTactics' ids trees rhs startPos endPos endPosition (before := true))
+              return (← findTactics' ctxt trees rhs startPos endPos endPosition (before := true))
           | _ => continue
 
   -- Only show tactic output for the most specific source spans possible, with a few exceptions
@@ -718,14 +733,14 @@ partial def findTactics
   -- Override states - some tactics show many intermediate states, which is overwhelming in rendered
   -- output. Get the right one to show for the whole thing, then adjust its positioning.
   if let some brak := Compat.rwTacticRightBracket? stx then
-    if let some (goals, _startPos, _endPos, _endPosition) ← findTactics ids trees brak then
+    if let some (goals, _startPos, _endPos, _endPosition) ← findTactics ctxt trees brak then
       return some (goals, startPos.byteIdx, endPos.byteIdx, endPosition)
 
-  findTactics' ids trees stx startPos endPos endPosition
+  findTactics' ctxt trees stx startPos endPos endPosition
 
 
 partial def highlight'
-    (ids : HashMap Lsp.RefIdent Lsp.RefIdent)
+    (ctxt : Context)
     (trees : PersistentArray Lean.Elab.InfoTree)
     (stx : Syntax)
     (tactics : Bool)
@@ -734,20 +749,20 @@ partial def highlight'
   withTraceNode `SubVerso.Highlighting.Code (fun _ => pure m!"Highlighting {stx}") do
   let mut tactics := tactics
   if tactics then
-    if let some (tacticInfo, startPos, endPos, position) ← findTactics ids trees stx then
+    if let some (tacticInfo, startPos, endPos, position) ← findTactics ctxt trees stx then
       HighlightM.openTactic tacticInfo startPos endPos position
       -- No nested tactics - the tactic search process should only have returned results
       -- on "leaf" nodes anyway
       tactics := false
   match stx with
   | `($e.%$tk$field:ident) =>
       withTraceNode `SubVerso.Highlighting.Code (fun _ => pure m!"Highlighting projection {e} {tk} {field}") do
-      highlight' ids trees e tactics
+      highlight' ctxt trees e tactics
       if let some ⟨pos, endPos⟩ := tk.getRange? then
         emitToken tk.getHeadInfo <| .mk .unknown <| (← getFileMap).source.extract pos endPos
       else
         emitString' "."
-      highlight' ids trees field tactics
+      highlight' ctxt trees field tactics
   | _ =>
     match stx with
     | .missing => pure () -- TODO emit unhighlighted string
@@ -764,20 +779,20 @@ partial def highlight'
             withTraceNode `SubVerso.Highlighting.Code (fun _ => pure m!"Perhaps a field? {y} {field}") do
             if (← infoExists trees field) then
               withTraceNode `SubVerso.Highlighting.Code (fun _ => pure m!"Yes, a field!") do
-              highlight' ids trees y tactics
+              highlight' ctxt trees y tactics
               emitToken' <| fakeToken .unknown "."
-              highlight' ids trees field tactics
+              highlight' ctxt trees field tactics
             else
               withTraceNode `SubVerso.Highlighting.Code (fun _ => pure m!"Not a field.") do
-              let k ← identKind ids trees ⟨stx⟩
+              let k ← identKind ctxt trees ⟨stx⟩
               withTraceNode `SubVerso.Highlighting.Code (fun _ => pure m!"Identifier token for {stx} is {repr k}") do
               emitToken i ⟨k, x.toString⟩
           | _ =>
-            let k ← identKind ids trees ⟨stx⟩
+            let k ← identKind ctxt trees ⟨stx⟩
             withTraceNode `SubVerso.Highlighting.Code (fun _ => pure m!"Identifier token for {stx} is {repr k}") do
             emitToken i ⟨k, x.toString⟩
         | _ =>
-          let k ← identKind ids trees ⟨stx⟩
+          let k ← identKind ctxt trees ⟨stx⟩
           withTraceNode `SubVerso.Highlighting.Code (fun _ => pure m!"Identifier token for {stx} is {repr k}") do
           emitToken i ⟨k, x.toString⟩
     | stx@(.atom i x) =>
@@ -787,7 +802,7 @@ partial def highlight'
         | some (n, _) => findDocString? (← getEnv) n
       let name := lookingAt.map (·.1)
       let occ := lookingAt.map fun (n, pos) => s!"{n}-{pos}"
-      if let .sort ← identKind ids trees ⟨stx⟩ then
+      if let .sort ← identKind ctxt trees ⟨stx⟩ then
         withTraceNode `SubVerso.Highlighting.Code (fun _ => pure m!"Sort") do
           emitToken i ⟨.sort, x⟩
         return
@@ -822,48 +837,48 @@ partial def highlight'
       match i, i' with
       | .original leading pos _ _, .original _ _ trailing endPos =>
         let info := .original leading pos trailing endPos
-        emitToken info ⟨← identKind ids trees ⟨stx⟩, s!".{x.toString}"⟩
+        emitToken info ⟨← identKind ctxt trees ⟨stx⟩, s!".{x.toString}"⟩
       | _, _ =>
-        highlight' ids trees dot tactics
-        highlight' ids trees name tactics
+        highlight' ctxt trees dot tactics
+        highlight' ctxt trees name tactics
     | .node info k@``Lean.Parser.Term.anonymousCtor #[opener@(.atom oi l), children@(.node _ _ contents), closer@(.atom ci r)] =>
-      if let some tk ← anonCtorKind ids trees stx then
+      if let some tk ← anonCtorKind ctxt trees stx then
         emitToken oi ⟨tk, l⟩
         for child in contents do
           match child with
           | .atom commaInfo "," =>
             emitToken commaInfo ⟨tk, ","⟩
           | _ =>
-            highlight' ids trees child tactics
+            highlight' ctxt trees child tactics
         emitToken ci ⟨tk, r⟩
       else
         let pos := stx.getPos?
-        highlight' ids trees opener tactics (lookingAt := pos.map (k, ·))
-        highlight' ids trees children tactics (lookingAt := pos.map (k, ·))
-        highlight' ids trees closer tactics (lookingAt := pos.map (k, ·))
+        highlight' ctxt trees opener tactics (lookingAt := pos.map (k, ·))
+        highlight' ctxt trees children tactics (lookingAt := pos.map (k, ·))
+        highlight' ctxt trees closer tactics (lookingAt := pos.map (k, ·))
     | .node _ `choice alts =>
       withTraceNode `SubVerso.Highlighting.Code (fun _ => pure m!"Choice node with {alts.size} alternatives") do
       -- Arbitrarily select one of the alternatives found by the parser. Otherwise, quotations of
       -- let-bindings with antiquotations as the bound variable leads to doubled bound variables,
       -- because the parser emits a choice node in the quotation. And it's never correct to show
       -- both alternatives!
       if h : alts.size > 0 then
-        highlight' ids trees alts[0]  tactics
+        highlight' ctxt trees alts[0]  tactics
     | stx@(.node _ k children) =>
       withTraceNode `SubVerso.Highlighting.Code (fun _ => pure m!"Other node, kind {k}, with {children.size} children") do
       let pos := stx.getPos?
       for child in children do
-        highlight' ids trees child tactics (lookingAt := pos.map (k, ·))
+        highlight' ctxt trees child tactics (lookingAt := pos.map (k, ·))
 
 def highlight (stx : Syntax) (messages : Array Message) (trees : PersistentArray Lean.Elab.InfoTree) : TermElabM Highlighted := do
   let modrefs := Lean.Server.findModuleRefs (← getFileMap) trees.toArray
   let ids := build modrefs
   let st ← HighlightState.ofMessages stx messages
-  let ((), {output := output, ..}) ← StateT.run (highlight' ids trees stx true) st
+  let ((), {output := output, ..}) ← StateT.run (highlight' ⟨ids, true⟩ trees stx true) st
   pure <| .fromOutput output
 
 def highlightProofState (ci : ContextInfo) (goals : List MVarId) (trees : PersistentArray Lean.Elab.InfoTree) : TermElabM (Array (Highlighted.Goal Highlighted)) := do
   let modrefs := Lean.Server.findModuleRefs (← getFileMap) trees.toArray
   let ids := build modrefs
-  let (hlGoals, _) ← StateT.run (highlightGoals ids ci goals) .empty
+  let (hlGoals, _) ← StateT.run (highlightGoals ⟨ids, false⟩ ci goals) .empty
   pure hlGoals

src/highlighting/SubVerso/Highlighting/Highlighted.lean

@@ -120,6 +120,7 @@ inductive Highlighted where
 deriving Repr, Inhabited, BEq, Hashable, ToJson, FromJson
 
 def Highlighted.empty : Highlighted := .seq #[]
+
 instance : Append Highlighted where
   append
     | .text str1, .text str2 => .text (str1 ++ str2)
@@ -128,6 +129,15 @@ instance : Append Highlighted where
     | x, .seq xs => .seq (#[x] ++ xs)
     | x, y => .seq #[x, y]
 
+partial def Highlighted.definedNames : Highlighted → NameSet
+  | .token ⟨tok, _⟩ =>
+    match tok with
+    | .const n _ _ true => NameSet.empty.insert n
+    | _ => {}
+  | .span _ hl | .tactics _ _ _ hl => hl.definedNames
+  | .seq hls => hls.map (·.definedNames) |>.foldr NameSet.append {}
+  | .text .. | .point .. => {}
+
 open Lean Syntax in
 open Highlighted in
 partial instance : Quote Highlighted where