From 4545e183d82fdeec03c42ee00f38bf19f04ffddf Mon Sep 17 00:00:00 2001
From: Joscha <joscha@plugh.de>
Date: Fri, 28 Jan 2022 21:06:01 +0100
Subject: [PATCH] fix: go to definition in modified file

---
 src/Lean/Server/References.lean | 12 ++++++++++++
 src/Lean/Server/Watchdog.lean   | 13 +++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/src/Lean/Server/References.lean b/src/Lean/Server/References.lean
index f8b78b8cb3..3597f5f068 100644
--- a/src/Lean/Server/References.lean
+++ b/src/Lean/Server/References.lean
@@ -227,6 +227,18 @@ def referringTo (self : References) (ident : RefIdent) (srcSearchPath : SearchPa
           result := result.push ⟨uri, range⟩
   result
 
+def definitionOf (self : References) (ident : RefIdent) (srcSearchPath : SearchPath)
+    : IO (Option Location) := do
+  for (module, refs) in self.allRefs.toList do
+    if let some info := refs.find? ident then
+      if let some definition := info.definition then
+        if let some path ← srcSearchPath.findModuleWithExt "lean" module then
+          -- Resolve symlinks (such as `src` in the build dir) so that files are
+          -- opened in the right folder
+          let uri := DocumentUri.ofPath <| ← IO.FS.realPath path
+          return some ⟨uri, definition⟩
+  return none
+
 def definitionsMatching (self : References) (srcSearchPath : SearchPath) (filter : Name → Option α)
     (maxAmount? : Option Nat := none) : IO $ Array (α × Location) := do
   let mut result := #[]
diff --git a/src/Lean/Server/Watchdog.lean b/src/Lean/Server/Watchdog.lean
index 6babbcd6e9..bed32092d9 100644
--- a/src/Lean/Server/Watchdog.lean
+++ b/src/Lean/Server/Watchdog.lean
@@ -368,6 +368,15 @@ end ServerM
 
 section RequestHandling
 
+def findDefinition (p : TextDocumentPositionParams) : ServerM <| Option Location := do
+  if let some path := p.textDocument.uri.toPath? then
+    let srcSearchPath := (← read).srcSearchPath
+    if let some module ← searchModuleNameOfFileName path srcSearchPath then
+      let references ← (← read).references.get
+      if let some ident := references.findAt? module p.position then
+        return ← references.definitionOf ident srcSearchPath
+  return none
+
 def handleReference (p : ReferenceParams) : ServerM (Array Location) := do
   if let some path := p.textDocument.uri.toPath? then
     let srcSearchPath := (← read).srcSearchPath
@@ -527,6 +536,10 @@ section MessageHandling
           code := ErrorCode.internalError
           message := s!"Failed to process request {id}: {e}"
         }
+    if method == "textDocument/definition" || method == "textDocument/declaration" then
+      let params ← parseParams TextDocumentPositionParams params
+      if let some definition ← findDefinition params then
+        (← read).hOut.writeLspResponse ⟨id, #[definition]⟩
     match method with
       | "textDocument/references" => handle ReferenceParams (Array Location) handleReference
       | "workspace/symbol" => handle WorkspaceSymbolParams (Array SymbolInformation) handleWorkspaceSymbol