PR Feedback, and simplify

Author: davidwengier

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/WrapWithTag/RemoteSpanMappingService.cs

@@ -38,20 +38,7 @@ public ValueTask<ImmutableArray<RazorMappedSpanResult>> MapSpansAsync(RazorPinne
 
     private async ValueTask<ImmutableArray<RazorMappedSpanResult>> MapSpansAsync(Solution solution, DocumentId generatedDocumentId, ImmutableArray<TextSpan> spans, CancellationToken cancellationToken)
     {
-        var generatedDocument = await solution.GetSourceGeneratedDocumentAsync(generatedDocumentId, cancellationToken).ConfigureAwait(false);
-        if (generatedDocument is null)
-        {
-            return [];
-        }
-
-        if (!generatedDocument.Project.TryGetHintNameFromGeneratedDocumentUri(generatedDocument.CreateUri(), out var hintName))
-        {
-            return [];
-        }
-
-        var projectSnapshot = _snapshotManager.GetSnapshot(generatedDocument.Project);
-
-        var razorDocument = await projectSnapshot.TryGetRazorDocumentFromGeneratedHintNameAsync(hintName, cancellationToken).ConfigureAwait(false);
+        var razorDocument = await TryGetRazorDocumentForGeneratedDocumentIdAsync(generatedDocumentId, solution, cancellationToken).ConfigureAwait(false);
         if (razorDocument is null)
         {
             return [];
@@ -92,20 +79,7 @@ private async ValueTask<ImmutableArray<RazorMappedEditResult>> MapTextChangesAsy
     {
         try
         {
-            var generatedDocument = await solution.GetSourceGeneratedDocumentAsync(generatedDocumentId, cancellationToken).ConfigureAwait(false);
-            if (generatedDocument is null)
-            {
-                return [];
-            }
-
-            if (!generatedDocument.Project.TryGetHintNameFromGeneratedDocumentUri(generatedDocument.CreateUri(), out var hintName))
-            {
-                return [];
-            }
-
-            var projectSnapshot = _snapshotManager.GetSnapshot(generatedDocument.Project);
-
-            var razorDocument = await projectSnapshot.TryGetRazorDocumentFromGeneratedHintNameAsync(hintName, cancellationToken).ConfigureAwait(false);
+            var razorDocument = await TryGetRazorDocumentForGeneratedDocumentIdAsync(generatedDocumentId, solution, cancellationToken).ConfigureAwait(false);
             if (razorDocument is null)
             {
                 return [];
@@ -135,4 +109,23 @@ private async ValueTask<ImmutableArray<RazorMappedEditResult>> MapTextChangesAsy
             return [];
         }
     }
+
+    private async Task<TextDocument?> TryGetRazorDocumentForGeneratedDocumentIdAsync(DocumentId generatedDocumentId, Solution solution, CancellationToken cancellationToken)
+    {
+        var generatedDocument = await solution.GetSourceGeneratedDocumentAsync(generatedDocumentId, cancellationToken).ConfigureAwait(false);
+        if (generatedDocument is null)
+        {
+            return null;
+        }
+
+        var identity = RazorGeneratedDocumentIdentity.Create(generatedDocument);
+        if (!identity.IsRazorSourceGeneratedDocument())
+        {
+            return null;
+        }
+
+        var projectSnapshot = _snapshotManager.GetSnapshot(generatedDocument.Project);
+
+        return await projectSnapshot.TryGetRazorDocumentFromGeneratedHintNameAsync(identity.HintName, cancellationToken).ConfigureAwait(false);
+    }
 }