Update document mapping to work with the source generator

src/Compiler/Microsoft.CodeAnalysis.Razor.Compiler/src/SourceGenerators/RazorGeneratorResult.cs

@@ -7,12 +7,14 @@
 
 namespace Microsoft.NET.Sdk.Razor.SourceGenerators
 {
-    internal sealed class RazorGeneratorResult(IReadOnlyList<TagHelperDescriptor> tagHelpers, ImmutableDictionary<string, (string hintName, RazorCodeDocument document)> documents)
+    internal sealed class RazorGeneratorResult(IReadOnlyList<TagHelperDescriptor> tagHelpers, ImmutableDictionary<string, (string hintName, RazorCodeDocument document)> filePathToDocument, ImmutableDictionary<string, string> hintNameToFilePath)
     {
         public IReadOnlyList<TagHelperDescriptor> TagHelpers => tagHelpers;
 
-        public RazorCodeDocument? GetCodeDocument(string physicalPath) => documents.TryGetValue(physicalPath, out var pair) ? pair.document : null;
+        public RazorCodeDocument? GetCodeDocument(string physicalPath) => filePathToDocument.TryGetValue(physicalPath, out var pair) ? pair.document : null;
 
-        public string? GetHintName(string physicalPath) => documents.TryGetValue(physicalPath, out var pair) ? pair.hintName : null;
+        public string? GetHintName(string physicalPath) => filePathToDocument.TryGetValue(physicalPath, out var pair) ? pair.hintName : null;
+
+        public string? GetFilePath(string hintName) => hintNameToFilePath.TryGetValue(hintName, out var filePath) ? filePath : null;
     }
 }

src/Compiler/Microsoft.CodeAnalysis.Razor.Compiler/src/SourceGenerators/RazorSourceGenerator.cs

@@ -3,10 +3,11 @@
 
 using System;
 using System.Collections.Generic;
-using System.Collections.Immutable;
 using System.IO;
 using System.Linq;
+using Microsoft.AspNetCore.Razor;
 using Microsoft.AspNetCore.Razor.Language;
+using Microsoft.AspNetCore.Razor.PooledObjects;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp;
 
@@ -348,8 +349,16 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
 
                 if (!isGeneratorSuppressed)
                 {
-                    var documentDictionary = documents.Select(p => KeyValuePair.Create(p.codeDocument.Source.FilePath!, (p.hintName, p.codeDocument))).ToImmutableDictionary();
-                    context.AddOutput(nameof(RazorGeneratorResult), new RazorGeneratorResult(tagHelpers, documentDictionary));
+                    using var filePathToDocument = new PooledDictionaryBuilder<string, (string, RazorCodeDocument)>();
+                    using var hintNameToFilePath = new PooledDictionaryBuilder<string, string>();
+
+                    foreach (var (hintName, codeDocument, _) in documents)
+                    {
+                        filePathToDocument.Add(codeDocument.Source.FilePath!, (hintName, codeDocument));
+                        hintNameToFilePath.Add(hintName, codeDocument.Source.FilePath!);
+                    }
+
+                    context.AddOutput(nameof(RazorGeneratorResult), new RazorGeneratorResult(tagHelpers, filePathToDocument.ToImmutable(), hintNameToFilePath.ToImmutable()));
                 }
             });
         }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Extensions/ProjectExtensions.cs

@@ -7,6 +7,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 
 namespace Microsoft.CodeAnalysis;
 
@@ -40,7 +41,7 @@ public static bool TryGetCSharpDocument(this Project project, Uri csharpDocument
     /// <summary>
     /// Finds source generated documents by iterating through all of them. In OOP there are better options!
     /// </summary>
-    public static async Task<Document?> TryGetSourceGeneratedDocumentFromHintNameAsync(this Project project, string? hintName, CancellationToken cancellationToken)
+    public static async Task<SourceGeneratedDocument?> TryGetSourceGeneratedDocumentFromHintNameAsync(this Project project, string? hintName, CancellationToken cancellationToken)
     {
         // TODO: use this when the location is case-insensitive on windows (https://github.com/dotnet/roslyn/issues/76869)
         //var generator = typeof(RazorSourceGenerator);
@@ -51,4 +52,16 @@ public static bool TryGetCSharpDocument(this Project project, Uri csharpDocument
         var generatedDocuments = await project.GetSourceGeneratedDocumentsAsync(cancellationToken).ConfigureAwait(false);
         return generatedDocuments.SingleOrDefault(d => d.HintName == hintName);
     }
+
+    /// <summary>
+    /// Finds source generated documents by iterating through all of them. In OOP there are better options!
+    /// </summary>
+    public static async Task<string?> TryGetHintNameFromGeneratedDocumentUriAsync(this Project project, Uri generatedDocumentUri, CancellationToken cancellationToken)
+    {
+        // TODO: Call SourceGeneratedDocumentUri.DeserializeIdentity: https://github.com/dotnet/razor/issues/11557
+
+        var generatedDocuments = await project.GetSourceGeneratedDocumentsAsync(cancellationToken).ConfigureAwait(false);
+        var document = generatedDocuments.SingleOrDefault(d => generatedDocumentUri.Equals(d.CreateUri()));
+        return document?.HintName;
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Completion/RemoteCompletionService.cs

@@ -171,7 +171,7 @@ private async ValueTask<Response> GetCompletionAsync(
             var generatedText = await generatedDocument.GetTextAsync(cancellationToken).ConfigureAwait(false);
             var change = generatedText.GetTextChange(provisionalTextEdit);
             generatedText = generatedText.WithChanges([change]);
-            generatedDocument = generatedDocument.WithText(generatedText);
+            generatedDocument = (SourceGeneratedDocument)generatedDocument.WithText(generatedText);
         }
 
         // This is, to say the least, not ideal. In future we're going to normalize on to Roslyn LSP types, and this can go.

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

@@ -2,11 +2,13 @@
 // 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.AspNetCore.Razor.Language;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
@@ -32,11 +34,10 @@ internal sealed class RemoteDocumentMappingService(
         LinePositionSpan generatedDocumentRange,
         CancellationToken cancellationToken)
     {
-        var razorDocumentUri = _filePathService.GetRazorDocumentUri(generatedDocumentUri);
-
         // For Html we just map the Uri, the range will be the same
         if (_filePathService.IsVirtualHtmlFile(generatedDocumentUri))
         {
+            var razorDocumentUri = _filePathService.GetRazorDocumentUri(generatedDocumentUri);
             return (razorDocumentUri, generatedDocumentRange);
         }
 
@@ -46,31 +47,20 @@ internal sealed class RemoteDocumentMappingService(
             return (generatedDocumentUri, generatedDocumentRange);
         }
 
-        var solution = originSnapshot.TextDocument.Project.Solution;
-        if (!solution.TryGetRazorDocument(razorDocumentUri, out var razorDocument))
-        {
-            return (generatedDocumentUri, generatedDocumentRange);
-        }
-
-        var razorDocumentSnapshot = _snapshotManager.GetSnapshot(razorDocument);
-
-        var razorCodeDocument = await razorDocumentSnapshot
-            .GetGeneratedOutputAsync(cancellationToken)
-            .ConfigureAwait(false);
-
+        var project = originSnapshot.TextDocument.Project;
+        var razorCodeDocument = await _snapshotManager.GetSnapshot(project).TryGetCodeDocumentFromGeneratedDocumentUriAsync(generatedDocumentUri, cancellationToken).ConfigureAwait(false);
         if (razorCodeDocument is null)
         {
             return (generatedDocumentUri, generatedDocumentRange);
         }
 
-        if (!razorCodeDocument.TryGetGeneratedDocument(generatedDocumentUri, _filePathService, out var generatedDocument))
-        {
-            return Assumed.Unreachable<(Uri, LinePositionSpan)>();
-        }
-
-        if (TryMapToHostDocumentRange(generatedDocument, generatedDocumentRange, MappingBehavior.Strict, out var mappedRange))
+        if (TryMapToHostDocumentRange(razorCodeDocument.GetCSharpDocument(), generatedDocumentRange, MappingBehavior.Strict, out var mappedRange))
         {
-            return (razorDocumentUri, 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);
         }
 
         return (generatedDocumentUri, generatedDocumentRange);

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

@@ -22,7 +22,7 @@ internal sealed class RemoteDocumentSnapshot : IDocumentSnapshot
     public RemoteProjectSnapshot ProjectSnapshot { get; }
 
     private RazorCodeDocument? _codeDocument;
-    private Document? _generatedDocument;
+    private SourceGeneratedDocument? _generatedDocument;
 
     public RemoteDocumentSnapshot(TextDocument textDocument, RemoteProjectSnapshot projectSnapshot)
     {
@@ -103,7 +103,7 @@ public IDocumentSnapshot WithText(SourceText text)
         return snapshotManager.GetSnapshot(newDocument);
     }
 
-    public async ValueTask<Document> GetGeneratedDocumentAsync(CancellationToken cancellationToken)
+    public async ValueTask<SourceGeneratedDocument> GetGeneratedDocumentAsync(CancellationToken cancellationToken)
     {
         if (_generatedDocument is not null)
         {

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

@@ -169,7 +169,7 @@ public bool TryGetDocument(string filePath, [NotNullWhen(true)] out IDocumentSna
         return generatorResult.GetCodeDocument(documentSnapshot.FilePath);
     }
 
-    internal async Task<Document?> GetGeneratedDocumentAsync(IDocumentSnapshot documentSnapshot, CancellationToken cancellationToken)
+    internal async Task<SourceGeneratedDocument?> GetGeneratedDocumentAsync(IDocumentSnapshot documentSnapshot, CancellationToken cancellationToken)
     {
         var generatorResult = await GetRazorGeneratorResultAsync(cancellationToken).ConfigureAwait(false);
         if (generatorResult is null)
@@ -184,6 +184,29 @@ public bool TryGetDocument(string filePath, [NotNullWhen(true)] out IDocumentSna
         return generatedDocument ?? throw new InvalidOperationException("Couldn't get the source generated document for a hint name that we got from the generator?");
     }
 
+    public async Task<RazorCodeDocument?> TryGetCodeDocumentFromGeneratedDocumentUriAsync(Uri generatedDocumentUri, CancellationToken cancellationToken)
+    {
+        if (await _project.TryGetHintNameFromGeneratedDocumentUriAsync(generatedDocumentUri, cancellationToken).ConfigureAwait(false) is not { } hintName)
+        {
+            return null;
+        }
+
+        return await TryGetCodeDocumentFromGeneratedHintNameAsync(hintName, cancellationToken).ConfigureAwait(false);
+    }
+
+    public async Task<RazorCodeDocument?> TryGetCodeDocumentFromGeneratedHintNameAsync(string generatedDocumentHintName, CancellationToken cancellationToken)
+    {
+        var runResult = await GetRazorGeneratorResultAsync(cancellationToken).ConfigureAwait(false);
+        if (runResult is null)
+        {
+            return null;
+        }
+
+        return runResult.GetFilePath(generatedDocumentHintName) is { } razorFilePath
+            ? runResult.GetCodeDocument(razorFilePath)
+            : null;
+    }
+
     private async Task<RazorGeneratorResult?> GetRazorGeneratorResultAsync(CancellationToken cancellationToken)
     {
         var result = await _project.GetSourceGeneratorRunResultAsync(cancellationToken).ConfigureAwait(false);

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/RemoteFilePathService.cs

@@ -23,6 +23,6 @@ public override Uri GetRazorDocumentUri(Uri virtualDocumentUri)
 
     public override bool IsVirtualCSharpFile(Uri uri)
     {
-        return uri.Scheme == "source-generated";
+        return uri.Scheme == "roslyn-source-generated";
     }
 }

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Cohost/CodeActions/ExtractToCodeBehindTests.cs

@@ -1,7 +1,6 @@
 // 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.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Test.Common;
 using Xunit.Abstractions;
@@ -11,7 +10,7 @@ namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost.CodeActions;
 
 public class ExtractToCodeBehindTests(FuseTestContext context, ITestOutputHelper testOutputHelper) : CohostCodeActionsEndpointTestBase(context, testOutputHelper)
 {
-    [FuseFact(Skip = "Need to map uri back to source generated document")]
+    [FuseFact]
     public async Task ExtractToCodeBehind()
     {
         await VerifyCodeActionAsync(

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Cohost/CohostFindAllReferencesEndpointTest.cs

@@ -17,7 +17,7 @@ namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
 
 public class CohostFindAllReferencesEndpointTest(FuseTestContext context, ITestOutputHelper testOutputHelper) : CohostEndpointTestBase(testOutputHelper), IClassFixture<FuseTestContext>
 {
-    [FuseTheory(Skip = "IFilePathService does not yet map generated documents")]
+    [FuseTheory]
     [CombinatorialData]
     public Task FindCSharpMember(bool supportsVSExtensions)
         => VerifyFindAllReferencesAsync("""
@@ -36,7 +36,7 @@ string M()
             """,
             supportsVSExtensions);
 
-    [FuseTheory(Skip = "IFilePathService does not yet map generated documents")]
+    [FuseTheory]
     [CombinatorialData]
     public async Task ComponentAttribute(bool supportsVSExtensions)
     {
@@ -60,7 +60,7 @@ await VerifyFindAllReferencesAsync(input, supportsVSExtensions,
             (FilePath("SurveyPrompt.razor"), surveyPrompt));
     }
 
-    [FuseTheory(Skip = "IFilePathService does not yet map generated documents")]
+    [FuseTheory]
     [CombinatorialData]
     public async Task OtherCSharpFile(bool supportsVSExtensions)
     {

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Cohost/CohostGoToDefinitionEndpointTest.cs

@@ -21,7 +21,7 @@ namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
 
 public class CohostGoToDefinitionEndpointTest(FuseTestContext context, ITestOutputHelper testOutputHelper) : CohostEndpointTestBase(testOutputHelper), IClassFixture<FuseTestContext>
 {
-    [FuseFact(Skip = "IFilePathService does not yet map generated documents")]
+    [FuseFact]
     public async Task CSharp_Method()
     {
         var input = """
@@ -40,7 +40,7 @@ public async Task CSharp_Method()
         await VerifyGoToDefinitionAsync(input);
     }
 
-    [FuseFact(Skip = "IFilePathService does not yet map generated documents")]
+    [FuseFact]
     public async Task CSharp_Local()
     {
         var input = """
@@ -87,7 +87,7 @@ public async Task CSharp_MetadataReference()
         Assert.Contains("public sealed class String", line);
     }
 
-    [FuseTheory(Skip = "IFilePathService does not yet map generated documents")]
+    [FuseTheory]
     [InlineData("$$IncrementCount")]
     [InlineData("In$$crementCount")]
     [InlineData("IncrementCount$$")]
@@ -107,7 +107,7 @@ public async Task Attribute_SameFile(string method)
         await VerifyGoToDefinitionAsync(input, FileKinds.Component);
     }
 
-    [FuseFact(Skip = "IFilePathService does not yet map generated documents")]
+    [FuseFact]
     public async Task AttributeValue_BindAfter()
     {
         var input = """
@@ -157,7 +157,7 @@ public async Task Component()
         Assert.Equal(range, location.Range);
     }
 
-    [FuseTheory(Skip = "IFilePathService does not yet map generated documents")]
+    [FuseTheory]
     [InlineData("Ti$$tle")]
     [InlineData("$$@bind-Title")]
     [InlineData("@$$bind-Title")]

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Cohost/CohostGoToImplementationEndpointTest.cs

@@ -19,7 +19,7 @@ namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
 
 public class CohostGoToImplementationEndpointTest(FuseTestContext context, ITestOutputHelper testOutputHelper) : CohostEndpointTestBase(testOutputHelper), IClassFixture<FuseTestContext>
 {
-    [FuseFact(Skip = "IFilePathService does not yet map generated documents")]
+    [FuseFact]
     public async Task CSharp_Method()
     {
         var input = """
@@ -40,7 +40,7 @@ public async Task CSharp_Method()
         await VerifyCSharpGoToImplementationAsync(input);
     }
 
-    [FuseFact(Skip = "IFilePathService does not yet map generated documents")]
+    [FuseFact]
     public async Task CSharp_Field()
     {
         var input = """
@@ -64,7 +64,7 @@ string GetX()
         await VerifyCSharpGoToImplementationAsync(input);
     }
 
-    [FuseFact(Skip = "IFilePathService does not yet map generated documents")]
+    [FuseFact]
     public async Task CSharp_Multiple()
     {
         var input = """