Use ISpanMapper before sending cross file results to LSP

src/EditorFeatures/TestUtilities/LanguageServer/AbstractLanguageServerProtocolTests.cs

@@ -12,10 +12,12 @@
 using System.Linq;
 using System.Text.RegularExpressions;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.Editor.UnitTests;
 using Microsoft.CodeAnalysis.Editor.UnitTests.Workspaces;
+using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.LanguageServer;
 using Microsoft.CodeAnalysis.LanguageServer.CustomProtocol;
@@ -66,6 +68,37 @@ public ImmutableArray<Document> GetDocuments(Uri documentUri)
             }
         }
 
+        private class TestSpanMapperProvider : IDocumentServiceProvider
+        {
+            TService IDocumentServiceProvider.GetService<TService>()
+                => (TService)(object)new TestSpanMapper();
+        }
+
+        internal class TestSpanMapper : ISpanMappingService
+        {
+            private static readonly LinePositionSpan s_mappedLinePosition = new LinePositionSpan(new LinePosition(0, 0), new LinePosition(0, 5));
+            private static readonly string s_mappedFilePath = "c:\\MappedFile.cs";
+
+            internal static readonly string GeneratedFileName = "GeneratedFile.cs";
+
+            internal static readonly LSP.Location MappedFileLocation = new LSP.Location
+            {
+                Range = ProtocolConversions.LinePositionToRange(s_mappedLinePosition),
+                Uri = new Uri(s_mappedFilePath)
+            };
+
+            public Task<ImmutableArray<MappedSpanResult>> MapSpansAsync(Document document, IEnumerable<TextSpan> spans, CancellationToken cancellationToken)
+            {
+                ImmutableArray<MappedSpanResult> mappedResult = default;
+                if (document.Name == GeneratedFileName)
+                {
+                    mappedResult = ImmutableArray.Create(new MappedSpanResult(s_mappedFilePath, s_mappedLinePosition, new TextSpan(0, 5)));
+                }
+
+                return Task.FromResult(mappedResult);
+            }
+        }
+
         protected virtual ExportProvider GetExportProvider()
         {
             var requestHelperTypes = DesktopTestHelpers.GetAllTypesImplementingGivenInterface(
@@ -258,6 +291,18 @@ protected TestWorkspace CreateXmlTestWorkspace(string xmlContent, out Dictionary
             return workspace;
         }
 
+        protected static void AddMappedDocument(Workspace workspace, string markup)
+        {
+            var generatedDocumentId = DocumentId.CreateNewId(workspace.CurrentSolution.ProjectIds.First());
+            var version = VersionStamp.Create();
+            var loader = TextLoader.From(TextAndVersion.Create(SourceText.From(markup), version, TestSpanMapper.GeneratedFileName));
+            var generatedDocumentInfo = DocumentInfo.Create(generatedDocumentId, TestSpanMapper.GeneratedFileName, SpecializedCollections.EmptyReadOnlyList<string>(),
+                SourceCodeKind.Regular, loader, $"C:\\{TestSpanMapper.GeneratedFileName}", true, new TestSpanMapperProvider());
+            var newSolution = workspace.CurrentSolution.AddDocument(generatedDocumentInfo);
+            workspace.TryApplyChanges(newSolution);
+            UpdateSolutionProvider((TestWorkspace)workspace, newSolution);
+        }
+
         private static void UpdateSolutionProvider(TestWorkspace workspace, Solution solution)
         {
             var provider = (TestLspSolutionProvider)workspace.ExportProvider.GetExportedValue<ILspSolutionProvider>();
@@ -274,7 +319,7 @@ private static void UpdateSolutionProvider(TestWorkspace workspace, Solution sol
                 foreach (var (name, spans) in testDocument.AnnotatedSpans)
                 {
                     var locationsForName = locations.GetOrValue(name, new List<LSP.Location>());
-                    locationsForName.AddRange(spans.Select(span => ProtocolConversions.TextSpanToLocation(span, text, new Uri(document.FilePath))));
+                    locationsForName.AddRange(spans.Select(span => ConvertTextSpanWithTextToLocation(span, text, new Uri(document.FilePath))));
 
                     // Linked files will return duplicate annotated Locations for each document that links to the same file.
                     // Since the test output only cares about the actual file, make sure we de-dupe before returning.
@@ -283,6 +328,17 @@ private static void UpdateSolutionProvider(TestWorkspace workspace, Solution sol
             }
 
             return locations;
+
+            static LSP.Location ConvertTextSpanWithTextToLocation(TextSpan span, SourceText text, Uri documentUri)
+            {
+                var location = new LSP.Location
+                {
+                    Uri = documentUri,
+                    Range = ProtocolConversions.TextSpanToRange(span, text),
+                };
+
+                return location;
+            }
         }
 
         // Private protected because LanguageServerProtocol is internal

src/Features/LanguageServer/Protocol/Extensions/ProtocolConversions.cs

@@ -2,11 +2,15 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+#nullable enable
+
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.DocumentHighlighting;
+using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.NavigateTo;
 using Microsoft.CodeAnalysis.Tags;
 using Microsoft.CodeAnalysis.Text;
@@ -61,7 +65,7 @@ internal static class ProtocolConversions
             { WellKnownTags.NuGet, LSP.CompletionItemKind.Text }
         };
 
-        public static Uri GetUriFromFilePath(string filePath)
+        public static Uri GetUriFromFilePath(string? filePath)
         {
             if (filePath is null)
             {
@@ -116,48 +120,74 @@ public static LSP.Range TextSpanToRange(TextSpan textSpan, SourceText text)
             return LinePositionToRange(linePosSpan);
         }
 
-        public static Task<LSP.Location> DocumentSpanToLocationAsync(DocumentSpan documentSpan, CancellationToken cancellationToken)
+        public static Task<LSP.Location?> DocumentSpanToLocationAsync(DocumentSpan documentSpan, CancellationToken cancellationToken)
             => TextSpanToLocationAsync(documentSpan.Document, documentSpan.SourceSpan, cancellationToken);
 
-        public static async Task<LSP.LocationWithText> DocumentSpanToLocationWithTextAsync(DocumentSpan documentSpan, ClassifiedTextElement text, CancellationToken cancellationToken)
+        public static async Task<LSP.LocationWithText?> DocumentSpanToLocationWithTextAsync(DocumentSpan documentSpan, ClassifiedTextElement text, CancellationToken cancellationToken)
         {
-            var sourceText = await documentSpan.Document.GetTextAsync(cancellationToken).ConfigureAwait(false);
+            var location = await TextSpanToLocationAsync(documentSpan.Document, documentSpan.SourceSpan, cancellationToken).ConfigureAwait(false);
 
-            var locationWithText = new LSP.LocationWithText
+            return location == null ? null : new LSP.LocationWithText
             {
-                Uri = documentSpan.Document.GetURI(),
-                Range = TextSpanToRange(documentSpan.SourceSpan, sourceText),
+                Uri = location.Uri,
+                Range = location.Range,
                 Text = text
             };
-
-            return locationWithText;
         }
 
-        public static LSP.Location RangeToLocation(LSP.Range range, string uriString)
+        public static async Task<LSP.Location?> TextSpanToLocationAsync(Document document, TextSpan textSpan, CancellationToken cancellationToken)
         {
-            return new LSP.Location()
+            var spanMappingService = document.Services.GetService<ISpanMappingService>();
+
+            if (spanMappingService == null)
             {
-                Range = range,
-                Uri = new Uri(uriString)
-            };
-        }
+                return await ConvertTextSpanToLocation(document, textSpan, cancellationToken).ConfigureAwait(false);
+            }
 
-        public static async Task<LSP.Location> TextSpanToLocationAsync(Document document, TextSpan span, CancellationToken cancellationToken)
-        {
-            var text = await document.GetTextAsync(cancellationToken).ConfigureAwait(false);
+            var mappedSpanResult = await spanMappingService.MapSpansAsync(document, SpecializedCollections.SingletonEnumerable(textSpan), cancellationToken).ConfigureAwait(false);
+            if (mappedSpanResult.IsDefaultOrEmpty)
+            {
+                return await ConvertTextSpanToLocation(document, textSpan, cancellationToken).ConfigureAwait(false);
+            }
 
-            return TextSpanToLocation(span, text, document.GetURI());
-        }
+            var mappedSpan = mappedSpanResult.Single();
+            if (mappedSpan.IsDefault)
+            {
+                return null;
+            }
 
-        public static LSP.Location TextSpanToLocation(TextSpan span, SourceText text, Uri documentUri)
-        {
-            var location = new LSP.Location
+            return new LSP.Location
             {
-                Uri = documentUri,
-                Range = TextSpanToRange(span, text),
+                Uri = GetUriFromFilePath(mappedSpan.FilePath),
+                Range = MappedSpanResultToRange(mappedSpan)
             };
 
-            return location;
+            static async Task<LSP.Location> ConvertTextSpanToLocation(Document document, TextSpan span, CancellationToken cancellationToken)
+            {
+                var text = await document.GetTextAsync(cancellationToken).ConfigureAwait(false);
+
+                return ConvertTextSpanWithTextToLocation(span, text, document.GetURI());
+            }
+
+            static LSP.Range MappedSpanResultToRange(MappedSpanResult mappedSpanResult)
+            {
+                return new LSP.Range
+                {
+                    Start = LinePositionToPosition(mappedSpanResult.LinePositionSpan.Start),
+                    End = LinePositionToPosition(mappedSpanResult.LinePositionSpan.End)
+                };
+            }
+
+            static LSP.Location ConvertTextSpanWithTextToLocation(TextSpan span, SourceText text, Uri documentUri)
+            {
+                var location = new LSP.Location
+                {
+                    Uri = documentUri,
+                    Range = TextSpanToRange(span, text),
+                };
+
+                return location;
+            }
         }
 
         public static LSP.DiagnosticSeverity DiagnosticSeverityToLspDiagnositcSeverity(DiagnosticSeverity severity)

src/Features/LanguageServer/Protocol/Handler/Definitions/AbstractGoToDefinitionHandler.cs

@@ -13,6 +13,8 @@
 using Microsoft.CodeAnalysis.MetadataAsSource;
 using Microsoft.CodeAnalysis.Navigation;
 using Microsoft.CodeAnalysis.PooledObjects;
+using Microsoft.CodeAnalysis.Shared.Extensions;
+using Roslyn.Utilities;
 using LSP = Microsoft.VisualStudio.LanguageServer.Protocol;
 
 namespace Microsoft.CodeAnalysis.LanguageServer.Handler
@@ -47,12 +49,8 @@ protected async Task<LSP.Location[]> GetDefinitionAsync(LSP.TextDocumentPosition
                         continue;
                     }
 
-                    var definitionText = await definition.Document.GetTextAsync(cancellationToken).ConfigureAwait(false);
-                    locations.Add(new LSP.Location
-                    {
-                        Uri = definition.Document.GetURI(),
-                        Range = ProtocolConversions.TextSpanToRange(definition.SourceSpan, definitionText),
-                    });
+                    var location = await ProtocolConversions.TextSpanToLocationAsync(definition.Document, definition.SourceSpan, cancellationToken).ConfigureAwait(false);
+                    locations.AddIfNotNull(location);
                 }
             }
             else if (document.SupportsSemanticModel && _metadataAsSourceFileService != null)

src/Features/LanguageServer/Protocol/Handler/References/FindImplementationsHandler.cs

@@ -50,11 +50,11 @@ public FindImplementationsHandler(ILspSolutionProvider solutionProvider) : base(
                 {
                     if (clientCapabilities?.HasVisualStudioLspCapability() == true)
                     {
-                        locations.Add(await ProtocolConversions.DocumentSpanToLocationWithTextAsync(sourceSpan, text, cancellationToken).ConfigureAwait(false));
+                        locations.AddIfNotNull(await ProtocolConversions.DocumentSpanToLocationWithTextAsync(sourceSpan, text, cancellationToken).ConfigureAwait(false));
                     }
                     else
                     {
-                        locations.Add(await ProtocolConversions.DocumentSpanToLocationAsync(sourceSpan, cancellationToken).ConfigureAwait(false));
+                        locations.AddIfNotNull(await ProtocolConversions.DocumentSpanToLocationAsync(sourceSpan, cancellationToken).ConfigureAwait(false));
                     }
                 }
             }

src/Features/LanguageServer/ProtocolUnitTests/Definitions/GoToDefinitionTests.cs

@@ -2,6 +2,8 @@
 // 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.Immutable;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -58,6 +60,31 @@ class B
             AssertLocationsEqual(locations["definition"], results);
         }
 
+        [Fact]
+        public async Task TestGotoDefinitionAsync_MappedFile()
+        {
+            var markup =
+@"class A
+{
+    string aString = 'hello';
+    void M()
+    {
+        var len = aString.Length;
+    }
+}";
+            using var workspace = CreateTestWorkspace(string.Empty, out var _);
+
+            AddMappedDocument(workspace, markup);
+
+            var position = new LSP.Position { Line = 5, Character = 18 };
+            var results = await RunGotoDefinitionAsync(workspace.CurrentSolution, new LSP.Location
+            {
+                Uri = new Uri($"C:\\{TestSpanMapper.GeneratedFileName}"),
+                Range = new LSP.Range { Start = position, End = position }
+            });
+            AssertLocationsEqual(ImmutableArray.Create(TestSpanMapper.MappedFileLocation), results);
+        }
+
         [Fact]
         public async Task TestGotoDefinitionAsync_InvalidLocation()
         {

src/Features/LanguageServer/ProtocolUnitTests/References/FindImplementationsTests.cs

@@ -2,6 +2,8 @@
 // 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.Immutable;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -62,6 +64,33 @@ class A : IA
             AssertLocationsEqual(locations["implementation"], results);
         }
 
+        [Fact]
+        public async Task TestFindImplementationAsync_MappedFile()
+        {
+            var markup =
+@"interface IA
+{
+    void M();
+}
+class A : IA
+{
+    void IA.M()
+    {
+    }
+}";
+            using var workspace = CreateTestWorkspace(string.Empty, out var _);
+
+            AddMappedDocument(workspace, markup);
+
+            var position = new LSP.Position { Line = 2, Character = 9 };
+            var results = await RunFindImplementationAsync(workspace.CurrentSolution, new LSP.Location
+            {
+                Uri = new Uri($"C:\\{TestSpanMapper.GeneratedFileName}"),
+                Range = new LSP.Range { Start = position, End = position }
+            });
+            AssertLocationsEqual(ImmutableArray.Create(TestSpanMapper.MappedFileLocation), results);
+        }
+
         [Fact]
         public async Task TestFindImplementationAsync_InvalidLocation()
         {

src/VisualStudio/LiveShare/Impl/AbstractGoToDefinitionWithFindUsagesServiceHandler.cs

@@ -6,6 +6,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -15,6 +16,7 @@
 using Microsoft.CodeAnalysis.FindSymbols;
 using Microsoft.CodeAnalysis.LanguageServer;
 using Microsoft.CodeAnalysis.MetadataAsSource;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.VisualStudio.LiveShare.LanguageServices;
 using Microsoft.VisualStudio.Shell;
 using LSP = Microsoft.VisualStudio.LanguageServer.Protocol;
@@ -49,7 +51,7 @@ public async Task<object> HandleAsync(LSP.TextDocumentPositionParams request, Re
             var locations = await GetDefinitionsWithFindUsagesServiceAsync(document, position, cancellationToken).ConfigureAwait(false);
 
             // No definition found - see if we can get metadata as source but that's only applicable for C#\VB.
-            if ((locations.Count == 0) && document.SupportsSemanticModel && this._metadataAsSourceService != null)
+            if ((locations.Count() == 0) && document.SupportsSemanticModel && this._metadataAsSourceService != null)
             {
                 var symbol = await SymbolFinder.FindSymbolAtPositionAsync(document, position, cancellationToken).ConfigureAwait(false);
                 if (symbol?.Locations.FirstOrDefault().IsInMetadata == true)
@@ -57,11 +59,10 @@ public async Task<object> HandleAsync(LSP.TextDocumentPositionParams request, Re
                     var declarationFile = await this._metadataAsSourceService.GetGeneratedFileAsync(document.Project, symbol, false, cancellationToken).ConfigureAwait(false);
 
                     var linePosSpan = declarationFile.IdentifierLocation.GetLineSpan().Span;
-                    locations.Add(new LSP.Location
+                    return new LSP.Location[]
                     {
-                        Uri = new Uri(declarationFile.FilePath),
-                        Range = ProtocolConversions.LinePositionToRange(linePosSpan)
-                    });
+                        new LSP.Location { Uri = new Uri(declarationFile.FilePath), Range = ProtocolConversions.LinePositionToRange(linePosSpan) }
+                    };
                 }
             }
 
@@ -72,7 +73,7 @@ public async Task<object> HandleAsync(LSP.TextDocumentPositionParams request, Re
         ///  Using the find usages service is more expensive than using the definitions service because a lot of unnecessary information is computed. However,
         ///  TypeScript doesn't provide an <see cref="IGoToDefinitionService"/> implementation that will return definitions so we must use <see cref="IFindUsagesService"/>.
         /// </summary>
-        private async Task<List<LSP.Location>> GetDefinitionsWithFindUsagesServiceAsync(Document document, int pos, CancellationToken cancellationToken)
+        private async Task<ImmutableArray<LSP.Location>> GetDefinitionsWithFindUsagesServiceAsync(Document document, int pos, CancellationToken cancellationToken)
         {
             var findUsagesService = document.Project.LanguageServices.GetRequiredService<IFindUsagesService>();
 
@@ -83,7 +84,7 @@ private async Task<List<LSP.Location>> GetDefinitionsWithFindUsagesServiceAsync(
             await ThreadHelper.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
             await findUsagesService.FindReferencesAsync(document, pos, context).ConfigureAwait(false);
 
-            var locations = new List<LSP.Location>();
+            using var _ = ArrayBuilder<LSP.Location>.GetInstance(out var locations);
 
             var definitions = context.GetDefinitions();
             if (definitions != null)
@@ -92,12 +93,12 @@ private async Task<List<LSP.Location>> GetDefinitionsWithFindUsagesServiceAsync(
                 {
                     foreach (var docSpan in definition.SourceSpans)
                     {
-                        locations.Add(await ProtocolConversions.DocumentSpanToLocationAsync(docSpan, cancellationToken).ConfigureAwait(false));
+                        locations.AddIfNotNull(await ProtocolConversions.DocumentSpanToLocationAsync(docSpan, cancellationToken).ConfigureAwait(false));
                     }
                 }
             }
 
-            return locations;
+            return locations.ToImmutable();
         }
     }
 }

src/Workspaces/Remote/ServiceHub/Services/LanguageServer.cs

@@ -138,18 +138,12 @@ private static async Task SearchDocumentAndReportSymbolsAsync(Document document,
 
             foreach (var result in results)
             {
-                var text = await result.NavigableItem.Document.GetTextAsync(cancellationToken).ConfigureAwait(false);
-
                 symbols.Add(new VSSymbolInformation()
                 {
                     Name = result.Name,
                     ContainerName = result.AdditionalInformation,
                     Kind = ProtocolConversions.NavigateToKindToSymbolKind(result.Kind),
-                    Location = new LSP.Location()
-                    {
-                        Uri = result.NavigableItem.Document.GetURI(),
-                        Range = ProtocolConversions.TextSpanToRange(result.NavigableItem.SourceSpan, text)
-                    },
+                    Location = await ProtocolConversions.TextSpanToLocationAsync(result.NavigableItem.Document, result.NavigableItem.SourceSpan, cancellationToken).ConfigureAwait(false),
                     Icon = new VisualStudio.Text.Adornments.ImageElement(result.NavigableItem.Glyph.GetImageId())
                 });
             }