Implement goto declaration (microsoft#1343)

* Goto declaration

* Some more tests

Author: Mikhail Arkhipov

src/Analysis/Ast/Impl/Extensions/MemberExtensions.cs

@@ -86,5 +86,21 @@ public static string GetName(this IMember m) {
             }
             return null;
         }
+
+        public static ILocatedMember GetRootDefinition(this ILocatedMember lm) {
+            if (!(lm is IImportedMember im) || im.Parent == null) {
+                return lm;
+            }
+
+            var parent = im.Parent;
+            for (; parent != null;) {
+                if (!(parent is IImportedMember im1) || im1.Parent == null) {
+                    break;
+                }
+                parent = im1.Parent;
+            }
+            return parent;
+        }
+
     }
 }

src/LanguageServer/Impl/Implementation/Server.Editor.cs

@@ -92,6 +92,15 @@ public async Task<Reference[]> GotoDefinition(TextDocumentPositionParams @params
             return reference != null ? new[] { reference } : Array.Empty<Reference>();
         }
 
+        public async Task<Location> GotoDeclaration(TextDocumentPositionParams @params, CancellationToken cancellationToken) {
+            var uri = @params.textDocument.uri;
+            _log?.Log(TraceEventType.Verbose, $"Goto Declaration in {uri} at {@params.position}");
+
+            var analysis = await Document.GetAnalysisAsync(uri, Services, CompletionAnalysisTimeout, cancellationToken);
+            var reference = new DeclarationSource(Services).FindDefinition(analysis, @params.position, out _);
+            return reference != null ? new Location { uri = reference.uri, range = reference.range} : null;
+        }
+
         public Task<Reference[]> FindReferences(ReferencesParams @params, CancellationToken cancellationToken) {
             var uri = @params.textDocument.uri;
             _log?.Log(TraceEventType.Verbose, $"References in {uri} at {@params.position}");

src/LanguageServer/Impl/Implementation/Server.cs

@@ -88,6 +88,7 @@ public Server(IServiceManager services) {
                 workspaceSymbolProvider = true,
                 documentSymbolProvider = true,
                 renameProvider = true,
+                declarationProvider = true,
                 documentOnTypeFormattingProvider = new DocumentOnTypeFormattingOptions {
                     firstTriggerCharacter = "\n",
                     moreTriggerCharacter = new[] { ";", ":" }

src/LanguageServer/Impl/LanguageServer.cs

@@ -233,6 +233,14 @@ public async Task<Reference[]> GotoDefinition(JToken token, CancellationToken ca
             }
         }
 
+        [JsonRpcMethod("textDocument/declaration")]
+        public async Task<Location> GotoDeclaration(JToken token, CancellationToken cancellationToken) {
+            using (_requestTimer.Time("textDocument/declaration")) {
+                await _prioritizer.DefaultPriorityAsync(cancellationToken);
+                return await _server.GotoDeclaration(ToObject<TextDocumentPositionParams>(token), GetToken(cancellationToken));
+            }
+        }
+
         [JsonRpcMethod("textDocument/references")]
         public async Task<Reference[]> FindReferences(JToken token, CancellationToken cancellationToken) {
             using (_requestTimer.Time("textDocument/references")) {

src/LanguageServer/Impl/Protocol/Classes.cs

@@ -448,6 +448,7 @@ public sealed class ServerCapabilities {
         public DocumentOnTypeFormattingOptions documentOnTypeFormattingProvider;
         public bool renameProvider;
         public DocumentLinkOptions documentLinkProvider;
+        public bool declarationProvider; // 3.14.0+
         public ExecuteCommandOptions executeCommandProvider;
         public object experimental;
     }

src/LanguageServer/Impl/Sources/DefinitionSource.cs

@@ -29,13 +29,41 @@
 using Microsoft.Python.Parsing.Ast;
 
 namespace Microsoft.Python.LanguageServer.Sources {
-    internal sealed class DefinitionSource {
+    /// <summary>
+    /// Implements location of symbol declaration, such as 'B' in 'from A import B'
+    /// statement in the file. For 'goto definition' behavior see <see cref="DefinitionSource"/>.
+    /// </summary>
+    internal sealed class DeclarationSource : DefinitionSourceBase {
+        public DeclarationSource(IServiceContainer services) : base(services) { }
+        protected override ILocatedMember GetDefiningMember(IMember m) => m as ILocatedMember;
+    }
+
+    /// <summary>
+    /// Implements location of symbol definition. For example, in 'from A import B'
+    /// locates actual code of 'B' in module A. For 'goto declaration' behavior
+    /// see <see cref="DeclarationSource"/>.
+    /// </summary>
+    internal sealed class DefinitionSource : DefinitionSourceBase {
+        public DefinitionSource(IServiceContainer services) : base(services) { }
+        protected override ILocatedMember GetDefiningMember(IMember m) => (m as ILocatedMember)?.GetRootDefinition();
+    }
+
+    internal abstract class DefinitionSourceBase {
         private readonly IServiceContainer _services;
 
-        public DefinitionSource(IServiceContainer services) {
+        protected DefinitionSourceBase(IServiceContainer services) {
             _services = services;
         }
 
+        protected abstract ILocatedMember GetDefiningMember(IMember m);
+
+        /// <summary>
+        /// Locates definition or declaration of a symbol at the provided location.
+        /// </summary>
+        /// <param name="analysis">Document analysis.</param>
+        /// <param name="location">Location in the document.</param>
+        /// <param name="definingMember">Member location or null of not found.</param>
+        /// <returns>Definition location (module URI and the text range).</returns>
         public Reference FindDefinition(IDocumentAnalysis analysis, SourceLocation location, out ILocatedMember definingMember) {
             definingMember = null;
             if (analysis?.Ast == null) {
@@ -171,11 +199,12 @@ private Reference TryFromVariable(string name, IDocumentAnalysis analysis, Sourc
             definingMember = v;
             if (statement is ImportStatement || statement is FromImportStatement) {
                 // If we are on the variable definition in this module,
-                // then goto definition should go to the parent, if any.
+                // then goto declaration should go to the parent, if any.
+                // Goto to definition navigates to the very root of the parent chain.
                 var indexSpan = v.Definition.Span.ToIndexSpan(analysis.Ast);
                 var index = location.ToIndex(analysis.Ast);
                 if (indexSpan.Start <= index && index < indexSpan.End) {
-                    var parent = (v as IImportedMember)?.Parent;
+                    var parent = GetDefiningMember((v as IImportedMember)?.Parent);
                     var definition = parent?.Definition ?? (v.Value as ILocatedMember)?.Definition;
                     if (definition != null && CanNavigateToModule(definition.DocumentUri)) {
                         return new Reference { range = definition.Span, uri = definition.DocumentUri };
@@ -227,7 +256,7 @@ private Reference FromMemberExpression(MemberExpression mex, IDocumentAnalysis a
         }
 
         private Reference FromMember(IMember m) {
-            var definition = (m as ILocatedMember)?.Definition;
+            var definition = GetDefiningMember(m)?.Definition;
             var moduleUri = definition?.DocumentUri;
             // Make sure module we are looking for is not a stub
             if (m is IPythonType t) {

src/LanguageServer/Impl/Sources/ReferenceSource.cs

@@ -54,7 +54,7 @@ public async Task<Reference[]> FindAllReferencesAsync(Uri uri, SourceLocation lo
                     return Array.Empty<Reference>();
                 }
 
-                var rootDefinition = GetRootDefinition(definingMember);
+                var rootDefinition = definingMember.GetRootDefinition();
                 var name = definingMember.GetName();
 
                 // If it is an implicitly declared variable, such as function or a class
@@ -83,7 +83,7 @@ private async Task<Reference[]> FindAllReferencesAsync(string name, IPythonModul
                     return Array.Empty<Reference>();
                 }
 
-                rootDefinition = GetRootDefinition(definingMember);
+                rootDefinition = definingMember.GetRootDefinition();
             }
 
             return rootDefinition.References
@@ -151,20 +151,5 @@ private async Task<bool> AnalyzeFiles(IModuleManagement moduleManagement, IEnume
 
             return analysisTasks.Count > 0;
         }
-
-        private ILocatedMember GetRootDefinition(ILocatedMember lm) {
-            if (!(lm is IImportedMember im) || im.Parent == null) {
-                return lm;
-            }
-
-            var parent = im.Parent;
-            for (; parent != null;) {
-                if (!(parent is IImportedMember im1) || im1.Parent == null) {
-                    break;
-                }
-                parent = im1.Parent;
-            }
-            return parent;
-        }
     }
 }

src/LanguageServer/Test/GoToDefinitionTests.cs

@@ -158,7 +158,7 @@ import logging as log
         }
 
         [TestMethod, Priority(0)]
-        public async Task GotoModuleSourceFromImport1() {
+        public async Task GotoModuleSourceFromImport() {
             const string code = @"from logging import A";
             var analysis = await GetAnalysisAsync(code, PythonVersions.LatestAvailable3X);
             var ds = new DefinitionSource(Services);
@@ -171,14 +171,34 @@ public async Task GotoModuleSourceFromImport1() {
         }
 
         [TestMethod, Priority(0)]
-        public async Task GotoModuleSourceFromImport2() {
+        public async Task GotoDefitionFromImport() {
             const string code = @"
 from MultiValues import t
 x = t
 ";
             var analysis = await GetAnalysisAsync(code);
             var ds = new DefinitionSource(Services);
 
+            var reference = ds.FindDefinition(analysis, new SourceLocation(3, 5), out _);
+            reference.Should().NotBeNull();
+            reference.range.Should().Be(2, 0, 2, 1);
+            reference.uri.AbsolutePath.Should().Contain("MultiValues.py");
+
+            reference = ds.FindDefinition(analysis, new SourceLocation(2, 25), out _);
+            reference.Should().NotBeNull();
+            reference.range.Should().Be(2, 0, 2, 1);
+            reference.uri.AbsolutePath.Should().Contain("MultiValues.py");
+        }
+
+        [TestMethod, Priority(0)]
+        public async Task GotoDeclarationFromImport() {
+            const string code = @"
+from MultiValues import t
+x = t
+";
+            var analysis = await GetAnalysisAsync(code);
+            var ds = new DeclarationSource(Services);
+
             var reference = ds.FindDefinition(analysis, new SourceLocation(3, 5), out _);
             reference.Should().NotBeNull();
             reference.range.Should().Be(1, 24, 1, 25);
@@ -202,16 +222,50 @@ public async Task GotoModuleSourceFromImportAs() {
             reference.uri.AbsolutePath.Should().NotContain("pyi");
         }
 
+        [TestMethod, Priority(0)]
+        public async Task GotoDefinitionFromImportAs() {
+            const string code = @"
+from logging import critical as crit
+x = crit
+";
+            var analysis = await GetAnalysisAsync(code, PythonVersions.LatestAvailable3X);
+            var ds = new DefinitionSource(Services);
+
+            var reference = ds.FindDefinition(analysis, new SourceLocation(3, 6), out _);
+            reference.Should().NotBeNull();
+            reference.range.start.line.Should().BeGreaterThan(500);
+            reference.uri.AbsolutePath.Should().Contain("logging");
+            reference.uri.AbsolutePath.Should().NotContain("pyi");
+        }
+
+        [TestMethod, Priority(0)]
+        public async Task GotoDeclarationFromImportAs() {
+            const string code = @"
+from logging import critical as crit
+x = crit
+";
+            var analysis = await GetAnalysisAsync(code, PythonVersions.LatestAvailable3X);
+            var ds = new DeclarationSource(Services);
+
+            var reference = ds.FindDefinition(analysis, new SourceLocation(3, 6), out _);
+            reference.Should().NotBeNull();
+            reference.range.Should().Be(1, 32, 1, 36);
+        }
+
         [TestMethod, Priority(0)]
         public async Task GotoBuiltinObject() {
             const string code = @"
 class A(object):
     pass
 ";
             var analysis = await GetAnalysisAsync(code);
-            var ds = new DefinitionSource(Services);
 
-            var reference = ds.FindDefinition(analysis, new SourceLocation(2, 12), out _);
+            var ds1 = new DefinitionSource(Services);
+            var reference = ds1.FindDefinition(analysis, new SourceLocation(2, 12), out _);
+            reference.Should().BeNull();
+
+            var ds2 = new DeclarationSource(Services);
+            reference = ds2.FindDefinition(analysis, new SourceLocation(2, 12), out _);
             reference.Should().BeNull();
         }