Razor span mapping in LSP for cohosting (#79677)

Follow up to #79604
Part of dotnet/razor#9519

This is the follow up to the above PR to support rename in LSP, but
thought I'd make the API slightly stronger while I was here since we
haven't consumed it on Razor yet. I also misunderstood the span mapping
API, and it needing to return an array of the same length as the input.

Author: davidwengier

src/LanguageServer/Protocol/Extensions/ProtocolConversions.cs

@@ -1015,6 +1015,20 @@ static string GetStyledText(TaggedText taggedText, bool isInCodeBlock)
             return null;
         }
 
+        if (textDocument is SourceGeneratedDocument sourceGeneratedDocument &&
+            textDocument.Project.Solution.Services.GetService<ISourceGeneratedDocumentSpanMappingService>() is { } sourceGeneratedSpanMappingService)
+        {
+            var result = await sourceGeneratedSpanMappingService.MapSpansAsync(sourceGeneratedDocument, textSpans, cancellationToken).ConfigureAwait(false);
+            if (result.IsDefaultOrEmpty)
+            {
+                return null;
+            }
+
+            Contract.ThrowIfFalse(textSpans.Length == result.Length,
+                $"The number of input spans {textSpans.Length} should match the number of mapped spans returned {result.Length}");
+            return result;
+        }
+
         var spanMappingService = document.DocumentServiceProvider.GetService<ISpanMappingService>();
         if (spanMappingService == null)
         {

src/LanguageServer/ProtocolUnitTests/Rename/RenameTests.cs

@@ -4,16 +4,20 @@
 
 #nullable disable
 
+using System;
 using System.Collections.Immutable;
+using System.Composition;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.CodeAnalysis.LanguageServer.Handler;
+using Microsoft.CodeAnalysis.Host;
+using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.Test.Utilities;
 using Microsoft.CodeAnalysis.Testing;
 using Microsoft.CodeAnalysis.Text;
+using Microsoft.VisualStudio.TestPlatform.ObjectModel;
 using Roslyn.LanguageServer.Protocol;
 using Roslyn.Test.Utilities;
-using Roslyn.Test.Utilities.TestGenerators;
 using Xunit;
 using Xunit.Abstractions;
 using LSP = Roslyn.LanguageServer.Protocol;
@@ -22,6 +26,8 @@ namespace Microsoft.CodeAnalysis.LanguageServer.UnitTests.Rename;
 
 public sealed class RenameTests(ITestOutputHelper testOutputHelper) : AbstractLanguageServerProtocolTests(testOutputHelper)
 {
+    protected override TestComposition Composition => base.Composition.AddParts(typeof(TestSourceGeneratedDocumentSpanMappingService));
+
     [Theory, CombinatorialData]
     public async Task TestRenameAsync(bool mutatingLspWorkspace)
     {
@@ -269,6 +275,9 @@ void M2()
 
         var results = await RunRenameAsync(testLspServer, CreateRenameParams(renameLocation, renameValue));
         AssertJsonEquals(expectedEdits.Concat(expectedGeneratedEdits), ((TextDocumentEdit[])results.DocumentChanges).SelectMany(e => e.Edits));
+
+        var service = Assert.IsType<TestSourceGeneratedDocumentSpanMappingService>(workspace.Services.GetService<ISourceGeneratedDocumentSpanMappingService>());
+        Assert.True(service.DidMapSpans);
     }
 
     [Theory, CombinatorialData]
@@ -316,6 +325,9 @@ void M2()
 
         var results = await RunRenameAsync(testLspServer, CreateRenameParams(renameLocation, renameValue));
         AssertJsonEquals(expectedEdits.Concat(expectedGeneratedEdits), ((TextDocumentEdit[])results.DocumentChanges).SelectMany(e => e.Edits));
+
+        var service = Assert.IsType<TestSourceGeneratedDocumentSpanMappingService>(workspace.Services.GetService<ISourceGeneratedDocumentSpanMappingService>());
+        Assert.True(service.DidMapSpans);
     }
 
     private static LSP.RenameParams CreateRenameParams(LSP.Location location, string newName)
@@ -330,4 +342,26 @@ private static async Task<WorkspaceEdit> RunRenameAsync(TestLspServer testLspSer
     {
         return await testLspServer.ExecuteRequestAsync<LSP.RenameParams, LSP.WorkspaceEdit>(LSP.Methods.TextDocumentRenameName, renameParams, CancellationToken.None);
     }
+
+    [ExportWorkspaceService(typeof(ISourceGeneratedDocumentSpanMappingService))]
+    [Shared]
+    [method: ImportingConstructor]
+    [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+    private class TestSourceGeneratedDocumentSpanMappingService() : ISourceGeneratedDocumentSpanMappingService
+    {
+        public bool DidMapSpans { get; private set; }
+
+        public Task<ImmutableArray<MappedTextChange>> GetMappedTextChangesAsync(SourceGeneratedDocument oldDocument, SourceGeneratedDocument newDocument, CancellationToken cancellationToken)
+        {
+            throw new System.NotImplementedException();
+        }
+
+        public Task<ImmutableArray<MappedSpanResult>> MapSpansAsync(SourceGeneratedDocument document, ImmutableArray<TextSpan> spans, CancellationToken cancellationToken)
+        {
+            DidMapSpans = true;
+            Assert.True(document.IsRazorSourceGeneratedDocument());
+
+            return Task.FromResult(ImmutableArray<MappedSpanResult>.Empty);
+        }
+    }
 }

src/Tools/ExternalAccess/Razor/Features/IRazorSourceGeneratedDocumentSpanMappingService.cs

@@ -11,7 +11,7 @@ namespace Microsoft.CodeAnalysis.ExternalAccess.Razor
 {
     internal interface IRazorSourceGeneratedDocumentSpanMappingService
     {
-        Task<ImmutableArray<RazorMappedEditResult>> GetMappedTextChangesAsync(Document oldDocument, Document newDocument, CancellationToken cancellationToken);
-        Task<ImmutableArray<RazorMappedSpanResult>> MapSpansAsync(Document document, ImmutableArray<TextSpan> spans, CancellationToken cancellationToken);
+        Task<ImmutableArray<RazorMappedEditResult>> GetMappedTextChangesAsync(SourceGeneratedDocument oldDocument, SourceGeneratedDocument newDocument, CancellationToken cancellationToken);
+        Task<ImmutableArray<RazorMappedSpanResult>> MapSpansAsync(SourceGeneratedDocument document, ImmutableArray<TextSpan> spans, CancellationToken cancellationToken);
     }
 }

src/Tools/ExternalAccess/Razor/Features/RazorGeneratedDocumentIdentity.cs

@@ -9,4 +9,17 @@ namespace Microsoft.CodeAnalysis.ExternalAccess.Razor;
 /// <summary>
 /// Wrapper for <see cref="SourceGeneratedDocumentIdentity" /> and <see cref="SourceGeneratorIdentity" />
 /// </summary>
-internal record struct RazorGeneratedDocumentIdentity(DocumentId DocumentId, string HintName, string FilePath, string GeneratorAssemblyName, string? GeneratorAssemblyPath, Version GeneratorAssemblyVersion, string GeneratorTypeName);
+internal record struct RazorGeneratedDocumentIdentity(DocumentId DocumentId, string HintName, string FilePath, string GeneratorAssemblyName, string? GeneratorAssemblyPath, Version GeneratorAssemblyVersion, string GeneratorTypeName)
+{
+    internal static RazorGeneratedDocumentIdentity Create(SourceGeneratedDocument document)
+        => Create(document.Identity);
+
+    internal static RazorGeneratedDocumentIdentity Create(SourceGeneratedDocumentIdentity identity)
+        => new(identity.DocumentId,
+            identity.HintName,
+            identity.FilePath,
+            identity.Generator.AssemblyName,
+            identity.Generator.AssemblyPath,
+            identity.Generator.AssemblyVersion,
+            identity.Generator.TypeName);
+}

src/Tools/ExternalAccess/Razor/Features/RazorSourceGeneratedDocumentSpanMappingService.cs renamed to src/Tools/ExternalAccess/Razor/Features/RazorSourceGeneratedDocumentSpanMappingServiceWrapper.cs

@@ -4,7 +4,6 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
-using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Composition;
 using System.Threading;
@@ -19,12 +18,12 @@ namespace Microsoft.CodeAnalysis.ExternalAccess.Razor;
 [ExportWorkspaceService(typeof(ISourceGeneratedDocumentSpanMappingService)), Shared]
 [method: ImportingConstructor]
 [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
-internal sealed class RazorSourceGeneratedDocumentSpanMappingService(
+internal sealed class RazorSourceGeneratedDocumentSpanMappingServiceWrapper(
     [Import(AllowDefault = true)] IRazorSourceGeneratedDocumentSpanMappingService? implementation) : ISourceGeneratedDocumentSpanMappingService
 {
     private readonly IRazorSourceGeneratedDocumentSpanMappingService? _implementation = implementation;
 
-    public async Task<ImmutableArray<MappedTextChange>> GetMappedTextChangesAsync(Document oldDocument, Document newDocument, CancellationToken cancellationToken)
+    public async Task<ImmutableArray<MappedTextChange>> GetMappedTextChangesAsync(SourceGeneratedDocument oldDocument, SourceGeneratedDocument newDocument, CancellationToken cancellationToken)
     {
         if (_implementation is null ||
             !oldDocument.IsRazorSourceGeneratedDocument() ||
@@ -56,7 +55,7 @@ public async Task<ImmutableArray<MappedTextChange>> GetMappedTextChangesAsync(Do
         return changesBuilder.ToImmutableAndClear();
     }
 
-    public async Task<ImmutableArray<MappedSpanResult>> MapSpansAsync(Document document, ImmutableArray<TextSpan> spans, CancellationToken cancellationToken)
+    public async Task<ImmutableArray<MappedSpanResult>> MapSpansAsync(SourceGeneratedDocument document, ImmutableArray<TextSpan> spans, CancellationToken cancellationToken)
     {
         if (_implementation is null ||
             !document.IsRazorSourceGeneratedDocument())
@@ -65,18 +64,17 @@ public async Task<ImmutableArray<MappedSpanResult>> MapSpansAsync(Document docum
         }
 
         var mappedSpans = await _implementation.MapSpansAsync(document, spans, cancellationToken).ConfigureAwait(false);
-        if (mappedSpans.IsDefaultOrEmpty)
+        if (mappedSpans.Length != spans.Length)
         {
             return [];
         }
 
         using var _ = ArrayBuilder<MappedSpanResult>.GetInstance(out var spansBuilder);
         foreach (var span in mappedSpans)
         {
-            if (!span.IsDefault)
-            {
-                spansBuilder.Add(new MappedSpanResult(span.FilePath, span.LinePositionSpan, span.Span));
-            }
+            spansBuilder.Add(span.IsDefault
+                ? default
+                : new MappedSpanResult(span.FilePath, span.LinePositionSpan, span.Span));
         }
 
         return spansBuilder.ToImmutableAndClear();

src/Tools/ExternalAccess/Razor/Features/RazorUri.cs

@@ -35,13 +35,6 @@ public static RazorGeneratedDocumentIdentity GetIdentityOfGeneratedDocument(Solu
 
         // Razor only cares about documents from its own generator, but it's better to just send them back the info they
         // need to check on their side, so we can avoid dual insertions if anything changes.
-        return new RazorGeneratedDocumentIdentity(
-            identity.DocumentId,
-            identity.HintName,
-            identity.FilePath,
-            identity.Generator.AssemblyName,
-            identity.Generator.AssemblyPath,
-            identity.Generator.AssemblyVersion,
-            identity.Generator.TypeName);
+        return RazorGeneratedDocumentIdentity.Create(identity);
     }
 }

src/Workspaces/Core/Portable/Workspace/Host/DocumentService/ISourceGeneratedDocumentSpanMappingService.cs

@@ -11,9 +11,9 @@ namespace Microsoft.CodeAnalysis.Host;
 
 internal interface ISourceGeneratedDocumentSpanMappingService : IWorkspaceService
 {
-    Task<ImmutableArray<MappedTextChange>> GetMappedTextChangesAsync(Document oldDocument, Document newDocument, CancellationToken cancellationToken);
+    Task<ImmutableArray<MappedTextChange>> GetMappedTextChangesAsync(SourceGeneratedDocument oldDocument, SourceGeneratedDocument newDocument, CancellationToken cancellationToken);
 
-    Task<ImmutableArray<MappedSpanResult>> MapSpansAsync(Document document, ImmutableArray<TextSpan> spans, CancellationToken cancellationToken);
+    Task<ImmutableArray<MappedSpanResult>> MapSpansAsync(SourceGeneratedDocument document, ImmutableArray<TextSpan> spans, CancellationToken cancellationToken);
 }
 
 internal record struct MappedTextChange(string MappedFilePath, TextChange TextChange);