Add support for implementation edges

Fixes dotnet#59617

src/Features/Lsif/Generator/Generator.cs

@@ -14,6 +14,7 @@
 using Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.ResultSetTracking;
 using Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.Writing;
 using Microsoft.CodeAnalysis.LanguageService;
+using Microsoft.CodeAnalysis.Shared.Extensions;
 using Roslyn.Utilities;
 using Methods = Microsoft.VisualStudio.LanguageServer.Protocol.Methods;
 
@@ -228,6 +229,45 @@ SymbolKind.RangeVariable or
                     {
                         var definitionResultsId = symbolResultsTracker.GetResultIdForSymbol(declaredSymbol, Methods.TextDocumentDefinitionName, () => new DefinitionResult(idFactory));
                         lsifJsonWriter.Write(new Item(definitionResultsId.As<DefinitionResult, Vertex>(), lazyRangeVertex.Value.GetId(), documentVertex.GetId(), idFactory));
+
+                        // If this declared symbol also implements an interface member, we count this as a definition of the interface member as well.
+                        // Note in C# there are estoeric cases where a method can implement an interface member even though the containing type does not
+                        // implement the interface, for example in this case:
+                        //
+                        //     interface I { void M(); }
+                        //     class Base { public void M() { } }
+                        //     class Derived : Base, I { }
+                        //
+                        // We don't worry about supporting these cases here.
+                        var implementedMembers = declaredSymbol.ExplicitOrImplicitInterfaceImplementations();
+
+                        foreach (var implementedMember in implementedMembers)
+                            MarkImplementationOfSymbol(implementedMember);
+
+                        // If this overrides a method, we'll also mark it the same way. We want to chase to the base virtual method, skipping over intermediate
+                        // methods so that way all overrides of the same method point to the same virtual method
+                        if (declaredSymbol.IsOverride)
+                        {
+                            var overridenMember = declaredSymbol.GetOverriddenMember();
+
+                            while (overridenMember?.GetOverriddenMember() != null)
+                                overridenMember = overridenMember.GetOverriddenMember();
+
+                            if (overridenMember != null)
+                                MarkImplementationOfSymbol(overridenMember);
+                        }
+
+                        void MarkImplementationOfSymbol(ISymbol baseMember)
+                        {
+                            // First we create a definition link for the reference results for the base member
+                            var referenceResultsId = symbolResultsTracker.GetResultSetReferenceResultId(baseMember.OriginalDefinition, idFactory);
+                            lsifJsonWriter.Write(new Item(referenceResultsId.As<ReferenceResult, Vertex>(), lazyRangeVertex.Value.GetId(), documentVertex.GetId(), idFactory, property: "definitions"));
+
+                            // Then also link the result set for the method to the moniker that it implements
+                            referenceResultsId = symbolResultsTracker.GetResultSetReferenceResultId(declaredSymbol.OriginalDefinition, idFactory);
+                            var implementedMemberMoniker = symbolResultsTracker.GetResultIdForSymbol<Moniker>(baseMember.OriginalDefinition, "moniker", () => throw new Exception("When we produced the resultSet, we should have already created a moniker for it."));
+                            lsifJsonWriter.Write(new Item(referenceResultsId.As<ReferenceResult, Vertex>(), implementedMemberMoniker, documentVertex.GetId(), idFactory, property: "referenceLinks"));
+                        }
                     }
 
                     if (referencedSymbol != null)
@@ -236,7 +276,7 @@ SymbolKind.RangeVariable or
                         // symbol but the range can point a different symbol's resultSet. This can happen if the token is
                         // both a definition of a symbol (where we will point to the definition) but also a reference to some
                         // other symbol.
-                        var referenceResultsId = symbolResultsTracker.GetResultIdForSymbol(referencedSymbol.OriginalDefinition, Methods.TextDocumentReferencesName, () => new ReferenceResult(idFactory));
+                        var referenceResultsId = symbolResultsTracker.GetResultSetReferenceResultId(referencedSymbol.OriginalDefinition, idFactory);
                         lsifJsonWriter.Write(new Item(referenceResultsId.As<ReferenceResult, Vertex>(), lazyRangeVertex.Value.GetId(), documentVertex.GetId(), idFactory, property: "references"));
                     }
 

src/Features/Lsif/Generator/Graph/Item.cs

@@ -5,7 +5,7 @@
 namespace Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.Graph
 {
     /// <summary>
-    /// Represents a single item that points to a range from a result. See https://github.com/Microsoft/language-server-protocol/blob/master/indexFormat/specification.md#request-textdocumentreferences
+    /// Represents a single item that points to a range or moniker from a result. See https://github.com/Microsoft/language-server-protocol/blob/master/indexFormat/specification.md#request-textdocumentreferences
     /// for an example of item edges.
     /// </summary>
     internal sealed class Item : Edge
@@ -19,5 +19,12 @@ public Item(Id<Vertex> outVertex, Id<Range> range, Id<LsifDocument> document, Id
             Document = document;
             Property = property;
         }
+
+        public Item(Id<Vertex> outVertex, Id<Moniker> moniker, Id<LsifDocument> document, IdFactory idFactory, string? property = null)
+            : base(label: "item", outVertex, new[] { moniker.As<Moniker, Vertex>() }, idFactory)
+        {
+            Document = document;
+            Property = property;
+        }
     }
 }

src/Features/Lsif/Generator/ResultSetTracking/IResultSetTrackerExtensions.cs

@@ -0,0 +1,23 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.Graph;
+using Microsoft.VisualStudio.LanguageServer.Protocol;
+
+namespace Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.ResultSetTracking
+{
+    internal static class IResultSetTrackerExtensions
+    {
+        /// <summary>
+        /// Returns the ID of the <see cref="ReferenceResult"/> for a <see cref="ResultSet"/>.
+        /// </summary>
+        public static Id<ReferenceResult> GetResultSetReferenceResultId(this IResultSetTracker tracker, ISymbol symbol, IdFactory idFactory)
+            => tracker.GetResultIdForSymbol(symbol, Methods.TextDocumentReferencesName, () => new ReferenceResult(idFactory));
+    }
+}

src/Features/Lsif/Generator/ResultSetTracking/SymbolHoldingResultSetTracker.cs

@@ -80,8 +80,8 @@ private TrackedResultSet GetTrackedResultSet(ISymbol symbol)
 
                 if (monikerVertex != null)
                 {
-                    _lsifJsonWriter.Write(monikerVertex);
-                    _lsifJsonWriter.Write(Edge.Create("moniker", trackedResultSet.Id, monikerVertex.GetId(), _idFactory));
+                    // Attach the moniker vertex for this result set
+                    _ = GetResultIdForSymbol(symbol, "moniker", () => monikerVertex);
                 }
             }
 
@@ -157,15 +157,15 @@ public TrackedResultSet(Id<ResultSet> id)
                 Id = id;
             }
 
-            public Id<T> GetResultId<T>(string edgeKind, Func<T> vertexCreator, ILsifJsonWriter lsifJsonWriter, IdFactory idFactory) where T : Vertex
+            public Id<T> GetResultId<T>(string edgeLabel, Func<T> vertexCreator, ILsifJsonWriter lsifJsonWriter, IdFactory idFactory) where T : Vertex
             {
                 lock (_edgeKindToVertexId)
                 {
-                    if (_edgeKindToVertexId.TryGetValue(edgeKind, out var existingId))
+                    if (_edgeKindToVertexId.TryGetValue(edgeLabel, out var existingId))
                     {
                         if (!existingId.HasValue)
                         {
-                            throw new Exception($"This ResultSet already has an edge of {edgeKind} as {nameof(ResultSetNeedsInformationalEdgeAdded)} was called with this edge kind.");
+                            throw new Exception($"This ResultSet already has an edge of {edgeLabel} as {nameof(ResultSetNeedsInformationalEdgeAdded)} was called with this edge label.");
                         }
 
                         // TODO: this is a violation of the type system here, really: we're assuming that all calls to this function with the same edge kind
@@ -174,10 +174,10 @@ public Id<T> GetResultId<T>(string edgeKind, Func<T> vertexCreator, ILsifJsonWri
                     }
 
                     var vertex = vertexCreator();
-                    _edgeKindToVertexId.Add(edgeKind, vertex.GetId().As<T, Vertex>());
+                    _edgeKindToVertexId.Add(edgeLabel, vertex.GetId().As<T, Vertex>());
 
                     lsifJsonWriter.Write(vertex);
-                    lsifJsonWriter.Write(Edge.Create(edgeKind, Id, vertex.GetId(), idFactory));
+                    lsifJsonWriter.Write(Edge.Create(edgeLabel, Id, vertex.GetId(), idFactory));
 
                     return vertex.GetId();
                 }

src/Features/Lsif/GeneratorTest/RangeResultSetTests.vb

@@ -5,6 +5,7 @@
 Imports Microsoft.CodeAnalysis.Editor.UnitTests.Workspaces
 Imports Microsoft.CodeAnalysis.Test.Utilities
 Imports Microsoft.VisualStudio.LanguageServer.Protocol
+Imports Roslyn.Test.Utilities
 
 Namespace Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.UnitTests
     <UseExportProvider>
@@ -179,5 +180,74 @@ Namespace Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.UnitTests
                 End If
             Next
         End Function
+
+        <Fact>
+        Public Async Function InterfaceImplementation() As Task
+            Dim lsif = Await TestLsifOutput.GenerateForWorkspaceAsync(
+                TestWorkspace.CreateWorkspace(
+                    <Workspace>
+                        <Project Language="C#" AssemblyName=<%= TestProjectAssemblyName %> FilePath="Z:\TestProject.csproj" CommonReferences="true">
+                            <Document Name="A.cs" FilePath="Z:\A.cs">
+interface I { void {|Base:M|}(); }
+class C : I { public void {|Implementation:M|}() { } }
+                            </Document>
+                        </Project>
+                    </Workspace>))
+
+            Await AssertImplementationCorrectlyLinked(lsif)
+        End Function
+
+        <Fact>
+        Public Async Function [Overrides]() As Task
+            Dim lsif = Await TestLsifOutput.GenerateForWorkspaceAsync(
+                TestWorkspace.CreateWorkspace(
+                    <Workspace>
+                        <Project Language="C#" AssemblyName=<%= TestProjectAssemblyName %> FilePath="Z:\TestProject.csproj" CommonReferences="true">
+                            <Document Name="A.cs" FilePath="Z:\A.cs">
+class C { public virtual void {|Base:M|}() { } }
+class D : C { public override void {|Implementation:M|}() { } }
+                            </Document>
+                        </Project>
+                    </Workspace>))
+
+            Await AssertImplementationCorrectlyLinked(lsif)
+        End Function
+
+        <Fact>
+        Public Async Function OverridesAlwaysPointsToInitialVirtualMethod() As Task
+            Dim lsif = Await TestLsifOutput.GenerateForWorkspaceAsync(
+                TestWorkspace.CreateWorkspace(
+                    <Workspace>
+                        <Project Language="C#" AssemblyName=<%= TestProjectAssemblyName %> FilePath="Z:\TestProject.csproj" CommonReferences="true">
+                            <Document Name="A.cs" FilePath="Z:\A.cs">
+class C { public virtual void {|Base:M|}() { } }
+class D : C { public override void {|Implementation:M|}() { } }
+class E : D { public override void {|Implementation:M|}() { } }
+                            </Document>
+                        </Project>
+                    </Workspace>))
+
+            Await AssertImplementationCorrectlyLinked(lsif)
+        End Function
+
+        Private Shared Async Function AssertImplementationCorrectlyLinked(lsif As TestLsifOutput) As Task
+            Dim interfaceMethodRange = Await lsif.GetAnnotatedRangeAsync("Base")
+            Dim implementationRanges = Await lsif.GetAnnotatedRangesAsync("Implementation")
+
+            ' The references vertex should have a definition items pointing ot the implementations
+            Dim resultSetVertex = lsif.GetLinkedVertices(Of Graph.ResultSet)(interfaceMethodRange, "next").Single()
+            Dim expectedMoniker = lsif.GetLinkedVertices(Of Graph.Moniker)(resultSetVertex, "moniker").Single()
+            Dim referencesVertex = lsif.GetLinkedVertices(Of Graph.ReferenceResult)(resultSetVertex, Methods.TextDocumentReferencesName).Single()
+            Dim definitionRangesForImplementingMethods = lsif.GetLinkedVertices(Of Graph.Range)(referencesVertex, Function(e) DirectCast(e, Graph.Item).Property = "definitions")
+            AssertEx.SetEqual(implementationRanges, definitionRangesForImplementingMethods)
+
+            ' The result set for the implementation method should point to the moniker of the base
+            For Each implementationRange In implementationRanges
+                resultSetVertex = lsif.GetLinkedVertices(Of Graph.ResultSet)(implementationRange, "next").Single()
+                referencesVertex = lsif.GetLinkedVertices(Of Graph.ReferenceResult)(resultSetVertex, Methods.TextDocumentReferencesName).Single()
+                Dim moniker = lsif.GetLinkedVertices(Of Graph.Moniker)(referencesVertex, Function(e) DirectCast(e, Graph.Item).Property = "referenceLinks").Single()
+                Assert.Same(expectedMoniker, moniker)
+            Next
+        End Function
     End Class
 End Namespace

src/Features/Lsif/GeneratorTest/Utilities/TestLsifJsonWriter.vb

@@ -66,12 +66,19 @@ Namespace Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.UnitTests.U
         ''' Returns all the vertices linked to the given vertex by the edge type.
         ''' </summary>
         Public Function GetLinkedVertices(Of T As Vertex)(vertex As Graph.Vertex, edgeLabel As String) As ImmutableArray(Of T)
+            Return GetLinkedVertices(Of T)(vertex, Function(e) e.Label = edgeLabel)
+        End Function
+
+        ''' <summary>
+        ''' Returns all the vertices linked to the given vertex by the edge predicate.
+        ''' </summary>
+        Public Function GetLinkedVertices(Of T As Vertex)(vertex As Graph.Vertex, predicate As Func(Of Edge, Boolean)) As ImmutableArray(Of T)
             SyncLock _gate
                 Dim builder = ImmutableArray.CreateBuilder(Of T)
 
                 Dim edges As List(Of Edge) = Nothing
                 If _edgesByOutVertex.TryGetValue(vertex, edges) Then
-                    Dim inVerticesId = edges.Where(Function(e) e.Label = edgeLabel).SelectMany(Function(e) e.InVertices)
+                    Dim inVerticesId = edges.Where(predicate).SelectMany(Function(e) e.InVertices)
 
                     For Each inVertexId In inVerticesId
                         ' This is an unsafe "cast" if you will converting the ID to the expected type;

src/Features/Lsif/GeneratorTest/Utilities/TestLsifOutput.vb

@@ -6,7 +6,9 @@ Imports System.Collections.Immutable
 Imports Microsoft.CodeAnalysis.Editor.UnitTests.Workspaces
 Imports Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.Graph
 Imports Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.Writing
+Imports Microsoft.CodeAnalysis.Text
 Imports LSP = Microsoft.VisualStudio.LanguageServer.Protocol
+Imports Roslyn.Utilities
 
 Namespace Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.UnitTests.Utilities
     Friend Class TestLsifOutput
@@ -43,6 +45,10 @@ Namespace Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.UnitTests.U
             Return _testLsifJsonWriter.GetElementById(id)
         End Function
 
+        Public Function GetLinkedVertices(Of T As Vertex)(vertex As Graph.Vertex, predicate As Func(Of Edge, Boolean)) As ImmutableArray(Of T)
+            Return _testLsifJsonWriter.GetLinkedVertices(Of T)(vertex, predicate)
+        End Function
+
         ''' <summary>
         ''' Returns all the vertices linked to the given vertex by the edge type.
         ''' </summary>
@@ -56,10 +62,7 @@ Namespace Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.UnitTests.U
             End Get
         End Property
 
-        ''' <summary>
-        ''' Returns the <see cref="Range" /> verticies in the output that corresponds to the selected range in the <see cref="TestWorkspace" />.
-        ''' </summary>
-        Public Async Function GetSelectedRangesAsync() As Task(Of IEnumerable(Of Graph.Range))
+        Private Async Function GetRangesAsync(selector As Func(Of TestHostDocument, IEnumerable(Of TextSpan))) As Task(Of IEnumerable(Of Graph.Range))
             Dim builder = ImmutableArray.CreateBuilder(Of Range)
 
             For Each testDocument In _workspace.Documents
@@ -69,7 +72,7 @@ Namespace Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.UnitTests.U
                                                         .Single()
                 Dim rangeVertices = GetLinkedVertices(Of Range)(documentVertex, "contains")
 
-                For Each selectedSpan In testDocument.SelectedSpans
+                For Each selectedSpan In selector(testDocument)
                     Dim document = _workspace.CurrentSolution.GetDocument(testDocument.Id)
                     Dim text = Await document.GetTextAsync()
                     Dim linePositionSpan = text.Lines.GetLinePositionSpan(selectedSpan)
@@ -85,10 +88,25 @@ Namespace Microsoft.CodeAnalysis.LanguageServerIndexFormat.Generator.UnitTests.U
             Return builder.ToImmutable()
         End Function
 
+        ''' <summary>
+        ''' Returns the <see cref="Range" /> verticies in the output that corresponds to the selected range in the <see cref="TestWorkspace" />.
+        ''' </summary>
+        Public Function GetSelectedRangesAsync() As Task(Of IEnumerable(Of Graph.Range))
+            Return GetRangesAsync(Function(testDocument) testDocument.SelectedSpans)
+        End Function
+
         Public Async Function GetSelectedRangeAsync() As Task(Of Graph.Range)
             Return (Await GetSelectedRangesAsync()).Single()
         End Function
 
+        Public Function GetAnnotatedRangesAsync(annotation As String) As Task(Of IEnumerable(Of Graph.Range))
+            Return GetRangesAsync(Function(testDocument) testDocument.AnnotatedSpans.GetValueOrDefault(annotation))
+        End Function
+
+        Public Async Function GetAnnotatedRangeAsync(annotation As String) As Task(Of Graph.Range)
+            Return (Await GetAnnotatedRangesAsync(annotation)).Single()
+        End Function
+
         Public Function GetFoldingRanges(document As Document) As LSP.FoldingRange()
             Dim documentVertex = _testLsifJsonWriter.Vertices _
                                                         .OfType(Of LsifDocument) _