Update logic for callers to the mapping service

Author: davidwengier

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/DocumentHighlighting/DocumentHighlightEndpoint.cs

@@ -73,14 +73,16 @@ public void ApplyCapabilities(VSInternalServerCapabilities serverCapabilities, V
             return null;
         }
 
-        var codeDocument = await documentContext.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
-
-        var generatedDocument = codeDocument.GetGeneratedDocument(positionInfo.LanguageKind);
-        foreach (var highlight in response)
+        if (positionInfo.LanguageKind is RazorLanguageKind.CSharp)
         {
-            if (_documentMappingService.TryMapToHostDocumentRange(generatedDocument, highlight.Range, out var mappedRange))
+            var codeDocument = await documentContext.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
+            var csharpDocument = codeDocument.GetRequiredCSharpDocument();
+            foreach (var highlight in response)
             {
-                highlight.Range = mappedRange;
+                if (_documentMappingService.TryMapToHostDocumentRange(csharpDocument, highlight.Range, out var mappedRange))
+                {
+                    highlight.Range = mappedRange;
+                }
             }
         }
 

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

@@ -47,13 +47,7 @@ internal sealed class LspDocumentMappingService(
         }
 
         var codeDocument = await documentContext.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
-
-        if (!codeDocument.TryGetGeneratedDocument(generatedDocumentUri, _filePathService, out var generatedDocument))
-        {
-            return Assumed.Unreachable<(Uri, LinePositionSpan)>();
-        }
-
-        if (TryMapToHostDocumentRange(generatedDocument, generatedDocumentRange, MappingBehavior.Strict, out var mappedRange))
+        if (TryMapToHostDocumentRange(codeDocument.GetRequiredCSharpDocument(), generatedDocumentRange, MappingBehavior.Strict, out var mappedRange))
         {
             return (razorDocumentUri, mappedRange);
         }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Mapping/RazorMapToDocumentRangesEndpoint.cs

@@ -14,16 +14,10 @@
 namespace Microsoft.AspNetCore.Razor.LanguageServer.Mapping;
 
 [RazorLanguageServerEndpoint(LanguageServerConstants.RazorMapToDocumentRangesEndpoint)]
-internal sealed class RazorMapToDocumentRangesEndpoint :
-    IRazorDocumentlessRequestHandler<RazorMapToDocumentRangesParams, RazorMapToDocumentRangesResponse?>,
-    ITextDocumentIdentifierHandler<RazorMapToDocumentRangesParams, Uri>
+internal sealed class RazorMapToDocumentRangesEndpoint(IDocumentMappingService documentMappingService)
+    : IRazorDocumentlessRequestHandler<RazorMapToDocumentRangesParams, RazorMapToDocumentRangesResponse?>, ITextDocumentIdentifierHandler<RazorMapToDocumentRangesParams, Uri>
 {
-    private readonly IDocumentMappingService _documentMappingService;
-
-    public RazorMapToDocumentRangesEndpoint(IDocumentMappingService documentMappingService)
-    {
-        _documentMappingService = documentMappingService;
-    }
+    private readonly IDocumentMappingService _documentMappingService = documentMappingService;
 
     public bool MutatesSolutionState { get; } = false;
 
@@ -46,29 +40,23 @@ public Uri GetTextDocumentIdentifier(RazorMapToDocumentRangesParams request)
         }
 
         var codeDocument = await documentContext.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
-        IRazorGeneratedDocument? generatedDocument = request.Kind switch
-        {
-            RazorLanguageKind.CSharp => codeDocument.GetRequiredCSharpDocument(),
-            RazorLanguageKind.Html => codeDocument.GetHtmlDocument(),
-            _ => throw new NotSupportedException($"Unsupported language kind '{request.Kind}'."),
-        };
-
-        if (generatedDocument is null)
-        {
-            return null;
-        }
+        var csharpDocument = codeDocument.GetRequiredCSharpDocument();
 
         var ranges = new LspRange[request.ProjectedRanges.Length];
         var spans = new RazorTextSpan[request.ProjectedRanges.Length];
 
         for (var i = 0; i < request.ProjectedRanges.Length; i++)
         {
-            var projectedRange = request.ProjectedRanges[i];
-            if (!_documentMappingService.TryMapToHostDocumentRange(generatedDocument, projectedRange, request.MappingBehavior, out var originalRange))
+            var originalRange = request.ProjectedRanges[i];
+            if (request.Kind is RazorLanguageKind.CSharp)
             {
-                // All language queries on unsupported documents return Html. This is equivalent to what pre-VSCode Razor was capable of.
-                ranges[i] = LspFactory.UndefinedRange;
-                continue;
+                var projectedRange = request.ProjectedRanges[i];
+                if (!_documentMappingService.TryMapToHostDocumentRange(csharpDocument, projectedRange, request.MappingBehavior, out originalRange))
+                {
+                    // All language queries on unsupported documents return Html. This is equivalent to what pre-VSCode Razor was capable of.
+                    ranges[i] = LspFactory.UndefinedRange;
+                    continue;
+                }
             }
 
             ranges[i] = originalRange;