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
@@ -231,8 +230,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
@@ -258,15 +258,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.
@@ -345,22 +350,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}"
@@ -374,6 +379,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?
@@ -443,7 +451,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

@@ -461,9 +461,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,7 +66,7 @@ def handleHover (p : HoverParams)
         | none => pure none
 
       -- now try info tree
-      if let some (ci, i) := snap.infoTree.hoverableInfoAt? hoverPos then
+      if let some (ci, i, _) := snap.infoTree.hoverableInfoAt? hoverPos then
         if let some range := i.range? then
           -- prefer info tree if at least as specific as parser docstring
           if stxDoc?.all fun (_, stxRange) => stxRange.includes range 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) (ci : Elab.ContextInfo) (i : Elab.Info) (children : PersistentArray InfoTree) 
     (infoTree? : Option InfoTree := none) : RequestM (Array LocationLink) := do
   let rc ← read
   let doc ← readDoc
@@ -126,6 +126,44 @@ 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 ()
+
+    let isExprGenerator := children.all fun tree => match tree with
+        | .node (Info.ofTermInfo info) _ =>
+          info.expr != expr
+        | .node (Info.ofMacroExpansionInfo _) _ =>
+          false
+        | _ => true
+
+
+    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
+      match expr.getAppFn with
+      | .const n _ =>
+        -- 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 result ←
+            if let some ei := i.toElabInfo? then do
+              if ei.elaborator != `Delab && ei.elaborator != `Lean.Elab.Term.elabApp && ei.elaborator != `Lean.Elab.Term.elabIdent then do
+                -- also include elaborator along with instance results
+                ci.runMetaM i.lctx <| locationLinksFromDecl i ei.elaborator
+              else pure default
+            else pure default
+          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 appArgs.size > instIdx then
+            let instArg := appArgs.get! instIdx
+            results := results.append <| ← (← extractInstances instArg).foldlM (init := result) fun acc n => do
+              pure $ acc.append (← ci.runMetaM i.lctx <| locationLinksFromDecl i n)
+        return results
+      | _ => pure ()
   | .ofFieldInfo fi =>
     if kind == type then
       let expr ← ci.runMetaM i.lctx do
@@ -157,8 +195,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 (ci, i, children) := snap.infoTree.hoverableInfoAt? (omitIdentApps := true) (includeStop := true /- #767 -/) hoverPos then
+        locationLinksOfInfo kind ci i children snap.infoTree
       else return #[]
 
 open RequestM in
@@ -214,7 +252,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 (ci, 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

@@ -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.ctx i.info i.children
       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

@@ -135,11 +135,18 @@ 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 (ContextInfo × Info × (PersistentArray InfoTree)) := Id.run do
+  let results := t.visitM (m := Id) (postNode := fun ctx info c 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])
+    if omitAppFns then
+      if info.stx.isOfKind ``Parser.Term.app && info.stx[0].isIdent then
+        results := results.filter (·.2.2.1.stx != info.stx[0])
+    if omitIdentApps then
+      -- if an identifier stands for an application (e.g. in the case of a typeclass projection), prefer the application
+      if info.stx.isIdent then
+        if let .ofTermInfo ti := info then
+          if ti.expr matches .app _ _ then
+            results := results.filter (·.2.2.1.stx != info.stx)
     unless results.isEmpty do
       return results  -- prefer innermost results
     /-
@@ -161,15 +168,15 @@ 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, c)]) |>.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 some (_, i, _) := res? then
     if let .ofTermInfo ti := i then
       if ti.expr.isSyntheticSorry then
         return none
@@ -323,7 +330,7 @@ where
       cs.any (hasNestedTactic pos tailPos)
     | _ => false
 
-partial def InfoTree.termGoalAt? (t : InfoTree) (hoverPos : String.Pos) : Option (ContextInfo × Info) :=
+partial def InfoTree.termGoalAt? (t : InfoTree) (hoverPos : String.Pos) : Option (ContextInfo × Info × (PersistentArray InfoTree)) :=
   -- 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)
 

src/Lean/Widget/Basic.lean

@@ -18,6 +18,7 @@ functionality is purpose-specific to showing the contents of infoview popups.
 structure InfoWithCtx where
   ctx  : Elab.ContextInfo
   info : Elab.Info
+  children : PersistentArray InfoTree
   deriving Inhabited, TypeName
 
 deriving instance TypeName for MessageData

src/Lean/Widget/InteractiveCode.lean

@@ -66,7 +66,7 @@ where
       | none   => go subTt
       | some i =>
         let t : SubexprInfo := {
-          info := WithRpcRef.mk { ctx, info := i }
+          info := WithRpcRef.mk { ctx, info := i, children := .empty}
           subexprPos := n
         }
         TaggedText.tag t (go subTt)