Cohost inlay hint support

eng/targets/Services.props

@@ -24,5 +24,6 @@
     <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" />
+    <ServiceHubService Include="Microsoft.VisualStudio.Razor.InlayHint" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteInlayHintService+Factory" />
   </ItemGroup>
 </Project>

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/InlayHints/InlayHintEndpoint.cs

@@ -6,16 +6,14 @@
 using Microsoft.AspNetCore.Razor.LanguageServer.EndpointContracts;
 using Microsoft.AspNetCore.Razor.LanguageServer.Hosting;
 using Microsoft.AspNetCore.Razor.Threading;
-using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer.InlayHints;
 
 [RazorLanguageServerEndpoint(Methods.TextDocumentInlayHintName)]
-internal sealed class InlayHintEndpoint(LanguageServerFeatureOptions featureOptions, IInlayHintService inlayHintService, IClientConnection clientConnection)
+internal sealed class InlayHintEndpoint(IInlayHintService inlayHintService, IClientConnection clientConnection)
     : IRazorRequestHandler<InlayHintParams, InlayHint[]?>, ICapabilitiesProvider
 {
-    private readonly LanguageServerFeatureOptions _featureOptions = featureOptions;
     private readonly IInlayHintService _inlayHintService = inlayHintService;
     private readonly IClientConnection _clientConnection = clientConnection;
 

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

@@ -189,6 +189,10 @@ static void AddHandlers(IServiceCollection services, LanguageServerFeatureOption
                 services.AddHandlerWithCapabilities<SignatureHelpEndpoint>();
                 services.AddHandlerWithCapabilities<LinkedEditingRangeEndpoint>();
                 services.AddHandlerWithCapabilities<FoldingRangeEndpoint>();
+
+                services.AddSingleton<IInlayHintService, InlayHintService>();
+                services.AddHandlerWithCapabilities<InlayHintEndpoint>();
+                services.AddHandler<InlayHintResolveEndpoint>();
             }
 
             services.AddHandler<WrapWithTagEndpoint>();
@@ -204,11 +208,6 @@ static void AddHandlers(IServiceCollection services, LanguageServerFeatureOption
             services.AddHandlerWithCapabilities<ProjectContextsEndpoint>();
             services.AddHandlerWithCapabilities<DocumentSymbolEndpoint>();
             services.AddHandlerWithCapabilities<MapCodeEndpoint>();
-
-            services.AddSingleton<IInlayHintService, InlayHintService>();
-
-            services.AddHandlerWithCapabilities<InlayHintEndpoint>();
-            services.AddHandler<InlayHintResolveEndpoint>();
         }
     }
 }

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

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
@@ -13,7 +14,7 @@ namespace Microsoft.CodeAnalysis.Razor.DocumentMapping;
 
 internal interface IRazorDocumentMappingService
 {
-    TextEdit[] GetHostDocumentEdits(IRazorGeneratedDocument generatedDocument, TextEdit[] generatedDocumentEdits);
+    IEnumerable<TextChange> GetHostDocumentEdits(IRazorGeneratedDocument generatedDocument, IEnumerable<TextChange> generatedDocumentEdits);
 
     bool TryMapToHostDocumentRange(IRazorGeneratedDocument generatedDocument, LinePositionSpan generatedDocumentRange, MappingBehavior mappingBehavior, out LinePositionSpan hostDocumentRange);
 

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

@@ -3,8 +3,10 @@
 
 using System;
 using System.Diagnostics.CodeAnalysis;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.Protocol;
@@ -16,6 +18,16 @@ namespace Microsoft.CodeAnalysis.Razor.DocumentMapping;
 
 internal static class IRazorDocumentMappingServiceExtensions
 {
+    public static TextEdit[] GetHostDocumentEdits(this IRazorDocumentMappingService service, IRazorGeneratedDocument generatedDocument, TextEdit[] generatedDocumentEdits)
+    {
+        var generatedDocumentSourceText = generatedDocument.GetGeneratedSourceText();
+        var documentText = generatedDocument.CodeDocument.AssumeNotNull().Source.Text;
+
+        var changes = generatedDocumentEdits.Select(generatedDocumentSourceText.GetTextChange);
+        var mappedChanges = service.GetHostDocumentEdits(generatedDocument, changes);
+        return mappedChanges.Select(documentText.GetTextEdit).ToArray();
+    }
+
     public static bool TryMapToHostDocumentRange(this IRazorDocumentMappingService service, IRazorGeneratedDocument generatedDocument, LinePositionSpan projectedRange, out LinePositionSpan originalRange)
         => service.TryMapToHostDocumentRange(generatedDocument, projectedRange, MappingBehavior.Strict, out originalRange);
 

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

@@ -0,0 +1,17 @@
+using System;
+using Microsoft.CodeAnalysis.Text;
+
+namespace Microsoft.AspNetCore.Razor.Language;
+
+internal static class IRazorGeneratedDocumentExtensions
+{
+    public static SourceText GetGeneratedSourceText(this IRazorGeneratedDocument generatedDocument)
+    {
+        if (generatedDocument.CodeDocument is not { } codeDocument)
+        {
+            throw new InvalidOperationException("Cannot use document mapping service on a generated document that has a null CodeDocument.");
+        }
+
+        return codeDocument.GetGeneratedSourceText(generatedDocument);
+    }
+}

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

@@ -0,0 +1,23 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using Microsoft.AspNetCore.Razor;
+using Microsoft.CodeAnalysis.Text;
+
+namespace Roslyn.LanguageServer.Protocol;
+
+internal static partial class RoslynLspExtensions
+{
+    public static Range GetRange(this SourceText text, TextSpan span)
+        => text.GetLinePositionSpan(span).ToRange();
+
+    public static TextSpan GetTextSpan(this SourceText text, Range range)
+        => text.GetTextSpan(range.Start.Line, range.Start.Character, range.End.Line, range.End.Character);
+
+    public static TextChange GetTextChange(this SourceText text, TextEdit edit)
+        => new(text.GetTextSpan(edit.Range), edit.NewText);
+
+    public static TextEdit GetTextEdit(this SourceText text, TextChange change)
+        => RoslynLspFactory.CreateTextEdit(text.GetRange(change.Span), change.NewText.AssumeNotNull());
+
+}

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/Extensions/TextDocumentExtensions.cs → src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Extensions/RoslynLspExtensions_TextDocumentIdentifier.cs

@@ -3,11 +3,10 @@
 
 using System;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost;
-using Microsoft.VisualStudio.LanguageServer.Protocol;
 
-namespace Microsoft.VisualStudio.Razor.LanguageClient.Extensions;
+namespace Roslyn.LanguageServer.Protocol;
 
-internal static class TextDocumentExtensions
+internal static partial class RoslynLspExtensions
 {
     /// <summary>
     /// Returns a copy of the passed in <see cref="TextDocumentIdentifier"/> with the passed in <see cref="Uri"/>.

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

@@ -1,6 +1,9 @@
 // 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 Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost;
+
 namespace Microsoft.VisualStudio.LanguageServer.Protocol;
 
 internal static partial class VsLspExtensions
@@ -9,4 +12,27 @@ internal static partial class VsLspExtensions
         => textDocumentIdentifier is VSTextDocumentIdentifier vsIdentifier
             ? vsIdentifier.ProjectContext
             : null;
+
+    /// <summary>
+    /// Returns a copy of the passed in <see cref="TextDocumentIdentifier"/> with the passed in <see cref="Uri"/>.
+    /// </summary>
+    public static TextDocumentIdentifier WithUri(this TextDocumentIdentifier textDocumentIdentifier, Uri uri)
+    {
+        if (textDocumentIdentifier is VSTextDocumentIdentifier vsTdi)
+        {
+            return new VSTextDocumentIdentifier
+            {
+                Uri = uri,
+                ProjectContext = vsTdi.ProjectContext
+            };
+        }
+
+        return new TextDocumentIdentifier
+        {
+            Uri = uri
+        };
+    }
+
+    public static RazorTextDocumentIdentifier ToRazorTextDocumentIdentifier(this TextDocumentIdentifier textDocumentIdentifier)
+        => new RazorTextDocumentIdentifier(textDocumentIdentifier.Uri, (textDocumentIdentifier as VSTextDocumentIdentifier)?.ProjectContext?.Id);
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Protocol/InlayHints/InlayHintDataWrapper.cs

@@ -0,0 +1,8 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using Roslyn.LanguageServer.Protocol;
+
+namespace Microsoft.CodeAnalysis.Razor.Protocol.InlayHints;
+
+internal record class InlayHintDataWrapper(TextDocumentIdentifier TextDocument, object? OriginalData, Position OriginalPosition);

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteInlayHintService.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.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
+using Roslyn.LanguageServer.Protocol;
+
+namespace Microsoft.CodeAnalysis.Razor.Remote;
+
+internal interface IRemoteInlayHintService : IRemoteJsonService
+{
+    ValueTask<InlayHint[]?> GetInlayHintsAsync(JsonSerializableRazorPinnedSolutionInfoWrapper solutionInfo, JsonSerializableDocumentId razorDocumentId, InlayHintParams inlayHintParams, bool displayAllOverride, CancellationToken cancellationToken);
+
+    ValueTask<InlayHint> ResolveHintAsync(JsonSerializableRazorPinnedSolutionInfoWrapper solutionInfo, JsonSerializableDocumentId razorDocumentId, InlayHint inlayHint, CancellationToken cancellationToken);
+}

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

@@ -27,6 +27,7 @@ internal static class RazorServices
     internal static readonly IEnumerable<(Type, Type?)> JsonServices =
         [
             (typeof(IRemoteSignatureHelpService), null),
+            (typeof(IRemoteInlayHintService), null),
         ];
 
     private const string ComponentName = "Razor";

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/InlayHints/RemoteInlayHintService.cs

@@ -0,0 +1,115 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor.Language;
+using Microsoft.AspNetCore.Razor.Language.Syntax;
+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.InlayHints;
+using Microsoft.CodeAnalysis.Razor.Remote;
+using Microsoft.CodeAnalysis.Razor.Workspaces;
+using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
+using Microsoft.CodeAnalysis.Text;
+using Roslyn.LanguageServer.Protocol;
+
+namespace Microsoft.CodeAnalysis.Remote.Razor;
+
+internal sealed partial class RemoteInlayHintService(in ServiceArgs args) : RazorDocumentServiceBase(in args), IRemoteInlayHintService
+{
+    internal sealed class Factory : FactoryBase<IRemoteInlayHintService>
+    {
+        protected override IRemoteInlayHintService CreateService(in ServiceArgs args)
+            => new RemoteInlayHintService(in args);
+    }
+
+    private readonly IRazorDocumentMappingService _documentMappingService = args.ExportProvider.GetExportedValue<IRazorDocumentMappingService>();
+    private readonly IFilePathService _filePathService = args.ExportProvider.GetExportedValue<IFilePathService>();
+
+    public ValueTask<InlayHint[]?> GetInlayHintsAsync(JsonSerializableRazorPinnedSolutionInfoWrapper solutionInfo, JsonSerializableDocumentId razorDocumentId, InlayHintParams inlayHintParams, bool displayAllOverride, CancellationToken cancellationToken)
+       => RunServiceAsync(
+            solutionInfo,
+            razorDocumentId,
+            context => GetInlayHintsAsync(context, inlayHintParams, displayAllOverride, cancellationToken),
+            cancellationToken);
+
+    private async ValueTask<InlayHint[]?> GetInlayHintsAsync(RemoteDocumentContext context, InlayHintParams inlayHintParams, bool displayAllOverride, CancellationToken cancellationToken)
+    {
+        var codeDocument = await context.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
+        var csharpDocument = codeDocument.GetCSharpDocument();
+
+        var span = inlayHintParams.Range.ToLinePositionSpan();
+
+        // We are given a range by the client, but our mapping only succeeds if the start and end of the range can both be mapped
+        // to C#. Since that doesn't logically match what we want from inlay hints, we instead get the minimum range of mappable
+        // C# to get hints for. We'll filter that later, to remove the sections that can't be mapped back.
+        if (!_documentMappingService.TryMapToGeneratedDocumentRange(csharpDocument, span, out var projectedLinePositionSpan) &&
+            !codeDocument.TryGetMinimalCSharpRange(span, out projectedLinePositionSpan))
+        {
+            // There's no C# in the range.
+            return null;
+        }
+
+        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+
+        var textDocument = inlayHintParams.TextDocument.WithUri(generatedDocument.CreateUri());
+        var range = projectedLinePositionSpan.ToRange();
+
+        var hints = await InlayHints.GetInlayHintsAsync(generatedDocument, textDocument, range, displayAllOverride, cancellationToken).ConfigureAwait(false);
+
+        if (hints is null)
+        {
+            return null;
+        }
+
+        using var inlayHintsBuilder = new PooledArrayBuilder<InlayHint>();
+        var razorSourceText = codeDocument.Source.Text;
+        var csharpSourceText = codeDocument.GetCSharpSourceText();
+        var syntaxTree = codeDocument.GetSyntaxTree();
+        foreach (var hint in hints)
+        {
+            if (csharpSourceText.TryGetAbsoluteIndex(hint.Position.ToLinePosition(), out var absoluteIndex) &&
+                _documentMappingService.TryMapToHostDocumentPosition(csharpDocument, absoluteIndex, out var hostDocumentPosition, out var hostDocumentIndex))
+            {
+                // We know this C# maps to Razor, but does it map to Razor that we like?
+                var node = syntaxTree.Root.FindInnermostNode(hostDocumentIndex);
+                if (node?.FirstAncestorOrSelf<MarkupTagHelperAttributeValueSyntax>() is not null)
+                {
+                    continue;
+                }
+
+                if (hint.TextEdits is not null)
+                {
+                    var changes = hint.TextEdits.Select(csharpSourceText.GetTextChange);
+                    var mappedChanges = _documentMappingService.GetHostDocumentEdits(csharpDocument, changes);
+                    hint.TextEdits = mappedChanges.Select(razorSourceText.GetTextEdit).ToArray();
+                }
+
+                hint.Data = new InlayHintDataWrapper(inlayHintParams.TextDocument, hint.Data, hint.Position);
+                hint.Position = hostDocumentPosition.ToPosition();
+
+                inlayHintsBuilder.Add(hint);
+            }
+        }
+
+        return inlayHintsBuilder.ToArray();
+    }
+
+    public ValueTask<InlayHint> ResolveHintAsync(JsonSerializableRazorPinnedSolutionInfoWrapper solutionInfo, JsonSerializableDocumentId razorDocumentId, InlayHint inlayHint, CancellationToken cancellationToken)
+       => RunServiceAsync(
+            solutionInfo,
+            razorDocumentId,
+            context => ResolveInlayHintAsync(context, inlayHint, cancellationToken),
+            cancellationToken);
+
+    private async ValueTask<InlayHint> ResolveInlayHintAsync(RemoteDocumentContext context, InlayHint inlayHint, CancellationToken cancellationToken)
+    {
+        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+
+        return await InlayHints.ResolveInlayHintAsync(generatedDocument, inlayHint, cancellationToken).ConfigureAwait(false);
+    }
+}

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

@@ -19,6 +19,12 @@ public static async Task<Document> GetGeneratedDocumentAsync(this VersionedDocum
         Debug.Assert(documentContext.Snapshot is RemoteDocumentSnapshot, "This method only works on document contexts created in the OOP process");
 
         var snapshot = (RemoteDocumentSnapshot)documentContext.Snapshot;
+
+        if (snapshot.TryGetGeneratedDocument(out var generatedDocument))
+        {
+            return generatedDocument;
+        }
+
         var razorDocument = snapshot.TextDocument;
         var solution = razorDocument.Project.Solution;
 
@@ -27,13 +33,25 @@ public static async Task<Document> GetGeneratedDocumentAsync(this VersionedDocum
         var projectKey = razorDocument.Project.ToProjectKey();
         var generatedFilePath = filePathService.GetRazorCSharpFilePath(projectKey, razorDocument.FilePath.AssumeNotNull());
         var generatedDocumentId = solution.GetDocumentIdsWithFilePath(generatedFilePath).First(d => d.ProjectId == razorDocument.Project.Id);
-        var generatedDocument = solution.GetDocument(generatedDocumentId).AssumeNotNull();
+        generatedDocument = solution.GetDocument(generatedDocumentId).AssumeNotNull();
 
         var csharpSourceText = await documentContext.GetCSharpSourceTextAsync(cancellationToken).ConfigureAwait(false);
 
         // HACK: We're not in the same solution fork as the LSP server that provides content for this document
         generatedDocument = generatedDocument.WithText(csharpSourceText);
 
-        return generatedDocument;
+        // Obviously this lock is not sufficient to avoid wasted work, but it does at least avoid mutating the snapshot
+        // any more than just a once of caching of the generated document, which is what is really happening with the set
+        // method call below.
+        lock (snapshot)
+        {
+            if (snapshot.TryGetGeneratedDocument(out var generatedDocument2))
+            {
+                return generatedDocument2;
+            }
+
+            snapshot.SetGeneratedDocument(generatedDocument);
+            return generatedDocument;
+        }
     }
 }