Allow LSP and cohosting to provide specialized methods to get a syntax tree

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Definition/RazorComponentDefinitionService.cs

@@ -5,10 +5,15 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor.Language;
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.GoToDefinition;
 using Microsoft.CodeAnalysis.Razor.Logging;
+using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer.Definition;
@@ -19,4 +24,9 @@ internal sealed class RazorComponentDefinitionService(
     ILoggerFactory loggerFactory)
     : AbstractRazorComponentDefinitionService(componentSearchEngine, documentMappingService, loggerFactory.GetOrCreateLogger<RazorComponentDefinitionService>())
 {
+    protected override ValueTask<SyntaxTree> GetCSharpSyntaxTreeAsync(IDocumentSnapshot documentSnapshot, RazorCodeDocument codeDocument, CancellationToken cancellationToken)
+    {
+        var csharpText = codeDocument.GetCSharpSourceText();
+        return new(CSharpSyntaxTree.ParseText(csharpText, cancellationToken: cancellationToken));
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/GoToDefinition/AbstractRazorComponentDefinitionService.cs

@@ -70,9 +70,10 @@ private async Task<LspRange> GetNavigateRangeAsync(IDocumentSnapshot documentSna
             _logger.LogInformation($"Attempting to get definition from an attribute directly.");
 
             var originCodeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+            var syntaxTree = await GetCSharpSyntaxTreeAsync(documentSnapshot, originCodeDocument, cancellationToken).ConfigureAwait(false);
 
             var range = await RazorComponentDefinitionHelpers
-                .TryGetPropertyRangeAsync(originCodeDocument, attributeDescriptor.GetPropertyName(), _documentMappingService, _logger, cancellationToken)
+                .TryGetPropertyRangeAsync(originCodeDocument, syntaxTree, attributeDescriptor.GetPropertyName(), _documentMappingService, _logger, cancellationToken)
                 .ConfigureAwait(false);
 
             if (range is not null)
@@ -87,4 +88,6 @@ private async Task<LspRange> GetNavigateRangeAsync(IDocumentSnapshot documentSna
         // at least then press F7 to go there.
         return VsLspFactory.DefaultRange;
     }
+
+    protected abstract ValueTask<SyntaxTree> GetCSharpSyntaxTreeAsync(IDocumentSnapshot documentSnapshot, RazorCodeDocument codeDocument, CancellationToken cancellationToken);
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/GoToDefinition/RazorComponentDefinitionHelpers.cs

@@ -8,7 +8,6 @@
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.Language.Syntax;
-using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Logging;
@@ -131,12 +130,13 @@ static bool TryGetTagName(RazorSyntaxNode node, [NotNullWhen(true)] out RazorSyn
 
     public static async Task<LspRange?> TryGetPropertyRangeAsync(
         RazorCodeDocument codeDocument,
+        SyntaxTree csharpSyntaxTree,
         string propertyName,
         IDocumentMappingService documentMappingService,
         ILogger logger,
         CancellationToken cancellationToken)
     {
-        // Parse the C# file and find the property that matches the name.
+        // Process the C# tree and find the property that matches the name.
         // We don't worry about parameter attributes here for two main reasons:
         //   1. We don't have symbolic information, so the best we could do would be checking for any
         //      attribute named Parameter, regardless of which namespace. It also means we would have
@@ -147,9 +147,8 @@ static bool TryGetTagName(RazorSyntaxNode node, [NotNullWhen(true)] out RazorSyn
         //      tag helper attribute. If they don't have the [Parameter] attribute then the Razor compiler
         //      will error, but allowing them to Go To Def on that property regardless, actually helps
         //      them fix the error.
-        var csharpText = codeDocument.GetCSharpSourceText();
-        var syntaxTree = CSharpSyntaxTree.ParseText(csharpText, cancellationToken: cancellationToken);
-        var root = await syntaxTree.GetRootAsync(cancellationToken).ConfigureAwait(false);
+
+        var root = await csharpSyntaxTree.GetRootAsync(cancellationToken).ConfigureAwait(false);
 
         // Since we know how the compiler generates the C# source we can be a little specific here, and avoid
         // long tree walks. If the compiler ever changes how they generate their code, the tests for this will break
@@ -169,6 +168,7 @@ static bool TryGetTagName(RazorSyntaxNode node, [NotNullWhen(true)] out RazorSyn
                 return null;
             }
 
+            var csharpText = codeDocument.GetCSharpSourceText();
             var range = csharpText.GetRange(property.Identifier.Span);
             if (documentMappingService.TryMapToHostDocumentRange(codeDocument.GetCSharpDocument(), range, out var originalRange))
             {

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

@@ -68,7 +68,7 @@ private async ValueTask<Response> GetHighlightsAsync(
         var csharpDocument = codeDocument.GetCSharpDocument();
         if (_documentMappingService.TryMapToGeneratedDocumentPosition(csharpDocument, index, out var mappedPosition, out _))
         {
-            var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+            var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
             var highlights = await DocumentHighlights.GetHighlightsAsync(generatedDocument, mappedPosition, cancellationToken).ConfigureAwait(false);
 

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/FoldingRanges/RemoteFoldingRangeService.cs

@@ -43,7 +43,7 @@ private async ValueTask<ImmutableArray<RemoteFoldingRange>> GetFoldingRangesAsyn
         ImmutableArray<RemoteFoldingRange> htmlRanges,
         CancellationToken cancellationToken)
     {
-        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         var csharpRanges = await ExternalHandlers.FoldingRanges.GetFoldingRangesAsync(generatedDocument, cancellationToken).ConfigureAwait(false);
 

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

@@ -2,10 +2,17 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System.Composition;
+using System.Diagnostics;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor;
+using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.GoToDefinition;
 using Microsoft.CodeAnalysis.Razor.Logging;
+using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
+using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor.GoToDefinition;
 
@@ -14,7 +21,25 @@ namespace Microsoft.CodeAnalysis.Remote.Razor.GoToDefinition;
 internal sealed class RazorComponentDefinitionService(
     IRazorComponentSearchEngine componentSearchEngine,
     IDocumentMappingService documentMappingService,
+    IFilePathService filePathService,
     ILoggerFactory loggerFactory)
     : AbstractRazorComponentDefinitionService(componentSearchEngine, documentMappingService, loggerFactory.GetOrCreateLogger<RazorComponentDefinitionService>())
 {
+    private readonly IFilePathService _filePathService = filePathService;
+
+    protected override async ValueTask<SyntaxTree> GetCSharpSyntaxTreeAsync(IDocumentSnapshot documentSnapshot, RazorCodeDocument codeDocument, CancellationToken cancellationToken)
+    {
+        Debug.Assert(documentSnapshot is RemoteDocumentSnapshot, "This method only works on document snapshots created in the OOP process");
+
+        var remoteSnapshot = (RemoteDocumentSnapshot)documentSnapshot;
+        var document = await remoteSnapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
+
+        if (document.TryGetSyntaxTree(out var syntaxTree))
+        {
+            return syntaxTree;
+        }
+
+        var tree = await document.GetSyntaxTreeAsync(cancellationToken).ConfigureAwait(false);
+        return tree.AssumeNotNull();
+    }
 }

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

@@ -62,14 +62,6 @@ protected override IRemoteGoToDefinitionService CreateService(in ServiceArgs arg
 
         var positionInfo = GetPositionInfo(codeDocument, hostDocumentIndex);
 
-        // First, see if this is a Razor component.
-        var componentLocation = await _componentDefinitionService.GetDefinitionAsync(context.Snapshot, positionInfo, ignoreAttributes: false, cancellationToken).ConfigureAwait(false);
-        if (componentLocation is not null)
-        {
-            // Convert from VS LSP Location to Roslyn. This can be removed when Razor moves fully onto Roslyn's LSP types.
-            return Results([RoslynLspFactory.CreateLocation(componentLocation.Uri, componentLocation.Range.ToLinePositionSpan())]);
-        }
-
         if (positionInfo.LanguageKind == RazorLanguageKind.Html)
         {
             // Sometimes Html can actually be mapped to C#, like for example component attributes, which map to
@@ -83,9 +75,17 @@ protected override IRemoteGoToDefinitionService CreateService(in ServiceArgs arg
             }
         }
 
-        // If it isn't a Razor component, and it isn't C#, let the server know to delegate to HTML.
-        if (positionInfo.LanguageKind != RazorLanguageKind.CSharp)
+        if (positionInfo.LanguageKind is RazorLanguageKind.Html or RazorLanguageKind.Razor)
         {
+            // First, see if this is a Razor component. We ignore attributes here, because they're better served by the C# handler.
+            var componentLocation = await _componentDefinitionService.GetDefinitionAsync(context.Snapshot, positionInfo, ignoreAttributes: true, cancellationToken).ConfigureAwait(false);
+            if (componentLocation is not null)
+            {
+                // Convert from VS LSP Location to Roslyn. This can be removed when Razor moves fully onto Roslyn's LSP types.
+                return Results([RoslynLspFactory.CreateLocation(componentLocation.Uri, componentLocation.Range.ToLinePositionSpan())]);
+            }
+
+            // If it isn't a Razor component, and it isn't C#, let the server know to delegate to HTML.
             return CallHtml;
         }
 
@@ -96,7 +96,7 @@ protected override IRemoteGoToDefinitionService CreateService(in ServiceArgs arg
         }
 
         // Finally, call into C#.
-        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         var locations = await ExternalHandlers.GoToDefinition
             .GetDefinitionsAsync(

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

@@ -52,7 +52,7 @@ protected override IRemoteInlayHintService CreateService(in ServiceArgs args)
             return null;
         }
 
-        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         var textDocument = inlayHintParams.TextDocument.WithUri(generatedDocument.CreateUri());
         var range = projectedLinePositionSpan.ToRange();
@@ -106,7 +106,7 @@ public ValueTask<InlayHint> ResolveHintAsync(JsonSerializableRazorPinnedSolution
 
     private async ValueTask<InlayHint> ResolveInlayHintAsync(RemoteDocumentContext context, InlayHint inlayHint, CancellationToken cancellationToken)
     {
-        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         return await InlayHints.ResolveInlayHintAsync(generatedDocument, inlayHint, cancellationToken).ConfigureAwait(false);
     }

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

@@ -8,7 +8,9 @@ namespace Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 
 internal class RemoteDocumentContext : VersionedDocumentContext
 {
-    public TextDocument TextDocument => ((RemoteDocumentSnapshot)Snapshot).TextDocument;
+    public TextDocument TextDocument => Snapshot.TextDocument;
+
+    public new RemoteDocumentSnapshot Snapshot => (RemoteDocumentSnapshot)base.Snapshot;
 
     public RemoteDocumentContext(Uri uri, RemoteDocumentSnapshot snapshot)
         // HACK: Need to revisit version and projectContext here I guess

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

@@ -1,13 +1,14 @@
 // 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 System.Diagnostics.CodeAnalysis;
 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.Workspaces;
 using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
@@ -88,14 +89,31 @@ public bool TryGetGeneratedOutput([NotNullWhen(true)] out RazorCodeDocument? res
         return result is not null;
     }
 
-    public async Task<Document> GetOrAddGeneratedDocumentAsync<TArg>(TArg arg, Func<TArg, Task<Document>> createGeneratedDocument)
+    public async Task<Document> GetGeneratedDocumentAsync(IFilePathService filePathService)
     {
         if (_generatedDocument is Document generatedDocument)
         {
             return generatedDocument;
         }
 
-        generatedDocument = await createGeneratedDocument(arg);
+        generatedDocument = await HACK_GenerateDocumentAsync(filePathService).ConfigureAwait(false);
         return InterlockedOperations.Initialize(ref _generatedDocument, generatedDocument);
     }
+
+    private async Task<Document> HACK_GenerateDocumentAsync(IFilePathService filePathService)
+    {
+        // TODO: A real implementation needs to get the SourceGeneratedDocument from the solution
+
+        var solution = TextDocument.Project.Solution;
+        var generatedFilePath = filePathService.GetRazorCSharpFilePath(Project.Key, FilePath.AssumeNotNull());
+        var projectId = TextDocument.Project.Id;
+        var generatedDocumentId = solution.GetDocumentIdsWithFilePath(generatedFilePath).First(d => d.ProjectId == projectId);
+        var generatedDocument = solution.GetRequiredDocument(generatedDocumentId);
+
+        var codeDocument = await this.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var csharpSourceText = codeDocument.GetCSharpSourceText();
+
+        // HACK: We're not in the same solution fork as the LSP server that provides content for this document
+        return generatedDocument.WithText(csharpSourceText);
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Rename/RemoteRenameService.cs

@@ -53,7 +53,7 @@ protected override IRemoteRenameService CreateService(in ServiceArgs args)
             return NoFurtherHandling;
         }
 
-        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         var razorEdit = await _renameService.TryGetRazorRenameEditsAsync(context, positionInfo, newName, cancellationToken).ConfigureAwait(false);
         if (razorEdit is not null)

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/SemanticTokens/RemoteCSharpSemanticTokensProvider.cs

@@ -4,13 +4,15 @@
 using System;
 using System.Collections.Immutable;
 using System.Composition;
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Telemetry;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost.Handlers;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.SemanticTokens;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
+using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor.SemanticTokens;
@@ -27,7 +29,9 @@ internal class RemoteCSharpSemanticTokensProvider(IFilePathService filePathServi
         using var _ = _telemetryReporter.TrackLspRequest(nameof(SemanticTokensRange.GetSemanticTokensAsync), Constants.ExternalAccessServerName, correlationId);
 
         // We have a razor document, lets find the generated C# document
-        var generatedDocument = await documentContext.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        Debug.Assert(documentContext is RemoteDocumentContext, "This method only works on document snapshots created in the OOP process");
+        var snapshot = (RemoteDocumentSnapshot)documentContext.Snapshot;
+        var generatedDocument = await snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         var data = await SemanticTokensRange.GetSemanticTokensAsync(
             generatedDocument,