Make go-to-definition on a typeclass projection application go to the instance(s)

src/Lean/Elab/Extra.lean

@@ -138,16 +138,17 @@ private inductive Tree where
   | term (ref : Syntax) (infoTrees : PersistentArray InfoTree) (val : Expr)
   /--
   `ref` is the original syntax that expanded into `binop%`.
-  `macroName` is the `macro_rule` that produce the expansion. We store this information
-  here to make sure "go to definition" behaves similarly to notation defined without using `binop%` helper elaborator.
   -/
-  | binop (ref : Syntax) (macroName : Name) (lazy : Bool) (f : Expr) (lhs rhs : Tree)
+  | binop (ref : Syntax) (lazy : Bool) (f : Expr) (lhs rhs : Tree)
   /--
   `ref` is the original syntax that expanded into `unop%`.
-  `macroName` is the `macro_rule` that produce the expansion. We store this information
-  here to make sure "go to definition" behaves similarly to notation defined without using `unop%` helper elaborator.
   -/
-  | unop (ref : Syntax) (macroName : Name) (f : Expr) (arg : Tree)
+  | unop (ref : Syntax) (f : Expr) (arg : Tree)
+  /--
+  Used for assembling the info tree. We store this information
+  to make sure "go to definition" behaves similarly to notation defined without using `binop%` helper elaborator.
+  -/
+  | macroExpansion (macroName : Name) (stx stx' : Syntax) (nested : Tree)
 
 
 private partial def toTree (s : Syntax) : TermElabM Tree := do
@@ -163,32 +164,30 @@ private partial def toTree (s : Syntax) : TermElabM Tree := do
 where
   go (s : Syntax) := do
     match s with
-    | `(binop% $f $lhs $rhs) => processBinOp (lazy := false) s .anonymous f lhs rhs
-    | `(binop_lazy% $f $lhs $rhs) => processBinOp (lazy := true) s .anonymous f lhs rhs
-    | `(unop% $f $arg) => processUnOp s .anonymous f arg
+    | `(binop% $f $lhs $rhs) => processBinOp (lazy := false) s f lhs rhs
+    | `(binop_lazy% $f $lhs $rhs) => processBinOp (lazy := true) s f lhs rhs
+    | `(unop% $f $arg) => processUnOp s f arg
     | `(($e)) =>
       if hasCDot e then
         processLeaf s
       else
         go e
     | _ =>
-       match (← liftMacroM <| expandMacroImpl? (← getEnv) s) with
-       | some (macroName, s?) =>
-         let s' ← liftMacroM <| liftExcept s?
-         match s' with
-         | `(binop% $f $lhs $rhs) => processBinOp (lazy := false) s macroName f lhs rhs
-         | `(binop_lazy% $f $lhs $rhs) => processBinOp (lazy := true) s macroName f lhs rhs
-         | `(unop% $f $arg) => processUnOp s .anonymous f arg
-         | _  => processLeaf s
-       | none => processLeaf s
-
-  processBinOp (ref : Syntax) (declName : Name) (f lhs rhs : Syntax) (lazy : Bool) := do
+      withRef s do
+        match (← liftMacroM <| expandMacroImpl? (← getEnv) s) with
+        | some (macroName, s?) =>
+          let s' ← liftMacroM <| liftExcept s?
+          withPushMacroExpansionStack s s' do
+            return .macroExpansion macroName s s' (← go s')
+        | none => processLeaf s
+
+  processBinOp (ref : Syntax) (f lhs rhs : Syntax) (lazy : Bool) := do
     let some f ← resolveId? f | throwUnknownConstant f.getId
-    return .binop (lazy := lazy) ref declName f (← go lhs) (← go rhs)
+    return .binop (lazy := lazy) ref f (← go lhs) (← go rhs)
 
-  processUnOp (ref : Syntax) (declName : Name) (f arg : Syntax) := do
+  processUnOp (ref : Syntax) (f arg : Syntax) := do
     let some f ← resolveId? f | throwUnknownConstant f.getId
-    return .unop ref declName f (← go arg)
+    return .unop ref f (← go arg)
 
   processLeaf (s : Syntax) := do
     let e ← elabTerm s none
@@ -229,8 +228,9 @@ where
    go (t : Tree) : StateRefT AnalyzeResult TermElabM Unit := do
      unless (← get).hasUncomparable do
        match t with
-       | .binop _ _ _ _ lhs rhs => go lhs; go rhs
-       | .unop _ _ _ arg => go arg
+       | .macroExpansion _ _ _ nested => go nested
+       | .binop _ _ _ lhs rhs => go lhs; go rhs
+       | .unop _ _ arg => go arg
        | .term _ _ val =>
          let type ← instantiateMVars (← inferType val)
          unless isUnknow type do
@@ -256,15 +256,20 @@ private def toExprCore (t : Tree) : TermElabM Expr := do
   match t with
   | .term _ trees e =>
     modifyInfoState (fun s => { s with trees := s.trees ++ trees }); return e
-  | .binop ref macroName true f lhs rhs  =>
-    withRef ref <| withInfoContext' ref (mkInfo := mkTermInfo macroName ref) do
-      mkBinOp f (← toExprCore lhs) (← mkFunUnit (← toExprCore rhs))
-  | .binop ref macroName false f lhs rhs =>
-    withRef ref <| withInfoContext' ref (mkInfo := mkTermInfo macroName ref) do
-      mkBinOp f (← toExprCore lhs) (← toExprCore rhs)
-  | .unop ref macroName f arg =>
-    withRef ref <| withInfoContext' ref (mkInfo := mkTermInfo macroName ref) do
+  | .binop ref lazy f lhs rhs =>
+    withRef ref <| withInfoContext' ref (mkInfo := mkTermInfo .anonymous ref) do
+      let lhs ← toExprCore lhs
+      let mut rhs ← toExprCore rhs
+      if lazy then
+        rhs ← mkFunUnit rhs
+      mkBinOp f lhs rhs
+  | .unop ref f arg =>
+    withRef ref <| withInfoContext' ref (mkInfo := mkTermInfo .anonymous ref) do
       mkUnOp f (← toExprCore arg)
+  | .macroExpansion macroName stx stx' nested =>
+    withRef stx <| withInfoContext' stx (mkInfo := mkTermInfo macroName stx) do
+      withMacroExpansion stx stx' do
+        toExprCore nested
 
 /--
   Auxiliary function to decide whether we should coerce `f`'s argument to `maxType` or not.
@@ -343,22 +348,22 @@ mutual
   where
     go (t : Tree) (f? : Option Expr) (lhs : Bool) (isPred : Bool) : TermElabM Tree := do
       match t with
-      | .binop ref macroName lazy f lhs rhs =>
+      | .binop ref lazy f lhs rhs =>
         /-
           We only keep applying coercions to `maxType` if `f` is predicate or
           `f` has a homogenous instance with `maxType`. See `hasHomogeneousInstance` for additional details.
 
           Remark: We assume `binrel%` elaborator is only used with homogenous predicates.
         -/
         if (← pure isPred <||> hasHomogeneousInstance f maxType) then
-          return .binop ref macroName lazy f (← go lhs f true false) (← go rhs f false false)
+          return .binop ref lazy f (← go lhs f true false) (← go rhs f false false)
         else
-          let r ← withRef ref <| withInfoContext' ref (mkInfo := mkTermInfo macroName ref) do
+          let r ← withRef ref do
             mkBinOp f (← toExpr lhs none) (← toExpr rhs none)
           let infoTrees ← getResetInfoTrees
           return .term ref infoTrees r
-      | .unop ref macroName f arg =>
-        return .unop ref macroName f (← go arg none false false)
+      | .unop ref f arg =>
+        return .unop ref f (← go arg none false false)
       | .term ref trees e =>
         let type ← instantiateMVars (← inferType e)
         trace[Elab.binop] "visiting {e} : {type} =?= {maxType}"
@@ -372,6 +377,9 @@ mutual
         else
           trace[Elab.binop] "added coercion: {e} : {type} => {maxType}"
           withRef ref <| return .term ref trees (← mkCoe maxType e)
+      | .macroExpansion macroName stx stx' nested =>
+        withRef stx <| withPushMacroExpansionStack stx stx' do
+          return .macroExpansion macroName stx stx' (← go nested f? lhs isPred)
 
   private partial def toExpr (tree : Tree) (expectedType? : Option Expr) : TermElabM Expr := do
     let r ← analyze tree expectedType?
@@ -441,7 +449,7 @@ def elabBinRelCore (noProp : Bool) (stx : Syntax) (expectedType? : Option Expr)
     -/
     let lhs ← withRef stx[2] <| toTree stx[2]
     let rhs ← withRef stx[3] <| toTree stx[3]
-    let tree := .binop (lazy := false) stx .anonymous f lhs rhs
+    let tree := .binop (lazy := false) stx f lhs rhs
     let r ← analyze tree none
     trace[Elab.binrel] "hasUncomparable: {r.hasUncomparable}, maxType: {r.max?}"
     if r.hasUncomparable || r.max?.isNone then

src/Lean/Elab/InfoTree/Main.lean

@@ -110,7 +110,7 @@ def TermInfo.format (ctx : ContextInfo) (info : TermInfo) : IO Format := do
 def CompletionInfo.format (ctx : ContextInfo) (info : CompletionInfo) : IO Format :=
   match info with
   | .dot i (expectedType? := expectedType?) .. => return f!"[.] {← i.format ctx} : {expectedType?}"
-  | .id stx _ _ lctx expectedType? => ctx.runMetaM lctx do return f!"[.] {stx} : {expectedType?} @ {formatStxRange ctx info.stx}"
+  | .id stx _ _ lctx expectedType? => ctx.runMetaM lctx do return f!"[.] {← ctx.ppSyntax lctx stx} : {expectedType?} @ {formatStxRange ctx info.stx}"
   | _ => return f!"[.] {info.stx} @ {formatStxRange ctx info.stx}"
 
 def CommandInfo.format (ctx : ContextInfo) (info : CommandInfo) : IO Format := do

src/Lean/Elab/Term.lean

@@ -457,9 +457,13 @@ def elabLevel (stx : Syntax) : TermElabM Level :=
   liftLevelM <| Level.elabLevel stx
 
 /-- Elaborate `x` with `stx` on the macro stack -/
+def withPushMacroExpansionStack (beforeStx afterStx : Syntax) (x : TermElabM α) : TermElabM α :=
+  withReader (fun ctx => { ctx with macroStack := { before := beforeStx, after := afterStx } :: ctx.macroStack }) x
+
+/-- Elaborate `x` with `stx` on the macro stack and produce macro expansion info -/
 def withMacroExpansion (beforeStx afterStx : Syntax) (x : TermElabM α) : TermElabM α :=
   withMacroExpansionInfo beforeStx afterStx do
-    withReader (fun ctx => { ctx with macroStack := { before := beforeStx, after := afterStx } :: ctx.macroStack }) x
+    withPushMacroExpansionStack beforeStx afterStx x
 
 /--
   Add the given metavariable to the list of pending synthetic metavariables.

src/Lean/Server/FileWorker/RequestHandling.lean

@@ -66,11 +66,11 @@ def handleHover (p : HoverParams)
         | none => pure none
 
       -- now try info tree
-      if let some (ci, i) := snap.infoTree.hoverableInfoAt? hoverPos then
-        if let some range := i.range? then
+      if let some ictx := snap.infoTree.hoverableInfoAt? hoverPos then
+        if let some range := ictx.info.range? then
           -- prefer info tree if at least as specific as parser docstring
           if stxDoc?.all fun (_, stxRange) => stxRange.includes range then
-            if let some hoverFmt ← i.fmtHover? ci then
+            if let some hoverFmt ← ictx.info.fmtHover? ictx.ctx then
               return mkHover (toString hoverFmt.fmt) range
 
       if let some (doc, range) := stxDoc? then
@@ -79,7 +79,7 @@ def handleHover (p : HoverParams)
       return none
 
 open Elab GoToKind in
-def locationLinksOfInfo (kind : GoToKind) (ci : Elab.ContextInfo) (i : Elab.Info)
+def locationLinksOfInfo (kind : GoToKind) (ictx : InfoWithCtx)
     (infoTree? : Option InfoTree := none) : RequestM (Array LocationLink) := do
   let rc ← read
   let doc ← readDoc
@@ -116,6 +116,9 @@ def locationLinksOfInfo (kind : GoToKind) (ci : Elab.ContextInfo) (i : Elab.Info
       return #[ll]
     return #[]
 
+  let i := ictx.info
+  let ci := ictx.ctx
+  let children := ictx.children
   match i with
   | .ofTermInfo ti =>
     let mut expr := ti.expr
@@ -126,6 +129,43 @@ def locationLinksOfInfo (kind : GoToKind) (ci : Elab.ContextInfo) (i : Elab.Info
     | Expr.const n .. => return ← ci.runMetaM i.lctx <| locationLinksFromDecl i n
     | Expr.fvar id .. => return ← ci.runMetaM i.lctx <| locationLinksFromBinder i id
     | _ => pure ()
+
+    -- Check whether this `TermInfo` node is directly responsible for its `.expr`.
+    -- This is the case iff all of its children represent strictly smaller subexpressions;
+    -- it is sufficient to check this of all direct children of this node (and that its elaborator didn't expand it as a macro)
+    let isExprGenerator := children.all fun
+      | .node (Info.ofTermInfo info) _ => info.expr != expr
+      | .node (Info.ofMacroExpansionInfo _) _ => false
+      | _ => true
+
+    -- don't go-to-instance if this `TermInfo` didn't directly generate its `.expr`
+    if kind != declaration && isExprGenerator then
+      -- go-to-definition on a projection application of a typeclass
+      -- should return all instances generated by TC
+      expr ← ci.runMetaM i.lctx do instantiateMVars expr
+      if let .const n _ := expr.getAppFn then
+        -- also include constant along with instance results
+        let mut results ← ci.runMetaM i.lctx <| locationLinksFromDecl i n
+        if let some info := ci.env.getProjectionFnInfo? n then
+          let mut elaborators := default
+          if let some ei := i.toElabInfo? then do
+            -- also include elaborator along with instance results, as this wouldn't be accessible otherwise
+            if ei.elaborator != `Delab -- prevent an error if this `TermInfo` came from the infoview
+              && ei.elaborator != `Lean.Elab.Term.elabApp && ei.elaborator != `Lean.Elab.Term.elabIdent -- don't include trivial elaborators
+              then do
+              elaborators ← ci.runMetaM i.lctx <| locationLinksFromDecl i ei.elaborator
+          let instIdx := info.numParams
+          let appArgs := expr.getAppArgs
+          let rec extractInstances : Expr → RequestM (Array Name)
+            | .const declName _ => do
+              if ← ci.runMetaM i.lctx do Lean.Meta.isInstance declName then pure #[declName] else pure #[]
+            | .app fn arg => do pure $ (← extractInstances fn).append (← extractInstances arg)
+            | _ => pure #[]
+          if let some instArg := appArgs.get? instIdx then
+            for inst in (← extractInstances instArg) do
+              results := results.append (← ci.runMetaM i.lctx <| locationLinksFromDecl i inst)
+            results := results.append elaborators -- put elaborators at the end of the results
+        return results
   | .ofFieldInfo fi =>
     if kind == type then
       let expr ← ci.runMetaM i.lctx do
@@ -157,8 +197,8 @@ def handleDefinition (kind : GoToKind) (p : TextDocumentPositionParams)
 
   withWaitFindSnap doc (fun s => s.endPos > hoverPos)
     (notFoundX := pure #[]) fun snap => do
-      if let some (ci, i) := snap.infoTree.hoverableInfoAt? (includeStop := true /- #767 -/) hoverPos then
-        locationLinksOfInfo kind ci i snap.infoTree
+      if let some infoWithCtx := snap.infoTree.hoverableInfoAt? (omitIdentApps := true) (includeStop := true /- #767 -/) hoverPos then
+        locationLinksOfInfo kind infoWithCtx snap.infoTree
       else return #[]
 
 open RequestM in
@@ -214,7 +254,7 @@ def getInteractiveTermGoal (p : Lsp.PlainTermGoalParams)
   let hoverPos := text.lspPosToUtf8Pos p.position
   withWaitFindSnap doc (fun s => s.endPos > hoverPos)
     (notFoundX := pure none) fun snap => do
-      if let some (ci, i@(Elab.Info.ofTermInfo ti)) := snap.infoTree.termGoalAt? hoverPos then
+      if let some {ctx := ci, info := i@(Elab.Info.ofTermInfo ti), ..} := snap.infoTree.termGoalAt? hoverPos then
         let ty ← ci.runMetaM i.lctx do
           instantiateMVars <| ti.expectedType?.getD (← Meta.inferType ti.expr)
         -- for binders, hide the last hypothesis (the binder itself)

src/Lean/Server/FileWorker/WidgetRequests.lean

@@ -15,7 +15,7 @@ import Lean.Server.FileWorker.RequestHandling
 /-! Registers all widget-related RPC procedures. -/
 
 namespace Lean.Widget
-open Server
+open Server Lean.Elab
 
 structure MsgToInteractive where
   msg : WithRpcRef MessageData
@@ -127,7 +127,7 @@ builtin_initialize
     (Array Lsp.LocationLink)
     fun ⟨kind, ⟨i⟩⟩ => RequestM.asTask do
       let rc ← read
-      let ls ← FileWorker.locationLinksOfInfo kind i.ctx i.info
+      let ls ← FileWorker.locationLinksOfInfo kind i
       if !ls.isEmpty then return ls
       -- TODO(WN): unify handling of delab'd (infoview) and elab'd (editor) applications
       let .ofTermInfo ti := i.info | return #[]

src/Lean/Server/InfoUtils.lean

@@ -8,6 +8,31 @@ import Lean.PrettyPrinter
 
 namespace Lean.Elab
 
+/-- Elaborator information with elaborator context.
+
+It can be thought of as a "thunked" elaboration computation that allows us
+to retroactively extract type information, symbol locations, etc.
+through arbitrary invocations of `runMetaM` (where the necessary context and state
+can be reconstructed from `ctx` and `info.lctx`).
+
+W.r.t. widgets, this is used to tag different parts of expressions in `ppExprTagged`.
+This is the input to the RPC call `Lean.Widget.InteractiveDiagnostics.infoToInteractive`.
+It carries over information about delaborated
+`Info` nodes in a `CodeWithInfos`, and the associated pretty-printing
+functionality is purpose-specific to showing the contents of infoview popups.
+
+For use in standard LSP go-to-definition (see `Lean.Server.FileWorker.locationLinksOfInfo`),
+all the elaborator information we need for similar tasks is already fully recoverable via
+the `InfoTree` structure (see `Lean.Elab.InfoTree.visitM`).
+There we use this as a convienience wrapper for queried nodes (e.g. the return value of
+`Lean.Elab.InfoTree.hoverableInfoAt?`). It also includes the children info nodes
+as additional context (this is unused in the RPC case, as delaboration has no notion of child nodes).
+-/
+structure InfoWithCtx where
+  ctx  : Elab.ContextInfo
+  info : Elab.Info
+  children : PersistentArray InfoTree
+
 /-- Visit nodes, passing in a surrounding context (the innermost one) and accumulating results on the way back up. -/
 partial def InfoTree.visitM [Monad m]
     (preNode  : ContextInfo → Info → (children : PersistentArray InfoTree) → m Unit := fun _ _ _ => pure ())
@@ -135,11 +160,16 @@ def InfoTree.smallestInfo? (p : Info → Bool) (t : InfoTree) : Option (ContextI
   infos.toArray.getMax? (fun a b => a.1 > b.1) |>.map fun (_, ci, i) => (ci, i)
 
 /-- Find an info node, if any, which should be shown on hover/cursor at position `hoverPos`. -/
-partial def InfoTree.hoverableInfoAt? (t : InfoTree) (hoverPos : String.Pos) (includeStop := false) (omitAppFns := false) : Option (ContextInfo × Info) := Id.run do
-  let results := t.visitM (m := Id) (postNode := fun ctx info _ results => do
+partial def InfoTree.hoverableInfoAt? (t : InfoTree) (hoverPos : String.Pos) (includeStop := false) (omitAppFns := false) (omitIdentApps := false) : Option InfoWithCtx := Id.run do
+  let results := t.visitM (m := Id) (postNode := fun ctx info children results => do
     let mut results := results.bind (·.getD [])
     if omitAppFns && info.stx.isOfKind ``Parser.Term.app && info.stx[0].isIdent then
-      results := results.filter (·.2.2.stx != info.stx[0])
+        results := results.filter (·.2.info.stx != info.stx[0])
+    if omitIdentApps && info.stx.isIdent then
+      -- if an identifier stands for an application (e.g. in the case of a typeclass projection), prefer the application
+      if let .ofTermInfo ti := info then
+        if ti.expr.isApp then
+          results := results.filter (·.2.info.stx != info.stx)
     unless results.isEmpty do
       return results  -- prefer innermost results
     /-
@@ -161,16 +191,16 @@ partial def InfoTree.hoverableInfoAt? (t : InfoTree) (hoverPos : String.Pos) (in
       Int.negOfNat (r.stop - r.start).byteIdx,
       -- prefer results for constants over variables (which overlap at declaration names)
       if info matches .ofTermInfo { expr := .fvar .., .. } then 0 else 1)
-    [(priority, ctx, info)]) |>.getD []
+    [(priority, {ctx, info, children})]) |>.getD []
   -- sort results by lexicographical priority
   let maxPrio? :=
     let _ := @lexOrd
     let _ := @leOfOrd.{0}
     let _ := @maxOfLe
     results.map (·.1) |>.maximum?
   let res? := results.find? (·.1 == maxPrio?) |>.map (·.2)
-  if let some (_, i) := res? then
-    if let .ofTermInfo ti := i then
+  if let some i := res? then
+    if let .ofTermInfo ti := i.info then
       if ti.expr.isSyntheticSorry then
         return none
   return res?
@@ -328,7 +358,7 @@ where
     stx.getNumArgs == 2 && stx[0].isToken "by" && stx[1].getNumArgs == 1 && stx[1][0].isMissing
 
 
-partial def InfoTree.termGoalAt? (t : InfoTree) (hoverPos : String.Pos) : Option (ContextInfo × Info) :=
+partial def InfoTree.termGoalAt? (t : InfoTree) (hoverPos : String.Pos) : Option InfoWithCtx :=
   -- In the case `f a b`, where `f` is an identifier, the term goal at `f` should be the goal for the full application `f a b`.
   hoverableInfoAt? t hoverPos (includeStop := true) (omitAppFns := true)