Cohost Go To Implementation

eng/targets/Services.props

@@ -29,5 +29,6 @@
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.GoToDefinition" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteGoToDefinitionService+Factory" />
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.Rename" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteRenameService+Factory" />
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.AutoInsert" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteAutoInsertService+Factory" />
+    <ServiceHubService Include="Microsoft.VisualStudio.Razor.GoToImplementation" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteGoToImplementationService+Factory" />
   </ItemGroup>
 </Project>

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

@@ -184,10 +184,10 @@ static void AddHandlers(IServiceCollection services, LanguageServerFeatureOption
             // Transient because it should only be used once and I'm hoping it doesn't stick around.
             services.AddTransient<IOnInitialized>(sp => sp.GetRequiredService<RazorConfigurationEndpoint>());
 
-            services.AddHandlerWithCapabilities<ImplementationEndpoint>();
-
             if (!featureOptions.UseRazorCohostServer)
             {
+                services.AddHandlerWithCapabilities<ImplementationEndpoint>();
+
                 services.AddSingleton<IRazorComponentDefinitionService, RazorComponentDefinitionService>();
                 services.AddHandlerWithCapabilities<DefinitionEndpoint>();
 

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

@@ -0,0 +1,19 @@
+// 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 RoslynLocation = Roslyn.LanguageServer.Protocol.Location;
+using RoslynPosition = Roslyn.LanguageServer.Protocol.Position;
+
+namespace Microsoft.CodeAnalysis.Razor.Remote;
+
+internal interface IRemoteGoToImplementationService : IRemoteJsonService
+{
+    ValueTask<RemoteResponse<RoslynLocation[]?>> GetImplementationAsync(
+        JsonSerializableRazorPinnedSolutionInfoWrapper solutionInfo,
+        JsonSerializableDocumentId razorDocumentId,
+        RoslynPosition position,
+        CancellationToken cancellationToken);
+}

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

@@ -32,6 +32,7 @@ internal static class RazorServices
             (typeof(IRemoteInlayHintService), null),
             (typeof(IRemoteDocumentSymbolService), null),
             (typeof(IRemoteRenameService), null),
+            (typeof(IRemoteGoToImplementationService), null),
         ];
 
     private const string ComponentName = "Razor";

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/GoToDefinition/RemoteGoToDefinitionService.cs

@@ -58,20 +58,7 @@ protected override IRemoteGoToDefinitionService CreateService(in ServiceArgs arg
             return NoFurtherHandling;
         }
 
-        var positionInfo = GetPositionInfo(codeDocument, hostDocumentIndex);
-
-        if (positionInfo.LanguageKind == RazorLanguageKind.Html)
-        {
-            // Sometimes Html can actually be mapped to C#, like for example component attributes, which map to
-            // C# properties, even though they appear entirely in a Html context. Since remapping is pretty cheap
-            // it's easier to just try mapping, and see what happens, rather than checking for specific syntax nodes.
-            if (DocumentMappingService.TryMapToGeneratedDocumentPosition(codeDocument.GetCSharpDocument(), positionInfo.HostDocumentIndex, out VsPosition? csharpPosition, out _))
-            {
-                // We're just gonna pretend this mapped perfectly normally onto C#. Moving this logic to the actual position info
-                // calculating code is possible, but could have untold effects, so opt-in is better (for now?)
-                positionInfo = positionInfo with { LanguageKind = RazorLanguageKind.CSharp, Position = csharpPosition };
-            }
-        }
+        var positionInfo = GetPositionInfo(codeDocument, hostDocumentIndex, preferCSharpOverHtml: true);
 
         if (positionInfo.LanguageKind is RazorLanguageKind.Html or RazorLanguageKind.Razor)
         {

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/GoToImplementation/RemoteGoToImplementationService.cs

@@ -0,0 +1,111 @@
+// 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.Razor.DocumentMapping;
+using Microsoft.CodeAnalysis.Razor.Protocol;
+using Microsoft.CodeAnalysis.Razor.Remote;
+using Microsoft.CodeAnalysis.Remote.Razor.DocumentMapping;
+using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
+using Microsoft.CodeAnalysis.Text;
+using Microsoft.VisualStudio.LanguageServer.Protocol;
+using Roslyn.LanguageServer.Protocol;
+using static Microsoft.CodeAnalysis.Razor.Remote.RemoteResponse<Roslyn.LanguageServer.Protocol.Location[]?>;
+using ExternalHandlers = Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost.Handlers;
+using RoslynLocation = Roslyn.LanguageServer.Protocol.Location;
+using RoslynPosition = Roslyn.LanguageServer.Protocol.Position;
+using VsPosition = Microsoft.VisualStudio.LanguageServer.Protocol.Position;
+
+namespace Microsoft.CodeAnalysis.Remote.Razor;
+
+internal sealed class RemoteGoToImplementationService(in ServiceArgs args) : RazorDocumentServiceBase(in args), IRemoteGoToImplementationService
+{
+    internal sealed class Factory : FactoryBase<IRemoteGoToImplementationService>
+    {
+        protected override IRemoteGoToImplementationService CreateService(in ServiceArgs args)
+            => new RemoteGoToImplementationService(in args);
+    }
+
+    protected override IDocumentPositionInfoStrategy DocumentPositionInfoStrategy => PreferAttributeNameDocumentPositionInfoStrategy.Instance;
+
+    public ValueTask<RemoteResponse<RoslynLocation[]?>> GetImplementationAsync(
+        JsonSerializableRazorPinnedSolutionInfoWrapper solutionInfo,
+        JsonSerializableDocumentId documentId,
+        RoslynPosition position,
+        CancellationToken cancellationToken)
+        => RunServiceAsync(
+            solutionInfo,
+            documentId,
+            context => GetImplementationAsync(context, position, cancellationToken),
+            cancellationToken);
+
+    private async ValueTask<RemoteResponse<RoslynLocation[]?>> GetImplementationAsync(
+        RemoteDocumentContext context,
+        RoslynPosition position,
+        CancellationToken cancellationToken)
+    {
+        var codeDocument = await context.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
+
+        if (!codeDocument.Source.Text.TryGetAbsoluteIndex(position, out var hostDocumentIndex))
+        {
+            return NoFurtherHandling;
+        }
+
+        var positionInfo = GetPositionInfo(codeDocument, hostDocumentIndex, preferCSharpOverHtml: true);
+
+        if (positionInfo.LanguageKind is RazorLanguageKind.Razor)
+        {
+            return NoFurtherHandling;
+        }
+
+        if (positionInfo.LanguageKind is RazorLanguageKind.Html)
+        {
+            return CallHtml;
+        }
+
+        if (!DocumentMappingService.TryMapToGeneratedDocumentPosition(codeDocument.GetCSharpDocument(), positionInfo.HostDocumentIndex, out var mappedPosition, out _))
+        {
+            // If we can't map to the generated C# file, we're done.
+            return NoFurtherHandling;
+        }
+
+        // Finally, call into C#.
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+
+        var locations = await ExternalHandlers.GoToImplementation
+            .FindImplementationsAsync(
+                generatedDocument,
+                mappedPosition,
+                supportsVisualStudioExtensions: true,
+                cancellationToken)
+            .ConfigureAwait(false);
+
+        if (locations is null and not [])
+        {
+            // C# didn't return anything, so we're done.
+            return NoFurtherHandling;
+        }
+
+        // Map the C# locations back to the Razor file.
+        using var mappedLocations = new PooledArrayBuilder<RoslynLocation>(locations.Length);
+
+        foreach (var location in locations)
+        {
+            var (uri, range) = location;
+
+            var (mappedDocumentUri, mappedRange) = await DocumentMappingService
+                .MapToHostDocumentUriAndRangeAsync(context.Snapshot, uri, range.ToLinePositionSpan(), cancellationToken)
+                .ConfigureAwait(false);
+
+            var mappedLocation = RoslynLspFactory.CreateLocation(mappedDocumentUri, mappedRange);
+
+            mappedLocations.Add(mappedLocation);
+        }
+
+        return Results(mappedLocations.ToArray());
+    }
+}

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

@@ -7,6 +7,7 @@
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
+using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
@@ -24,31 +25,57 @@ internal abstract class RazorDocumentServiceBase(in ServiceArgs args) : RazorBro
     protected virtual IDocumentPositionInfoStrategy DocumentPositionInfoStrategy { get; } = DefaultDocumentPositionInfoStrategy.Instance;
 
     protected DocumentPositionInfo GetPositionInfo(RazorCodeDocument codeDocument, int hostDocumentIndex)
+        => GetPositionInfo(codeDocument, hostDocumentIndex, preferCSharpOverHtml: false);
+
+    protected DocumentPositionInfo GetPositionInfo(RazorCodeDocument codeDocument, int hostDocumentIndex, bool preferCSharpOverHtml)
     {
-        return DocumentPositionInfoStrategy.GetPositionInfo(DocumentMappingService, codeDocument, hostDocumentIndex);
+        var positionInfo = DocumentPositionInfoStrategy.GetPositionInfo(DocumentMappingService, codeDocument, hostDocumentIndex);
+
+        if (preferCSharpOverHtml && positionInfo.LanguageKind == RazorLanguageKind.Html)
+        {
+            // Sometimes Html can actually be mapped to C#, like for example component attributes, which map to
+            // C# properties, even though they appear entirely in a Html context. Since remapping is pretty cheap
+            // it's easier to just try mapping, and see what happens, rather than checking for specific syntax nodes.
+            if (DocumentMappingService.TryMapToGeneratedDocumentPosition(codeDocument.GetCSharpDocument(), positionInfo.HostDocumentIndex, out VsPosition? csharpPosition, out _))
+            {
+                // We're just gonna pretend this mapped perfectly normally onto C#. Moving this logic to the actual position info
+                // calculating code is possible, but could have untold effects, so opt-in is better (for now?)
+
+                // TODO: Not using a with operator here because it doesn't work in OOP for some reason.
+                positionInfo = new DocumentPositionInfo(RazorLanguageKind.CSharp, csharpPosition, positionInfo.HostDocumentIndex);
+            }
+        }
+
+        return positionInfo;
     }
 
     protected bool TryGetDocumentPositionInfo(RazorCodeDocument codeDocument, RoslynPosition position, out DocumentPositionInfo positionInfo)
+        => TryGetDocumentPositionInfo(codeDocument, position, preferCSharpOverHtml: false, out positionInfo);
+
+    protected bool TryGetDocumentPositionInfo(RazorCodeDocument codeDocument, RoslynPosition position, bool preferCSharpOverHtml, out DocumentPositionInfo positionInfo)
     {
         if (!codeDocument.Source.Text.TryGetAbsoluteIndex(position, out var hostDocumentIndex))
         {
             positionInfo = default;
             return false;
         }
 
-        positionInfo = GetPositionInfo(codeDocument, hostDocumentIndex);
+        positionInfo = GetPositionInfo(codeDocument, hostDocumentIndex, preferCSharpOverHtml);
         return true;
     }
 
     protected bool TryGetDocumentPositionInfo(RazorCodeDocument codeDocument, VsPosition position, out DocumentPositionInfo positionInfo)
+        => TryGetDocumentPositionInfo(codeDocument, position, preferCSharpOverHtml: false, out positionInfo);
+
+    protected bool TryGetDocumentPositionInfo(RazorCodeDocument codeDocument, VsPosition position, bool preferCSharpOverHtml, out DocumentPositionInfo positionInfo)
     {
         if (!codeDocument.Source.Text.TryGetAbsoluteIndex(position, out var hostDocumentIndex))
         {
             positionInfo = default;
             return false;
         }
 
-        positionInfo = GetPositionInfo(codeDocument, hostDocumentIndex);
+        positionInfo = GetPositionInfo(codeDocument, hostDocumentIndex, preferCSharpOverHtml);
         return true;
     }