Cohost document highlighting

eng/targets/Services.props

@@ -23,5 +23,6 @@
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.UriPresentation" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteUriPresentationService+Factory" />
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.FoldingRange" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteFoldingRangeService+Factory" />
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.SignatureHelp" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteSignatureHelpService+Factory" />
+    <ServiceHubService Include="Microsoft.VisualStudio.Razor.DocumentHighlight" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteDocumentHighlightService+Factory" />
   </ItemGroup>
 </Project>

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

@@ -178,14 +178,14 @@ static void AddHandlers(IServiceCollection services, LanguageServerFeatureOption
             services.AddTransient<IOnInitialized>(sp => sp.GetRequiredService<RazorConfigurationEndpoint>());
 
             services.AddHandlerWithCapabilities<ImplementationEndpoint>();
-            services.AddHandlerWithCapabilities<DocumentHighlightEndpoint>();
             services.AddHandlerWithCapabilities<OnAutoInsertEndpoint>();
 
             services.AddHandlerWithCapabilities<RenameEndpoint>();
             services.AddHandlerWithCapabilities<DefinitionEndpoint>();
 
             if (!featureOptions.UseRazorCohostServer)
             {
+                services.AddHandlerWithCapabilities<DocumentHighlightEndpoint>();
                 services.AddHandlerWithCapabilities<SignatureHelpEndpoint>();
                 services.AddHandlerWithCapabilities<LinkedEditingRangeEndpoint>();
                 services.AddHandlerWithCapabilities<FoldingRangeEndpoint>();

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

@@ -4,6 +4,7 @@
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
+using RLSP = Roslyn.LanguageServer.Protocol;
 
 namespace Microsoft.CodeAnalysis.Razor.Workspaces;
 
@@ -12,6 +13,9 @@ internal static class LinePositionExtensions
     public static Position ToPosition(this LinePosition linePosition)
         => new Position(linePosition.Line, linePosition.Character);
 
+    public static RLSP.Position ToRLSPPosition(this LinePosition linePosition)
+        => new RLSP.Position(linePosition.Line, linePosition.Character);
+
     public static bool TryGetAbsoluteIndex(this LinePosition position, SourceText sourceText, ILogger logger, out int absoluteIndex)
         => sourceText.TryGetAbsoluteIndex(position.Line, position.Character, logger, out absoluteIndex);
 

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

@@ -3,6 +3,7 @@
 
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
+using RLSP = Roslyn.LanguageServer.Protocol;
 
 namespace Microsoft.CodeAnalysis.Razor.Workspaces;
 
@@ -15,6 +16,13 @@ public static Range ToRange(this LinePositionSpan linePositionSpan)
             End = linePositionSpan.End.ToPosition()
         };
 
+    public static RLSP.Range ToRLSPRange(this LinePositionSpan linePositionSpan)
+        => new RLSP.Range
+        {
+            Start = linePositionSpan.Start.ToRLSPPosition(),
+            End = linePositionSpan.End.ToRLSPPosition()
+        };
+
     public static TextSpan ToTextSpan(this LinePositionSpan linePositionSpan, SourceText sourceText)
         => sourceText.GetTextSpan(linePositionSpan.Start.Line, linePositionSpan.Start.Character, linePositionSpan.End.Line, linePositionSpan.End.Character);
 

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

@@ -7,6 +7,7 @@
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
+using RLSP = Roslyn.LanguageServer.Protocol;
 
 namespace Microsoft.CodeAnalysis.Razor.Workspaces;
 
@@ -15,6 +16,9 @@ internal static class PositionExtensions
     public static LinePosition ToLinePosition(this Position position)
         => new LinePosition(position.Line, position.Character);
 
+    public static LinePosition ToLinePosition(this RLSP.Position position)
+        => new LinePosition(position.Line, position.Character);
+
     public static bool TryGetAbsoluteIndex(this Position position, SourceText sourceText, ILogger? logger, out int absoluteIndex)
     {
         if (position is null)

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

@@ -5,6 +5,7 @@
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 using Range = Microsoft.VisualStudio.LanguageServer.Protocol.Range;
+using RLSP = Roslyn.LanguageServer.Protocol;
 
 namespace Microsoft.CodeAnalysis.Razor.Workspaces;
 
@@ -110,6 +111,9 @@ public static TextSpan ToTextSpan(this Range range, SourceText sourceText)
     public static LinePositionSpan ToLinePositionSpan(this Range range)
         => new LinePositionSpan(range.Start.ToLinePosition(), range.End.ToLinePosition());
 
+    public static LinePositionSpan ToLinePositionSpan(this RLSP.Range range)
+        => new LinePositionSpan(range.Start.ToLinePosition(), range.End.ToLinePosition());
+
     public static bool IsUndefined(this Range range)
     {
         if (range is null)

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Protocol/DocumentHighlight/RemoteDocumentHighlight.cs

@@ -0,0 +1,28 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System.Runtime.Serialization;
+using Microsoft.CodeAnalysis.Razor.Workspaces;
+using Microsoft.CodeAnalysis.Text;
+using Microsoft.VisualStudio.LanguageServer.Protocol;
+using RLSP = Roslyn.LanguageServer.Protocol;
+
+namespace Microsoft.CodeAnalysis.Razor.Protocol.DocumentHighlight;
+
+using DocumentHighlight = VisualStudio.LanguageServer.Protocol.DocumentHighlight;
+
+[DataContract]
+internal readonly record struct RemoteDocumentHighlight(
+    [property: DataMember(Order = 0)] LinePositionSpan Position,
+    [property: DataMember(Order = 1)] DocumentHighlightKind Kind)
+{
+    public static RemoteDocumentHighlight FromRLSPDocumentHighlight(RLSP.DocumentHighlight h)
+        => new RemoteDocumentHighlight(h.Range.ToLinePositionSpan(), (DocumentHighlightKind)h.Kind);
+
+    public static DocumentHighlight ToLspDocumentHighlight(RemoteDocumentHighlight r)
+        => new DocumentHighlight
+        {
+            Range = r.Position.ToRange(),
+            Kind = r.Kind
+        };
+}

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteDocumentHighlightService.cs

@@ -0,0 +1,15 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
+using Microsoft.CodeAnalysis.Razor.Protocol.DocumentHighlight;
+using Microsoft.CodeAnalysis.Text;
+
+namespace Microsoft.CodeAnalysis.Razor.Remote;
+
+internal interface IRemoteDocumentHighlightService
+{
+    ValueTask<RemoteResponse<RemoteDocumentHighlight[]?>> GetHighlightsAsync(RazorPinnedSolutionInfoWrapper solutionInfo, DocumentId documentId, LinePosition position, CancellationToken cancellationToken);
+}

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteUriPresentationService.cs

@@ -2,7 +2,6 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System;
-using System.Runtime.Serialization;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
@@ -12,10 +11,6 @@ namespace Microsoft.CodeAnalysis.Razor.Remote;
 
 internal interface IRemoteUriPresentationService
 {
-    ValueTask<Response> GetPresentationAsync(RazorPinnedSolutionInfoWrapper solutionInfo, DocumentId razorDocumentId, LinePositionSpan span, Uri[]? uris, CancellationToken cancellationToken);
+    ValueTask<RemoteResponse<TextChange?>> GetPresentationAsync(RazorPinnedSolutionInfoWrapper solutionInfo, DocumentId razorDocumentId, LinePositionSpan span, Uri[]? uris, CancellationToken cancellationToken);
 
-    [DataContract]
-    public record struct Response(
-        [property: DataMember(Order = 0)] bool CallHtml,
-        [property: DataMember(Order = 1)] TextChange? TextChange);
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/RazorServices.cs

@@ -1,37 +1,51 @@
 // 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.Collections.Generic;
 using Microsoft.AspNetCore.Razor.Serialization.MessagePack.Resolvers;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 
 namespace Microsoft.CodeAnalysis.Razor.Remote;
 
 internal static class RazorServices
 {
-    private const string ComponentName = "Razor";
-
-    public static readonly RazorServiceDescriptorsWrapper Descriptors = new(
-        ComponentName,
-        featureDisplayNameProvider: feature => $"Razor {feature} Feature",
-        additionalFormatters: [],
-        additionalResolvers: TopLevelResolvers.All,
-        interfaces:
+    // Internal for testing
+    internal static readonly IEnumerable<(Type, Type?)> MessagePackServices =
         [
             (typeof(IRemoteLinkedEditingRangeService), null),
             (typeof(IRemoteTagHelperProviderService), null),
             (typeof(IRemoteClientInitializationService), null),
             (typeof(IRemoteSemanticTokensService), null),
             (typeof(IRemoteHtmlDocumentService), null),
             (typeof(IRemoteUriPresentationService), null),
-            (typeof(IRemoteFoldingRangeService), null)
-        ]);
+            (typeof(IRemoteFoldingRangeService), null),
+            (typeof(IRemoteDocumentHighlightService), null),
+        ];
+
+    // Internal for testing
+    internal static readonly IEnumerable<(Type, Type?)> JsonServices =
+        [
+            (typeof(IRemoteSignatureHelpService), null),
+        ];
+
+    private const string ComponentName = "Razor";
+
+    public static readonly RazorServiceDescriptorsWrapper Descriptors = new(
+        ComponentName,
+        featureDisplayNameProvider: GetFeatureDisplayName,
+        additionalFormatters: [],
+        additionalResolvers: TopLevelResolvers.All,
+        interfaces: MessagePackServices);
 
     public static readonly RazorServiceDescriptorsWrapper JsonDescriptors = new(
         ComponentName, // Needs to match the above because so much of our ServiceHub infrastructure is convention based
-        featureDisplayNameProvider: feature => $"Razor {feature} Feature",
+        featureDisplayNameProvider: GetFeatureDisplayName,
         jsonConverters: RazorServiceDescriptorsWrapper.GetLspConverters(),
-        interfaces:
-        [
-            (typeof(IRemoteSignatureHelpService), null),
-        ]);
+        interfaces: JsonServices);
+
+    private static string GetFeatureDisplayName(string feature)
+    {
+        return $"Razor {feature} Feature";
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/RemoteResponse.cs

@@ -0,0 +1,16 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System.Runtime.Serialization;
+
+namespace Microsoft.CodeAnalysis.Razor.Remote;
+
+[DataContract]
+internal record struct RemoteResponse<T>(
+    [property: DataMember(Order = 0)] bool StopHandling,
+    [property: DataMember(Order = 1)] T Result)
+{
+    public static RemoteResponse<T> CallHtml => new(StopHandling: false, Result: default!);
+    public static RemoteResponse<T> NoFurtherHandling => new(StopHandling: true, Result: default!);
+    public static RemoteResponse<T> Results(T result) => new(StopHandling: false, Result: result);
+}

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentHighlight/RemoteDocumentHighlightService.cs

@@ -0,0 +1,92 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor.Language;
+using Microsoft.AspNetCore.Razor.PooledObjects;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost.Handlers;
+using Microsoft.CodeAnalysis.Razor.DocumentMapping;
+using Microsoft.CodeAnalysis.Razor.Protocol;
+using Microsoft.CodeAnalysis.Razor.Protocol.DocumentHighlight;
+using Microsoft.CodeAnalysis.Razor.Remote;
+using Microsoft.CodeAnalysis.Razor.Workspaces;
+using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
+using Microsoft.CodeAnalysis.Text;
+using Response = Microsoft.CodeAnalysis.Razor.Remote.RemoteResponse<Microsoft.CodeAnalysis.Razor.Protocol.DocumentHighlight.RemoteDocumentHighlight[]?>;
+
+namespace Microsoft.CodeAnalysis.Remote.Razor;
+
+internal sealed partial class RemoteDocumentHighlightService(in ServiceArgs args) : RazorDocumentServiceBase(in args), IRemoteDocumentHighlightService
+{
+    internal sealed class Factory : FactoryBase<IRemoteDocumentHighlightService>
+    {
+        protected override IRemoteDocumentHighlightService CreateService(in ServiceArgs args)
+            => new RemoteDocumentHighlightService(in args);
+    }
+
+    private readonly IRazorDocumentMappingService _documentMappingService = args.ExportProvider.GetExportedValue<IRazorDocumentMappingService>();
+    private readonly IFilePathService _filePathService = args.ExportProvider.GetExportedValue<IFilePathService>();
+
+    public ValueTask<Response> GetHighlightsAsync(
+        RazorPinnedSolutionInfoWrapper solutionInfo,
+        DocumentId razorDocumentId,
+        LinePosition position,
+        CancellationToken cancellationToken)
+        => RunServiceAsync(
+            solutionInfo,
+            razorDocumentId,
+            context => GetHighlightsAsync(context, position, cancellationToken),
+            cancellationToken);
+
+    private async ValueTask<Response> GetHighlightsAsync(
+        RemoteDocumentContext context,
+        LinePosition position,
+        CancellationToken cancellationToken)
+    {
+        var sourceText = await context.GetSourceTextAsync(cancellationToken).ConfigureAwait(false);
+        if (!sourceText.TryGetAbsoluteIndex(position.Line, position.Character, out var index))
+        {
+            return Response.NoFurtherHandling;
+        }
+
+        var codeDocument = await context.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
+
+        var languageKind = _documentMappingService.GetLanguageKind(codeDocument, index, rightAssociative: true);
+        if (languageKind is RazorLanguageKind.Html)
+        {
+            return Response.CallHtml;
+        }
+        else if (languageKind is RazorLanguageKind.Razor)
+        {
+            return Response.NoFurtherHandling;
+        }
+
+        var csharpDocument = codeDocument.GetCSharpDocument();
+        if (_documentMappingService.TryMapToGeneratedDocumentPosition(csharpDocument, index, out var mappedPosition, out _))
+        {
+            var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+
+            var highlights = await DocumentHighlights.GetHighlightsAsync(generatedDocument, mappedPosition, cancellationToken).ConfigureAwait(false);
+
+            if (highlights is not null)
+            {
+                using var results = new PooledArrayBuilder<RemoteDocumentHighlight>();
+
+                foreach (var highlight in highlights)
+                {
+                    if (_documentMappingService.TryMapToHostDocumentRange(csharpDocument, highlight.Range.ToLinePositionSpan(), out var mappedRange))
+                    {
+                        highlight.Range = mappedRange.ToRLSPRange();
+                        results.Add(RemoteDocumentHighlight.FromRLSPDocumentHighlight(highlight));
+                    }
+                }
+
+                return Response.Results(results.ToArray());
+            }
+        }
+
+        return Response.NoFurtherHandling;
+    }
+}