add TermInfo.isExprGenerator field to identify when TermInfo node is responsible for its expr.

Author: rish987

src/Lean/Elab/InfoTree/Main.lean

@@ -267,20 +267,20 @@ def resolveGlobalNameWithInfos [MonadResolveName m] [MonadEnv m] [MonadError m]
 
 /-- Use this to descend a node on the infotree that is being built.
 
-It saves the current list of trees `t₀` and resets it and then runs `x >>= mkInfo`, producing either an `i : Info` or a hole id.
+It saves the current list of trees `t₀` and resets it and then runs `(x >>= mkInfo) <*> getInfoTrees`, producing either an `i : Info` or a hole id.
 Running `x >>= mkInfo` will modify the trees state and produce a new list of trees `t₁`.
 In the `i : Info` case, `t₁` become the children of a node `node i t₁` that is appended to `t₀`.
  -/
-def withInfoContext' [MonadFinally m] (x : m α) (mkInfo : α → m (Sum Info MVarId)) : m α := do
+def withInfoContext' [MonadFinally m] (x : m α) (mkInfo : α → m ((PersistentArray InfoTree) → Sum Info MVarId)) : m α := do
   if (← getInfoState).enabled then
     let treesSaved ← getResetInfoTrees
     Prod.fst <$> MonadFinally.tryFinally' x fun a? => do
       match a? with
       | none   => modifyInfoTrees fun _ => treesSaved
       | some a =>
-        let info ← mkInfo a
+        let infoFn ← mkInfo a
         modifyInfoTrees fun trees =>
-          match info with
+          match infoFn trees with
           | Sum.inl info   => treesSaved.push <| InfoTree.node info trees
           | Sum.inr mvarId => treesSaved.push <| InfoTree.hole mvarId
   else

src/Lean/Elab/InfoTree/Types.lean

@@ -36,6 +36,7 @@ structure TermInfo extends ElabInfo where
   lctx : LocalContext -- The local context when the term was elaborated.
   expectedType? : Option Expr
   expr : Expr
+  isExprGenerator : Bool := false -- Is this node's `stx` directly responsible for generating `expr`?
   isBinder : Bool := false
   deriving Inhabited
 

src/Lean/Elab/Term.lean

@@ -915,12 +915,23 @@ def isTacticOrPostponedHole? (e : Expr) : TermElabM (Option MVarId) := do
     | _                                  => return none
   | _ => pure none
 
-def mkTermInfo (elaborator : Name) (stx : Syntax) (e : Expr) (expectedType? : Option Expr := none) (lctx? : Option LocalContext := none) (isBinder := false) : TermElabM (Sum Info MVarId) := do
+def mkTermInfo (elaborator : Name) (stx : Syntax) (e : Expr) (expectedType? : Option Expr := none) (lctx? : Option LocalContext := none) (isBinder := false) : TermElabM ((PersistentArray InfoTree) → Sum Info MVarId) := do
   match (← isTacticOrPostponedHole? e) with
-  | some mvarId => return Sum.inr mvarId
+  | some mvarId => return fun _ => Sum.inr mvarId
   | none =>
     let e := removeSaveInfoAnnotation e
-    return Sum.inl <| Info.ofTermInfo { elaborator, lctx := lctx?.getD (← getLCtx), expr := e, stx, expectedType?, isBinder }
+    let lctx ← getLCtx
+    return fun trees =>
+      -- NOTE: this can probably be made more efficient, but for now
+      -- a sufficient check for whether this node is responsible for `e` is to make sure
+      -- that no term subtree has the same expression
+      let isExprGenerator :=
+        trees.all fun tree => match tree with
+        | .node (Info.ofTermInfo info) _ =>
+          info.expr != e
+        | _ => true
+
+      Sum.inl <| Info.ofTermInfo { elaborator, lctx := lctx?.getD lctx, expr := e, isExprGenerator, stx, expectedType?, isBinder }
 
 /--
 Pushes a new leaf node to the info tree associating the expression `e` to the syntax `stx`.
@@ -949,7 +960,7 @@ def addTermInfo (stx : Syntax) (e : Expr) (expectedType? : Option Expr := none)
 def addTermInfo' (stx : Syntax) (e : Expr) (expectedType? : Option Expr := none) (lctx? : Option LocalContext := none) (elaborator := Name.anonymous) (isBinder := false) : TermElabM Unit :=
   discard <| addTermInfo stx e expectedType? lctx? elaborator isBinder
 
-def withInfoContext' (stx : Syntax) (x : TermElabM Expr) (mkInfo : Expr → TermElabM (Sum Info MVarId)) : TermElabM Expr := do
+def withInfoContext' (stx : Syntax) (x : TermElabM Expr) (mkInfo : Expr → TermElabM ((PersistentArray InfoTree) → Sum Info MVarId)) : TermElabM Expr := do
   if (← read).inPattern then
     let e ← x
     return mkPatternWithRef e stx

src/Lean/Server/FileWorker/RequestHandling.lean

@@ -123,7 +123,7 @@ def locationLinksOfInfo (kind : GoToKind) (ci : Elab.ContextInfo) (i : Elab.Info
     | Expr.fvar id .. => return ← ci.runMetaM i.lctx <| locationLinksFromBinder i id
     | _ => pure ()
 
-    if kind != declaration then
+    if kind != declaration && ti.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