Fix Code Lens around source generated files (#79992)

To best understand this fix, we first need to discuss how CodeLens works
in Visual Studio. The editor calls our implementation of
IAsyncCodeLensDataPointProvider which exists in a ServiceHub process
created for CodeLens which is not our usual ServiceHub process. That
implementation uses the callback feature to marshal a call back to
devenv, where from there we then do an OOP operation to our ServiceHub
process via the usual mechanism. This is where we compute the actual
results, which then get marshalled back to devenv and then back to the
CodeLens ServiceHub process, where we finally send the data via the
CodeLens APIs.

The core type of this is a ReferenceLocationDescriptor which represented
a reference location in a specific location in a document. However,
because this struct returned from our ServiceHub to devenv, and then
from devenv to the CodeLens service hub, it needs to be serialized and
deserialized twice. The serialization channels are different though, so
whereas a DocumentId can serialize between devenv and our ServiceHub
process, it would get tripped up when it serializes between devenv and
the CodeLens ServiceHub process. Because of this, the type instead
decomposed a DocumentId into the GUIDs and manually
serialized/deserialized it.

This serialization/deserialization was broken once DocumentIds got a "is
source generated" bit, since we'd lose that bit which would break
things. I could have added that bit in and updated the serialization but
the question we hit was "why can't I just pass a DocumentId?" Well, the
answer was because we might try sending that DocumentId to the CodeLens
ServiceHub process, but upon closer inspection, the result of that was
never used. The flow was that we'd take the results with document IDs
out of our ServiceHub process, then do some filtering and mapping for
Razor scenarios (which used those document IDs); at this point the
results don't need document IDs anymore. So to do this I split out the
"mapping" as a separate operation for clarity, and then used different
types, one with the ID and one without, for the different flows.

There is still one bug left here: code lens will now show references in
source generated documents, but you can't open those since we aren't
returning a file path VS will understand there. That's fixable too, but
this is a significant improvement so let's get this done first. I'm not
sure if fixing that might require other code changes in the CodeLens
system...

Fixes #79699
Fixes #76608
Fixes #74104

Author: jasonmalinowski

src/EditorFeatures/TestUtilities/CodeLens/AbstractCodeLensTest.cs

@@ -32,7 +32,7 @@ protected static async Task RunCountTest(XElement input, int cap = 0)
                 foreach (var span in annotatedSpan.Value)
                 {
                     var declarationSyntaxNode = syntaxNode.FindNode(span);
-                    var result = await new CodeLensReferencesService().GetReferenceCountAsync(workspace.CurrentSolution, annotatedDocument.Id,
+                    var result = await workspace.Services.GetRequiredService<ICodeLensReferencesService>().GetReferenceCountAsync(workspace.CurrentSolution, annotatedDocument.Id,
                         declarationSyntaxNode, cap, CancellationToken.None);
                     Assert.NotNull(result);
                     Assert.Equal(expected, result.Value.Count);
@@ -59,7 +59,7 @@ protected static async Task RunReferenceTest(XElement input)
                 foreach (var span in annotatedSpan.Value)
                 {
                     var declarationSyntaxNode = syntaxNode.FindNode(span);
-                    var result = await new CodeLensReferencesService().FindReferenceLocationsAsync(workspace.CurrentSolution,
+                    var result = await workspace.Services.GetRequiredService<ICodeLensReferencesService>().FindReferenceLocationsAsync(workspace.CurrentSolution,
                         annotatedDocument.Id, declarationSyntaxNode, CancellationToken.None);
                     Assert.True(result.HasValue);
                     Assert.Equal(expected, result.Value.Length);
@@ -85,7 +85,7 @@ protected static async Task RunMethodReferenceTest(XElement input)
                 foreach (var span in annotatedSpan.Value)
                 {
                     var declarationSyntaxNode = syntaxNode.FindNode(span);
-                    var result = await new CodeLensReferencesService().FindReferenceMethodsAsync(workspace.CurrentSolution,
+                    var result = await workspace.Services.GetRequiredService<ICodeLensReferencesService>().FindReferenceMethodsAsync(workspace.CurrentSolution,
                         annotatedDocument.Id, declarationSyntaxNode, CancellationToken.None);
                     Assert.True(result.HasValue);
                     Assert.Equal(expected, result.Value.Length);
@@ -111,7 +111,7 @@ protected static async Task RunFullyQualifiedNameTest(XElement input)
                 foreach (var span in annotatedSpan.Value)
                 {
                     var declarationSyntaxNode = syntaxNode.FindNode(span);
-                    var actual = await new CodeLensReferencesService().GetFullyQualifiedNameAsync(workspace.CurrentSolution,
+                    var actual = await workspace.Services.GetRequiredService<ICodeLensReferencesService>().GetFullyQualifiedNameAsync(workspace.CurrentSolution,
                         annotatedDocument.Id, declarationSyntaxNode, CancellationToken.None);
                     Assert.Equal(expected, actual);
                 }

src/Features/Core/Portable/CodeLens/CodeLensHelpers.cs

@@ -16,8 +16,16 @@ internal static class CodeLensHelpers
         if (TryGetGuid("RoslynDocumentIdGuid", out var documentIdGuid) &&
             TryGetGuid("RoslynProjectIdGuid", out var projectIdGuid))
         {
+            var isSourceGenerated = false;
+
+            if (descriptorProperties.TryGetValue("RoslynDocumentIsSourceGenerated", out var isSourceGeneratedObj))
+            {
+                if (isSourceGeneratedObj is bool isSourceGeneratedBoolean)
+                    isSourceGenerated = isSourceGeneratedBoolean;
+            }
+
             var projectId = ProjectId.CreateFromSerialized(projectIdGuid);
-            return DocumentId.CreateFromSerialized(projectId, documentIdGuid);
+            return DocumentId.CreateFromSerialized(projectId, documentIdGuid, isSourceGenerated, debugName: null);
         }
 
         return null;

src/Features/Core/Portable/CodeLens/CodeLensReferencesService.cs

@@ -10,9 +10,11 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Classification;
 using Microsoft.CodeAnalysis.FindSymbols;
 using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.LanguageService;
+using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Text;
@@ -98,7 +100,7 @@ public async ValueTask<VersionStamp> GetProjectCodeLensVersionAsync(Solution sol
             maxSearchResults, cancellationToken).ConfigureAwait(false);
     }
 
-    private static async Task<ReferenceLocationDescriptor> GetDescriptorOfEnclosingSymbolAsync(Solution solution, Location location, CancellationToken cancellationToken)
+    private static async Task<ReferenceLocationDescriptorAndDocument> GetDescriptorOfEnclosingSymbolAsync(Solution solution, Location location, CancellationToken cancellationToken)
     {
         var document = solution.GetDocument(location.SourceTree);
 
@@ -130,41 +132,36 @@ private static async Task<ReferenceLocationDescriptor> GetDescriptorOfEnclosingS
         var spanStart = token.Span.Start - textLine.Span.Start;
         var line = textLine.ToString();
 
-        var beforeLine1 = textLine.LineNumber > 0 ? text.Lines[textLine.LineNumber - 1].ToString() : string.Empty;
-        var beforeLine2 = textLine.LineNumber - 1 > 0
-            ? text.Lines[textLine.LineNumber - 2].ToString()
-            : string.Empty;
-        var afterLine1 = textLine.LineNumber < text.Lines.Count - 1
-            ? text.Lines[textLine.LineNumber + 1].ToString()
-            : string.Empty;
-        var afterLine2 = textLine.LineNumber + 1 < text.Lines.Count - 1
-            ? text.Lines[textLine.LineNumber + 2].ToString()
-            : string.Empty;
+        var beforeLine1 = GetLineTextOrEmpty(text.Lines, textLine.LineNumber - 1);
+        var beforeLine2 = GetLineTextOrEmpty(text.Lines, textLine.LineNumber - 2);
+        var afterLine1 = GetLineTextOrEmpty(text.Lines, textLine.LineNumber + 1);
+        var afterLine2 = GetLineTextOrEmpty(text.Lines, textLine.LineNumber + 2);
         var referenceSpan = new TextSpan(spanStart, token.Span.Length);
 
         var symbol = semanticModel.GetDeclaredSymbol(node, cancellationToken);
         var glyph = symbol?.GetGlyph();
         var startLinePosition = location.GetLineSpan().StartLinePosition;
-        var documentId = solution.GetDocument(location.SourceTree)?.Id;
-
-        return new ReferenceLocationDescriptor(
-            longName,
-            semanticModel.Language,
-            glyph,
-            token.Span.Start,
-            token.Span.Length,
-            startLinePosition.Line,
-            startLinePosition.Character,
-            documentId.ProjectId.Id,
-            documentId.Id,
-            document.FilePath,
-            line.TrimEnd(),
-            referenceSpan.Start,
-            referenceSpan.Length,
-            beforeLine1.TrimEnd(),
-            beforeLine2.TrimEnd(),
-            afterLine1.TrimEnd(),
-            afterLine2.TrimEnd());
+
+        return new ReferenceLocationDescriptorAndDocument
+        {
+            Descriptor = new ReferenceLocationDescriptor(
+                longName,
+                semanticModel.Language,
+                glyph,
+                token.Span.Start,
+                token.Span.Length,
+                startLinePosition.Line,
+                startLinePosition.Character,
+                document.FilePath,
+                line.TrimEnd(),
+                referenceSpan.Start,
+                referenceSpan.Length,
+                beforeLine1.TrimEnd(),
+                beforeLine2.TrimEnd(),
+                afterLine1.TrimEnd(),
+                afterLine2.TrimEnd()),
+            DocumentId = document.Id
+        };
     }
 
     private static SyntaxNode GetEnclosingCodeElementNode(Document document, SyntaxToken token, ICodeLensDisplayInfoService langServices, CancellationToken cancellationToken)
@@ -199,7 +196,7 @@ private static SyntaxNode GetEnclosingCodeElementNode(Document document, SyntaxT
         return langServices.GetDisplayNode(node);
     }
 
-    public async Task<ImmutableArray<ReferenceLocationDescriptor>?> FindReferenceLocationsAsync(Solution solution, DocumentId documentId, SyntaxNode syntaxNode, CancellationToken cancellationToken)
+    public async Task<ImmutableArray<ReferenceLocationDescriptorAndDocument>?> FindReferenceLocationsAsync(Solution solution, DocumentId documentId, SyntaxNode syntaxNode, CancellationToken cancellationToken)
     {
         return await FindAsync(solution, documentId, syntaxNode,
             async progress =>
@@ -214,6 +211,137 @@ private static SyntaxNode GetEnclosingCodeElementNode(Document document, SyntaxT
             }, onCapped: null, searchCap: 0, cancellationToken: cancellationToken).ConfigureAwait(false);
     }
 
+    public async Task<ImmutableArray<ReferenceLocationDescriptor>> MapReferenceLocationsAsync(Solution solution, ImmutableArray<ReferenceLocationDescriptorAndDocument> referenceLocations, ClassificationOptions classificationOptions, CancellationToken cancellationToken)
+    {
+        using var _ = ArrayBuilder<ReferenceLocationDescriptor>.GetInstance(out var list);
+        foreach (var descriptorAndDocument in referenceLocations)
+        {
+            var document = await solution.GetDocumentAsync(descriptorAndDocument.DocumentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+            if (document == null)
+            {
+                continue;
+            }
+
+            var descriptor = descriptorAndDocument.Descriptor;
+            var span = new TextSpan(descriptor.SpanStart, descriptor.SpanLength);
+            var results = await SpanMappingHelper.TryGetMappedSpanResultAsync(document, [span], cancellationToken).ConfigureAwait(false);
+            if (results is null)
+            {
+                // for normal document, just add one as they are
+                list.Add(descriptor);
+                continue;
+            }
+
+            var mappedSpans = results.GetValueOrDefault();
+
+            // external component violated contracts. the mapper should preserve input order/count. 
+            // since we gave in 1 span, it should return 1 span back
+            Contract.ThrowIfTrue(mappedSpans.IsDefaultOrEmpty);
+
+            var result = mappedSpans[0];
+            if (result.IsDefault)
+            {
+                // it is allowed for mapper to return default 
+                // if it can't map the given span to any usable span
+                continue;
+            }
+
+            var excerpter = document.DocumentServiceProvider.GetService<IDocumentExcerptService>();
+            if (excerpter == null)
+            {
+                if (document.IsRazorSourceGeneratedDocument())
+                {
+                    // HACK: Razor doesn't have has a workspace level excerpt service, but if we just return a simple descriptor here,
+                    // the user at least sees something, can navigate, and Razor can improve this later if necessary. Until
+                    // https://github.com/dotnet/roslyn/issues/79699 is fixed this won't get hit anyway.
+                    list.Add(new ReferenceLocationDescriptor(
+                        descriptor.LongDescription,
+                        descriptor.Language,
+                        descriptor.Glyph,
+                        result.Span.Start,
+                        result.Span.Length,
+                        result.LinePositionSpan.Start.Line,
+                        result.LinePositionSpan.Start.Character,
+                        result.FilePath,
+                        descriptor.ReferenceLineText,
+                        descriptor.ReferenceStart,
+                        descriptor.ReferenceLength,
+                        "",
+                        "",
+                        "",
+                        ""));
+                }
+                continue;
+            }
+
+            var referenceExcerpt = await excerpter.TryExcerptAsync(document, span, ExcerptMode.SingleLine, classificationOptions, cancellationToken).ConfigureAwait(false);
+            var tooltipExcerpt = await excerpter.TryExcerptAsync(document, span, ExcerptMode.Tooltip, classificationOptions, cancellationToken).ConfigureAwait(false);
+
+            var (text, start, length) = GetReferenceInfo(referenceExcerpt, descriptor);
+            var (before1, before2, after1, after2) = GetReferenceTexts(referenceExcerpt, tooltipExcerpt, descriptor);
+
+            list.Add(new ReferenceLocationDescriptor(
+                descriptor.LongDescription,
+                descriptor.Language,
+                descriptor.Glyph,
+                result.Span.Start,
+                result.Span.Length,
+                result.LinePositionSpan.Start.Line,
+                result.LinePositionSpan.Start.Character,
+                result.FilePath,
+                text,
+                start,
+                length,
+                before1,
+                before2,
+                after1,
+                after2));
+        }
+
+        return list.ToImmutableAndClear();
+    }
+
+    private static (string text, int start, int length) GetReferenceInfo(ExcerptResult? reference, ReferenceLocationDescriptor descriptor)
+    {
+        if (reference.HasValue)
+        {
+            return (reference.Value.Content.ToString().TrimEnd(),
+                    reference.Value.MappedSpan.Start,
+                    reference.Value.MappedSpan.Length);
+        }
+
+        return (descriptor.ReferenceLineText, descriptor.ReferenceStart, descriptor.ReferenceLength);
+    }
+
+    private static (string before1, string before2, string after1, string after2) GetReferenceTexts(ExcerptResult? reference, ExcerptResult? tooltip, ReferenceLocationDescriptor descriptor)
+    {
+        if (reference == null || tooltip == null)
+        {
+            return (descriptor.BeforeReferenceText1, descriptor.BeforeReferenceText2, descriptor.AfterReferenceText1, descriptor.AfterReferenceText2);
+        }
+
+        var lines = tooltip.Value.Content.Lines;
+        var mappedLine = lines.GetLineFromPosition(tooltip.Value.MappedSpan.Start);
+        var index = mappedLine.LineNumber;
+        if (index < 0)
+        {
+            return (descriptor.BeforeReferenceText1, descriptor.BeforeReferenceText2, descriptor.AfterReferenceText1, descriptor.AfterReferenceText2);
+        }
+
+        return (GetLineTextOrEmpty(lines, index - 1), GetLineTextOrEmpty(lines, index - 2),
+                GetLineTextOrEmpty(lines, index + 1), GetLineTextOrEmpty(lines, index + 2));
+    }
+
+    private static string GetLineTextOrEmpty(TextLineCollection lines, int index)
+    {
+        if (index < 0 || index >= lines.Count)
+        {
+            return string.Empty;
+        }
+
+        return lines[index].ToString().TrimEnd();
+    }
+
     private static ISymbol GetEnclosingMethod(SemanticModel semanticModel, Location location, CancellationToken cancellationToken)
     {
         var enclosingSymbol = semanticModel.GetEnclosingSymbol(location.SourceSpan.Start, cancellationToken);

src/Features/Core/Portable/CodeLens/ICodeLensReferencesService.cs

@@ -5,6 +5,7 @@
 using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Classification;
 using Microsoft.CodeAnalysis.Host;
 
 namespace Microsoft.CodeAnalysis.CodeLens;
@@ -25,7 +26,12 @@ internal interface ICodeLensReferencesService : IWorkspaceService
     /// <summary>
     /// Given a document and syntax node, returns a collection of locations where the located node is referenced.
     /// </summary>
-    Task<ImmutableArray<ReferenceLocationDescriptor>?> FindReferenceLocationsAsync(Solution solution, DocumentId documentId, SyntaxNode? syntaxNode, CancellationToken cancellationToken);
+    Task<ImmutableArray<ReferenceLocationDescriptorAndDocument>?> FindReferenceLocationsAsync(Solution solution, DocumentId documentId, SyntaxNode? syntaxNode, CancellationToken cancellationToken);
+
+    /// <summary>
+    /// Maps the list of referenced locations through span mapping services to the final list of locations; locations may be in files that aren't regular documents.
+    /// </summary>
+    Task<ImmutableArray<ReferenceLocationDescriptor>> MapReferenceLocationsAsync(Solution solution, ImmutableArray<ReferenceLocationDescriptorAndDocument> referenceLocations, ClassificationOptions classificationOptions, CancellationToken cancellationToken);
 
     /// <summary>
     /// Given a document and syntax node, returns a collection of locations of methods that refer to the located node.

src/Features/Core/Portable/CodeLens/IRemoteCodeLensReferencesService.cs

@@ -12,7 +12,7 @@ namespace Microsoft.CodeAnalysis.CodeLens;
 internal interface IRemoteCodeLensReferencesService
 {
     ValueTask<ReferenceCount?> GetReferenceCountAsync(Checksum solutionChecksum, DocumentId documentId, TextSpan textSpan, int maxResultCount, CancellationToken cancellationToken);
-    ValueTask<ImmutableArray<ReferenceLocationDescriptor>?> FindReferenceLocationsAsync(Checksum solutionChecksum, DocumentId documentId, TextSpan textSpan, CancellationToken cancellationToken);
+    ValueTask<ImmutableArray<ReferenceLocationDescriptorAndDocument>?> FindReferenceLocationsAsync(Checksum solutionChecksum, DocumentId documentId, TextSpan textSpan, CancellationToken cancellationToken);
     ValueTask<ImmutableArray<ReferenceMethodDescriptor>?> FindReferenceMethodsAsync(Checksum solutionChecksum, DocumentId documentId, TextSpan textSpan, CancellationToken cancellationToken);
     ValueTask<string?> GetFullyQualifiedNameAsync(Checksum solutionChecksum, DocumentId documentId, TextSpan textSpan, CancellationToken cancellationToken);
 }