From b76dd1a8e300091441733b05270fa0749816c596 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 15 Jul 2021 18:48:14 +0200
Subject: [PATCH] feat: go-to-definition for local variables

---
 src/Lean/Elab/InfoTree.lean                   |  4 ++--
 .../Server/FileWorker/RequestHandling.lean    | 21 +++++++++++++++++--
 tests/lean/interactive/goTo.lean              |  7 +++++++
 tests/lean/interactive/goTo.lean.expected.out |  9 ++++++++
 4 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/src/Lean/Elab/InfoTree.lean b/src/Lean/Elab/InfoTree.lean
index 151a5b3b9c..1403ec1a9e 100644
--- a/src/Lean/Elab/InfoTree.lean
+++ b/src/Lean/Elab/InfoTree.lean
@@ -102,12 +102,12 @@ inductive InfoTree where
   | hole (mvarId : MVarId) -- The elaborator creates holes (aka metavariables) for tactics and postponed terms
   deriving Inhabited
 
-partial def InfoTree.findInfo? (p : Info → Bool) (t : InfoTree) : Option InfoTree :=
+partial def InfoTree.findInfo? (p : Info → Bool) (t : InfoTree) : Option Info :=
   match t with
   | context _ t => findInfo? p t
   | node i ts   =>
     if p i then
-      some t
+      some i
     else
       ts.findSome? (findInfo? p)
   | _ => none
diff --git a/src/Lean/Server/FileWorker/RequestHandling.lean b/src/Lean/Server/FileWorker/RequestHandling.lean
index 3f88127959..1135cceca4 100644
--- a/src/Lean/Server/FileWorker/RequestHandling.lean
+++ b/src/Lean/Server/FileWorker/RequestHandling.lean
@@ -90,6 +90,21 @@ partial def handleDefinition (kind : GoToKind) (p : TextDocumentPositionParams)
       return #[ll]
     return #[]
 
+  let locationLinksFromBinder (t : InfoTree) (i : Elab.Info) (id : FVarId) := do
+    if let some i' := t.findInfo? fun
+        | Info.ofTermInfo { isBinder := true, expr := Expr.fvar id' .., .. } => id' == id
+        | _ => false then
+      if let some r := i'.range? then
+        let r := r.toLspRange text
+        let ll : LocationLink := {
+          originSelectionRange? := (·.toLspRange text) <$> i.range?
+          targetUri := p.textDocument.uri
+          targetRange := r
+          targetSelectionRange := r
+        }
+        return #[ll]
+    return #[]
+
   withWaitFindSnap doc (fun s => s.endPos > hoverPos)
     (notFoundX := pure #[]) fun snap => do
       for t in snap.cmdState.infoState.trees do
@@ -99,8 +114,10 @@ partial def handleDefinition (kind : GoToKind) (p : TextDocumentPositionParams)
             if kind = type then
               expr ← ci.runMetaM i.lctx do
                 Expr.getAppFn (← Meta.instantiateMVars (← Meta.inferType expr))
-            if let some n := expr.constName? then
-              return ← ci.runMetaM i.lctx <| locationLinksFromDecl i n
+            match expr with
+            | Expr.const n .. => return ← ci.runMetaM i.lctx <| locationLinksFromDecl i n
+            | Expr.fvar id .. => return ← ci.runMetaM i.lctx <| locationLinksFromBinder t i id
+            | _ => pure ()
           if let Info.ofFieldInfo fi := i then
             if kind = type then
               let expr ← ci.runMetaM i.lctx do
diff --git a/tests/lean/interactive/goTo.lean b/tests/lean/interactive/goTo.lean
index 859c5eaddb..c90a7b9c9d 100644
--- a/tests/lean/interactive/goTo.lean
+++ b/tests/lean/interactive/goTo.lean
@@ -44,3 +44,10 @@ def Baz (α : Type) := α
 
 #check fun (b : Baz Nat) => b
                           --^ textDocument/typeDefinition
+
+example : Nat :=
+  let a := 1
+--v textDocument/definition
+  a + b
+where
+  b := 2
diff --git a/tests/lean/interactive/goTo.lean.expected.out b/tests/lean/interactive/goTo.lean.expected.out
index 194dbd1b6b..0e559e5c69 100644
--- a/tests/lean/interactive/goTo.lean.expected.out
+++ b/tests/lean/interactive/goTo.lean.expected.out
@@ -87,3 +87,12 @@
   "originSelectionRange":
   {"start": {"line": 44, "character": 28},
    "end": {"line": 44, "character": 29}}}]
+{"textDocument": {"uri": "file://goTo.lean"},
+ "position": {"line": 50, "character": 2}}
+[{"targetUri": "file://goTo.lean",
+  "targetSelectionRange":
+  {"start": {"line": 48, "character": 6}, "end": {"line": 48, "character": 7}},
+  "targetRange":
+  {"start": {"line": 48, "character": 6}, "end": {"line": 48, "character": 7}},
+  "originSelectionRange":
+  {"start": {"line": 50, "character": 2}, "end": {"line": 50, "character": 3}}}]