Allow source generated documents to be mapped by Razor

Author: davidwengier

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

@@ -0,0 +1,17 @@
+// 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.Collections.Immutable;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Text;
+
+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);
+    }
+}

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

@@ -3,16 +3,21 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Runtime.Serialization;
 using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.ExternalAccess.Razor;
 
+[DataContract]
 internal readonly struct RazorMappedSpanResult
 {
+    [DataMember(Order = 0)]
     public readonly string FilePath;
 
+    [DataMember(Order = 1)]
     public readonly LinePositionSpan LinePositionSpan;
 
+    [DataMember(Order = 2)]
     public readonly TextSpan Span;
 
     public RazorMappedSpanResult(string filePath, LinePositionSpan linePositionSpan, TextSpan span)
@@ -30,7 +35,10 @@ public RazorMappedSpanResult(string filePath, LinePositionSpan linePositionSpan,
     public bool IsDefault => FilePath == null;
 }
 
-internal readonly record struct RazorMappedEditResult(string FilePath, TextChange[] TextChanges)
+[DataContract]
+internal readonly record struct RazorMappedEditResult(
+    [property: DataMember(Order = 0)] string FilePath,
+    [property: DataMember(Order = 1)] TextChange[] TextChanges)
 {
     public bool IsDefault => FilePath == null || TextChanges == null;
 }

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

@@ -0,0 +1,85 @@
+
+// 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.Collections.Immutable;
+using System.Composition;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Host;
+using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.PooledObjects;
+using Microsoft.CodeAnalysis.Text;
+
+namespace Microsoft.CodeAnalysis.ExternalAccess.Razor;
+
+[ExportWorkspaceService(typeof(ISourceGeneratedDocumentSpanMappingService)), Shared]
+[method: ImportingConstructor]
+[method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+internal sealed class RazorSourceGeneratedDocumentSpanMappingService(
+    [Import(AllowDefault = true)] IRazorSourceGeneratedDocumentSpanMappingService? implementation) : ISourceGeneratedDocumentSpanMappingService
+{
+    private readonly IRazorSourceGeneratedDocumentSpanMappingService? _implementation = implementation;
+
+    public async Task<ImmutableArray<MappedTextChange>> GetMappedTextChangesAsync(Document oldDocument, Document newDocument, CancellationToken cancellationToken)
+    {
+        if (_implementation is null ||
+            !oldDocument.IsRazorSourceGeneratedDocument() ||
+            !newDocument.IsRazorSourceGeneratedDocument())
+        {
+            return [];
+        }
+
+        var mappedChanges = await _implementation.GetMappedTextChangesAsync(oldDocument, newDocument, cancellationToken).ConfigureAwait(false);
+        if (mappedChanges.IsDefaultOrEmpty)
+        {
+            return [];
+        }
+
+        using var _ = ArrayBuilder<MappedTextChange>.GetInstance(out var changesBuilder);
+        foreach (var change in mappedChanges)
+        {
+            if (change.IsDefault)
+            {
+                continue;
+            }
+
+            foreach (var textChange in change.TextChanges)
+            {
+                changesBuilder.Add(new MappedTextChange(change.FilePath, textChange));
+            }
+        }
+
+        return changesBuilder.ToImmutableAndClear();
+
+    }
+
+    public async Task<ImmutableArray<MappedSpanResult>> MapSpansAsync(Document document, ImmutableArray<TextSpan> spans, CancellationToken cancellationToken)
+    {
+        if (_implementation is null ||
+            !document.IsRazorSourceGeneratedDocument())
+        {
+            return [];
+        }
+
+        var mappedSpans = await _implementation.MapSpansAsync(document, spans, cancellationToken).ConfigureAwait(false);
+        if (mappedSpans.IsDefaultOrEmpty)
+        {
+            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));
+            }
+        }
+
+        return spansBuilder.ToImmutableAndClear();
+    }
+}

src/VisualStudio/Core/Def/ProjectSystem/VisualStudioWorkspaceImpl.cs

@@ -15,6 +15,7 @@
 using EnvDTE;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Diagnostics;
+using Microsoft.CodeAnalysis.EditAndContinue;
 using Microsoft.CodeAnalysis.Editor;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
@@ -626,7 +627,7 @@ internal override void ApplyMappedFileChanges(SolutionChanges solutionChanges)
         // Instead, we invoke this in JTF run which will mitigate deadlocks when the ConfigureAwait(true)
         // tries to switch back to the main thread in the LSP client.
         // Link to LSP client bug for ConfigureAwait(true) - https://devdiv.visualstudio.com/DevDiv/_workitems/edit/1216657
-        var mappedChanges = _threadingContext.JoinableTaskFactory.Run(() => GetMappedTextChangesAsync(solutionChanges));
+        var mappedChanges = _threadingContext.JoinableTaskFactory.Run(() => GetMappedTextChangesAsync(solutionChanges, CancellationToken.None));
 
         // Group the mapped text changes by file, then apply all mapped text changes for the file.
         foreach (var changesForFile in mappedChanges)
@@ -641,7 +642,7 @@ internal override void ApplyMappedFileChanges(SolutionChanges solutionChanges)
 
         return;
 
-        async Task<MultiDictionary<string, (TextChange TextChange, ProjectId ProjectId)>> GetMappedTextChangesAsync(SolutionChanges solutionChanges)
+        async Task<MultiDictionary<string, (TextChange TextChange, ProjectId ProjectId)>> GetMappedTextChangesAsync(SolutionChanges solutionChanges, CancellationToken cancellationToken)
         {
             var filePathToMappedTextChanges = new MultiDictionary<string, (TextChange TextChange, ProjectId ProjectId)>();
             foreach (var projectChanges in solutionChanges.GetProjectChanges())
@@ -656,14 +657,36 @@ internal override void ApplyMappedFileChanges(SolutionChanges solutionChanges)
 
                     var newDocument = projectChanges.NewProject.GetRequiredDocument(changedDocumentId);
                     var mappedTextChanges = await mappingService.GetMappedTextChangesAsync(
-                        oldDocument, newDocument, CancellationToken.None).ConfigureAwait(false);
+                        oldDocument, newDocument, cancellationToken).ConfigureAwait(false);
                     foreach (var (filePath, textChange) in mappedTextChanges)
                     {
                         filePathToMappedTextChanges.Add(filePath, (textChange, projectChanges.ProjectId));
                     }
                 }
             }
 
+            var sourceGeneratedDocumentMappingService = Services.GetService<ISourceGeneratedDocumentSpanMappingService>();
+            if (sourceGeneratedDocumentMappingService is not null)
+            {
+                // Since we're mapping changes to source generated documents, we have to ensure the old solution has run the generators
+                // so the mapper has something to compare to.
+                foreach (var (docId, state) in solutionChanges.NewSolution.CompilationState.FrozenSourceGeneratedDocumentStates.States)
+                {
+                    _ = await solutionChanges.OldSolution.GetRequiredDocumentAsync(docId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+                }
+
+                foreach (var docId in solutionChanges.GetExplicitlyChangedSourceGeneratedDocuments())
+                {
+                    var oldDocument = solutionChanges.OldSolution.GetRequiredSourceGeneratedDocumentForAlreadyGeneratedId(docId);
+                    var newDocument = solutionChanges.NewSolution.GetRequiredSourceGeneratedDocumentForAlreadyGeneratedId(docId);
+                    var mappedTextChanges = await sourceGeneratedDocumentMappingService.GetMappedTextChangesAsync(oldDocument, newDocument, cancellationToken).ConfigureAwait(false);
+                    foreach (var (filePath, textChange) in mappedTextChanges)
+                    {
+                        filePathToMappedTextChanges.Add(filePath, (textChange, docId.ProjectId));
+                    }
+                }
+            }
+
             return filePathToMappedTextChanges;
         }
 

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

@@ -0,0 +1,19 @@
+// 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.Collections.Immutable;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Text;
+
+namespace Microsoft.CodeAnalysis.Host;
+
+internal interface ISourceGeneratedDocumentSpanMappingService : IWorkspaceService
+{
+    Task<ImmutableArray<MappedTextChange>> GetMappedTextChangesAsync(Document oldDocument, Document newDocument, CancellationToken cancellationToken);
+
+    Task<ImmutableArray<MappedSpanResult>> MapSpansAsync(Document document, ImmutableArray<TextSpan> spans, CancellationToken cancellationToken);
+}
+
+internal record struct MappedTextChange(string MappedFilePath, TextChange TextChange);

src/Workspaces/Core/Portable/Workspace/Solution/SolutionChanges.cs

@@ -16,6 +16,9 @@ public readonly struct SolutionChanges
     private readonly Solution _newSolution;
     private readonly Solution _oldSolution;
 
+    internal Solution OldSolution => _oldSolution;
+    internal Solution NewSolution => _newSolution;
+
     internal SolutionChanges(Solution newSolution, Solution oldSolution)
     {
         _newSolution = newSolution;