feat: improve semantic highlighting

src/Lean/Server/FileWorker/SemanticHighlighting.lean

@@ -817,10 +817,11 @@ without elaboration information.
 partial def collectSyntaxBasedSemanticTokens (text : FileMap) : (stx : Syntax) → Array LeanSemanticToken
   | `($e.$id:ident)    =>
     let tokens := collectSyntaxBasedSemanticTokens text e
-    emitSyntaxToken tokens id .property
+    -- Prefer the elaborated declaration kind over this syntactic fallback.
+    emitSyntaxToken tokens id .property (priority := 4)
   | `($e |>.$field:ident) =>
     let tokens := collectSyntaxBasedSemanticTokens text e
-    emitSyntaxToken tokens field .property
+    emitSyntaxToken tokens field .property (priority := 4)
   | stx => Id.run do
     if noHighlightKinds.contains stx.getKind then
       return #[]
@@ -862,25 +863,136 @@ partial def collectSyntaxBasedSemanticTokens (text : FileMap) : (stx : Syntax) 
       return tokens
     return emitSyntaxToken tokens stx (keywordSemanticTokenMap.getD val .keyword)
 
+private def typeForSortResult? (type : Expr) : Option SemanticTokenType :=
+  (Meta.typeFormerTypeLevelQuick type).map fun u =>
+    if u.isAlwaysZero then .enum else .type
+
+private partial def typeExprResult? (env : Environment) (lctx : LocalContext) (type : Expr)
+    (bvarTypes : List Expr := []) : Option SemanticTokenType :=
+  match type.getAppFn with
+  | .sort _ => typeForSortResult? type
+  | .forallE _ domain body _ => typeExprResult? env lctx body (domain :: bvarTypes)
+  | .const declName us =>
+    match env.find? declName with
+    | some info => typeForSortResult? (info.instantiateTypeLevelParams us)
+    | none => none
+  | .fvar fvarId =>
+    match lctx.find? fvarId with
+    | some localDecl => typeForSortResult? localDecl.type
+    | none => none
+  | .bvar idx => bvarTypes[idx]?.bind typeForSortResult?
+  | _ => none
+
+private def classifyLocal (env : Environment) (ti : Elab.TermInfo) (localDecl : LocalDecl) :
+    Option SemanticTokenType := Id.run do
+  let localType := localDecl.type
+  -- Recall that `isAuxDecl` is used for declarations currently being elaborated, e.g. recursive
+  -- functions and inductive types while their constructors are being processed.
+  if localDecl.isAuxDecl then
+    if let some tokenType := typeForSortResult? localType then
+      return some tokenType
+    else
+      return some .function
+  if let some tokenType := typeForSortResult? localType then
+    if localType.isSort && tokenType == .type then
+      return some .typeParameter
+    else
+      return some tokenType
+  if localDecl.isImplementationDetail then
+    return none
+  if typeExprResult? env ti.lctx localType == some .enum then
+    return some .variable
+  if ti.isBinder then
+    return some .parameter
+  return some .variable
+
+private def classifyConstant (env : Environment) (info : ConstantInfo) (us : List Level := []) :
+    SemanticTokenType :=
+  if env.isProjectionFn info.name then
+    .property
+  else match info with
+  | .ctorInfo _ => .enumMember
+  | .inductInfo _ =>
+    if isClass env info.name then
+      .interface
+    else if isStructure env info.name then
+      .struct
+    else if let some tokenType := typeForSortResult? (info.instantiateTypeLevelParams us) then
+      tokenType
+    else
+      .type
+  | .thmInfo _ | .recInfo _ => .function
+  | _ =>
+    if let some tokenType := typeForSortResult? (info.instantiateTypeLevelParams us) then
+      tokenType
+    else
+      .function
+
+private def findConstInfoWithEnv? (ctx : Elab.ContextInfo) (declName : Name) :
+    Option (Environment × ConstantInfo) :=
+  match ctx.env.find? declName with
+  | some info => some (ctx.env, info)
+  | none => ctx.cmdEnv?.bind fun env => (env.find? declName).map (env, ·)
+
+private def getSemanticTokenSyntax (stx : Syntax) : Syntax :=
+  if stx.isOfKind Parser.Term.identProjKind then
+    stx[0]
+  else match stx with
+  | `(Parser.Command.declId|$id$[.{$_ls,*}]?) => id.raw
+  | `(.$id:ident) | `(.$id:ident.{$_us,*}) => id.raw
+  | _ => stx
+
+private def getAppFnSyntax? (stx : Syntax) : Option Syntax :=
+  if stx.isOfKind ``Parser.Term.app then some stx[0] else none
+
+private def collectTermInfoBasedSemanticToken (ctx : Elab.ContextInfo) (ti : Elab.TermInfo) :
+    Option LeanSemanticToken := Id.run do
+  let stx := getSemanticTokenSyntax ti.stx
+  if ti.expr.isSort then
+    unless stx.isOfKind ``Parser.Term.type || stx.isOfKind ``Parser.Term.sort || stx.isOfKind ``Parser.Term.prop do
+      return none
+    let stx := stx[0]
+    let .original .. := stx.getHeadInfo
+      | return none
+    return LeanSemanticToken.ofSyntax? stx .class
+  match ti.expr.getAppFn with
+  | Expr.fvar fvarId .. =>
+    let some localDecl := ti.lctx.find? fvarId
+      | return none
+    let some stx :=
+      if let `($_:ident) := stx then
+        some stx
+      else if localDecl.isAuxDecl then
+        getAppFnSyntax? stx
+      else
+        none
+      | return none
+    let .original .. := stx.getHeadInfo
+      | return none
+    let `($_:ident) := stx
+      | return none
+    if localDecl.isAuxDecl then
+      if let some declName := ti.lctx.auxDeclToFullName.get? fvarId then
+        let env := ctx.cmdEnv?.getD ctx.env
+        if let some info := env.find? declName then
+          return LeanSemanticToken.ofSyntax? stx (classifyConstant env info)
+    return (classifyLocal ctx.env ti localDecl).bind (LeanSemanticToken.ofSyntax? stx ·)
+  | Expr.const declName us =>
+    let `($_:ident) := stx
+      | return none
+    let .original .. := stx.getHeadInfo
+      | return none
+    match findConstInfoWithEnv? ctx declName with
+    | some (env, info) => return LeanSemanticToken.ofSyntax? stx (classifyConstant env info us)
+    | none => return none
+  | _ => return none
+
 /-- Collects all semantic tokens from the given `Elab.InfoTree`. -/
 def collectInfoBasedSemanticTokens (i : Elab.InfoTree) : Array LeanSemanticToken :=
-  List.toArray <| i.deepestNodes fun _ info _ => do
+  List.toArray <| i.deepestNodes fun ctx info _ => do
     let .ofTermInfo ti := info
       | none
-    let .original .. := ti.stx.getHeadInfo
-      | none
-    if let `($_:ident) := ti.stx then
-      if let Expr.fvar fvarId .. := ti.expr then
-        if let some localDecl := ti.lctx.find? fvarId then
-          -- Recall that `isAuxDecl` is an auxiliary declaration used to elaborate a recursive definition.
-          if localDecl.isAuxDecl then
-            if ti.isBinder then
-              return (← LeanSemanticToken.ofSyntax? ti.stx .function)
-          else if ! localDecl.isImplementationDetail then
-            return (← LeanSemanticToken.ofSyntax? ti.stx .variable)
-    if ti.stx.getKind == Parser.Term.identProjKind then
-      return (← LeanSemanticToken.ofSyntax? ti.stx .property)
-    none
+    collectTermInfoBasedSemanticToken ctx ti
 
 /--
 A debugging utility for inspecting sets of collected tokens, classified by line and sorted by

tests/server_interactive/semanticTokens.lean

@@ -1,8 +1,31 @@
-def foo1 (x : Nat) := x.succ
-def foo2 (x : Nat) := x |>.succ
-theorem foo3 (x : Nat) : True :=
+/-! This file tests semantic token kinds. -/
+
+def dotNotation (x : Nat) := x.succ
+def dottedConstructor (x : Nat) : List Nat := .cons.{0} (x |>.succ) .nil
+theorem withLocalLet (x : Nat) : True :=
   let y := x.succ
-  True.intro
+  .intro
+
+universe u
+def TypeAlias := Type
+  0
+def SortAlias := Sort (u + 1)
+abbrev NatAlias := Nat
+def InferredList (α : Type) := List α
+def ExplicitList (α : Type) : Type := InferredList α
+def singletonNat (x : NatAlias) : ExplicitList NatAlias := [x]
+
+structure RecursiveStruct where
+  xs : List RecursiveStruct
+class RecursiveClass (α : Type) where
+  xs : List (RecursiveClass α)
+
+theorem withHypotheses (p q : Prop) (h1 : p ∧ q) (h2 : p → q)
+    (h3 : ∀ x : Nat, x = x) : q ∧ 0 = 0 :=
+  ⟨h2 h1.1, h3 0⟩
+theorem withPolymorphicHypothesis (h : ∀ p : Prop, p → p) (f : ∀ α : Type, α → α) :
+    f Nat 0 = f Nat 0 :=
+  h _ rfl
 
 #eval 1
 --^ collectDiagnostics