Specialize methods to get a Razor Uri from a generated document Uri, etc.

Author: davidwengier

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/LspEditMappingService.cs

@@ -15,11 +15,12 @@
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer;
 
-internal class LspEditMappingService(
+internal sealed class LspEditMappingService(
     IDocumentMappingService documentMappingService,
     IFilePathService filePathService,
     IDocumentContextFactory documentContextFactory) : AbstractEditMappingService(documentMappingService, filePathService)
 {
+    private readonly IFilePathService _filePathService = filePathService;
     private readonly IDocumentContextFactory _documentContextFactory = documentContextFactory;
 
     protected override bool TryGetDocumentContext(IDocumentSnapshot contextDocumentSnapshot, Uri razorDocumentUri, VSProjectContext? projectContext, [NotNullWhen(true)] out DocumentContext? documentContext)
@@ -31,4 +32,9 @@ protected override bool TryGetDocumentContext(IDocumentSnapshot contextDocumentS
 
         return true;
     }
+
+    protected override Task<Uri?> GetRazorDocumentUriAsync(IDocumentSnapshot contextDocumentSnapshot, Uri uri, CancellationToken cancellationToken)
+    {
+        return Task.FromResult<Uri?>(_filePathService.GetRazorDocumentUri(uri));
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/DocumentMapping/AbstractEditMappingService.cs

@@ -53,29 +53,34 @@ private async Task<Dictionary<string, TextEdit[]>> RemapDocumentEditsAsync(IDocu
 
         foreach (var (uriString, edits) in changes)
         {
-            var uri = new Uri(uriString);
+            var generatedDocumentUri = new Uri(uriString);
 
             // Check if the edit is actually for a generated document, because if not we don't need to do anything
-            if (!_filePathService.IsVirtualDocumentUri(uri))
+            if (!_filePathService.IsVirtualDocumentUri(generatedDocumentUri))
             {
                 remappedChanges[uriString] = edits;
                 continue;
             }
 
-            if (!TryGetDocumentContext(contextDocumentSnapshot, uri, projectContext: null, out var documentContext))
+            var razorDocumentUri = await GetRazorDocumentUriAsync(contextDocumentSnapshot, generatedDocumentUri, cancellationToken).ConfigureAwait(false);
+            if (razorDocumentUri is null)
+            {
+                continue;
+            }
+
+            if (!TryGetDocumentContext(contextDocumentSnapshot, razorDocumentUri, projectContext: null, out var documentContext))
             {
                 continue;
             }
 
             var codeDocument = await documentContext.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
-            var remappedEdits = RemapTextEditsCore(uri, codeDocument, edits);
+            var remappedEdits = RemapTextEditsCore(generatedDocumentUri, codeDocument, edits);
             if (remappedEdits.Length == 0)
             {
                 // Nothing to do.
                 continue;
             }
 
-            var razorDocumentUri = _filePathService.GetRazorDocumentUri(uri);
             remappedChanges[razorDocumentUri.AbsoluteUri] = remappedEdits;
         }
 
@@ -123,7 +128,11 @@ private async Task<TextDocumentEdit[]> RemapTextDocumentEditsAsync(IDocumentSnap
                 continue;
             }
 
-            var razorDocumentUri = _filePathService.GetRazorDocumentUri(generatedDocumentUri);
+            var razorDocumentUri = await GetRazorDocumentUriAsync(contextDocumentSnapshot, generatedDocumentUri, cancellationToken).ConfigureAwait(false);
+            if (razorDocumentUri is null)
+            {
+                continue;
+            }
 
             if (!TryGetDocumentContext(contextDocumentSnapshot, razorDocumentUri, entry.TextDocument.GetProjectContext(), out var documentContext))
             {
@@ -154,4 +163,6 @@ private async Task<TextDocumentEdit[]> RemapTextDocumentEditsAsync(IDocumentSnap
     }
 
     protected abstract bool TryGetDocumentContext(IDocumentSnapshot contextDocumentSnapshot, Uri razorDocumentUri, VSProjectContext? projectContext, [NotNullWhen(true)] out DocumentContext? documentContext);
+
+    protected abstract Task<Uri?> GetRazorDocumentUriAsync(IDocumentSnapshot contextDocumentSnapshot, Uri uri, CancellationToken cancellationToken);
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentMapping/RemoteDocumentMappingService.cs

@@ -2,12 +2,9 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System;
-using System.Collections.Generic;
 using System.Composition;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.AspNetCore.Razor;
-using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
@@ -55,11 +52,8 @@ internal sealed class RemoteDocumentMappingService(
 
         if (TryMapToHostDocumentRange(razorCodeDocument.GetRequiredCSharpDocument(), generatedDocumentRange, MappingBehavior.Strict, out var mappedRange))
         {
-            var solution = project.Solution;
-            var filePath = razorCodeDocument.Source.FilePath;
-            var documentId = solution.GetDocumentIdsWithFilePath(filePath).First();
-            var document = solution.GetAdditionalDocument(documentId).AssumeNotNull();
-            return (document.CreateUri(), mappedRange);
+            var razorDocumentUri = project.Solution.GetRazorDocumentUri(razorCodeDocument);
+            return (razorDocumentUri, mappedRange);
         }
 
         return (generatedDocumentUri, generatedDocumentRange);

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentMapping/RemoteEditMappingService.cs

@@ -4,6 +4,8 @@
 using System;
 using System.Composition;
 using System.Diagnostics.CodeAnalysis;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
@@ -39,4 +41,21 @@ protected override bool TryGetDocumentContext(IDocumentSnapshot contextDocumentS
         documentContext = new RemoteDocumentContext(razorDocumentUri, razorDocumentSnapshot);
         return true;
     }
+
+    protected override async Task<Uri?> GetRazorDocumentUriAsync(IDocumentSnapshot contextDocumentSnapshot, Uri generatedDocumentUri, CancellationToken cancellationToken)
+    {
+        if (contextDocumentSnapshot is not RemoteDocumentSnapshot originSnapshot)
+        {
+            throw new InvalidOperationException("RemoteEditMappingService can only be used with RemoteDocumentSnapshot instances.");
+        }
+
+        var project = originSnapshot.TextDocument.Project;
+        var razorCodeDocument = await _snapshotManager.GetSnapshot(project).TryGetCodeDocumentFromGeneratedDocumentUriAsync(generatedDocumentUri, cancellationToken).ConfigureAwait(false);
+        if (razorCodeDocument is null)
+        {
+            return null;
+        }
+
+        return project.Solution.GetRazorDocumentUri(razorCodeDocument);
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/ProjectSystem/Extensions.cs

@@ -1,8 +1,11 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
+using System;
 using System.Linq;
 using Microsoft.AspNetCore.Razor;
+using Microsoft.AspNetCore.Razor.Language;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 
@@ -26,4 +29,12 @@ document.FilePath is string filePath &&
 
     public static bool ContainsRazorDocuments(this Project project)
         => project.AdditionalDocuments.Any(static d => d.IsRazorDocument());
+
+    public static Uri GetRazorDocumentUri(this Solution solution, RazorCodeDocument codeDocument)
+    {
+        var filePath = codeDocument.Source.FilePath;
+        var documentId = solution.GetDocumentIdsWithFilePath(filePath).First();
+        var document = solution.GetAdditionalDocument(documentId).AssumeNotNull();
+        return document.CreateUri();
+    }
 }