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

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

@@ -1,11 +1,6 @@
 // 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.Linq;
-using System.Text;
-using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.GoToDefinition;
 using Microsoft.CodeAnalysis.Razor.Logging;

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

@@ -69,10 +69,8 @@ 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 range = await RazorComponentDefinitionHelpers
-                .TryGetPropertyRangeAsync(originCodeDocument, attributeDescriptor.GetPropertyName(), _documentMappingService, _logger, cancellationToken)
+                .TryGetPropertyRangeAsync(documentSnapshot, attributeDescriptor.GetPropertyName(), _documentMappingService, _logger, cancellationToken)
                 .ConfigureAwait(false);
 
             if (range is not null)

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

@@ -8,10 +8,10 @@
 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;
+using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 using LspRange = Microsoft.VisualStudio.LanguageServer.Protocol.Range;
@@ -130,13 +130,13 @@ static bool TryGetTagName(RazorSyntaxNode node, [NotNullWhen(true)] out RazorSyn
     }
 
     public static async Task<LspRange?> TryGetPropertyRangeAsync(
-        RazorCodeDocument codeDocument,
+        IDocumentSnapshot documentSnapshot,
         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,10 @@ 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 csharpSyntaxTree = await documentSnapshot.GetCSharpSyntaxTreeAsync(cancellationToken).ConfigureAwait(false);
+        var root = await csharpSyntaxTree.GetRootAsync(cancellationToken).ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().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 +170,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.Razor.Workspaces/ProjectSystem/DocumentSnapshot.cs

@@ -3,8 +3,10 @@
 
 using System;
 using System.Diagnostics.CodeAnalysis;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
+using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
@@ -62,4 +64,11 @@ public IDocumentSnapshot WithText(SourceText text)
     {
         return new DocumentSnapshot(ProjectInternal, State.WithText(text, VersionStamp.Create()));
     }
+
+    public async Task<SyntaxTree> GetCSharpSyntaxTreeAsync(CancellationToken cancellationToken)
+    {
+        var codeDocument = await GetGeneratedOutputAsync().ConfigureAwait(false);
+        var csharpText = codeDocument.GetCSharpSourceText();
+        return CSharpSyntaxTree.ParseText(csharpText, cancellationToken: cancellationToken);
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/IDocumentSnapshot.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System.Diagnostics.CodeAnalysis;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Text;
@@ -20,6 +21,12 @@ internal interface IDocumentSnapshot
     Task<VersionStamp> GetTextVersionAsync();
     Task<RazorCodeDocument> GetGeneratedOutputAsync();
 
+    /// <summary>
+    /// Gets the Roslyn syntax tree for the generated C# for this Razor document
+    /// </summary>
+    /// <remarks>Using this from the LSP server side of things is not ideal. Use sparingly :)</remarks>
+    Task<SyntaxTree> GetCSharpSyntaxTreeAsync(CancellationToken cancellationToken);
+
     bool TryGetText([NotNullWhen(true)] out SourceText? result);
     bool TryGetTextVersion(out VersionStamp result);
     bool TryGetGeneratedOutput([NotNullWhen(true)] out RazorCodeDocument? result);

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/ImportDocumentSnapshot.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Text;
@@ -75,4 +76,7 @@ public bool TryGetGeneratedOutput([NotNullWhen(true)] out RazorCodeDocument? res
 
     public IDocumentSnapshot WithText(SourceText text)
         => throw new NotSupportedException();
+
+    public Task<SyntaxTree> GetCSharpSyntaxTreeAsync(CancellationToken cancellationToken)
+        => throw new NotSupportedException();
 }

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

@@ -11,7 +11,6 @@
 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 static Microsoft.VisualStudio.LanguageServer.Protocol.VsLspExtensions;
@@ -29,7 +28,6 @@ protected override IRemoteDocumentHighlightService CreateService(in ServiceArgs
     }
 
     private readonly IDocumentMappingService _documentMappingService = args.ExportProvider.GetExportedValue<IDocumentMappingService>();
-    private readonly IFilePathService _filePathService = args.ExportProvider.GetExportedValue<IFilePathService>();
 
     public ValueTask<Response> GetHighlightsAsync(
         RazorPinnedSolutionInfoWrapper solutionInfo,
@@ -68,7 +66,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().ConfigureAwait(false);
 
             var highlights = await DocumentHighlights.GetHighlightsAsync(generatedDocument, mappedPosition, cancellationToken).ConfigureAwait(false);
 

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

@@ -9,7 +9,6 @@
 using Microsoft.CodeAnalysis.Razor.FoldingRanges;
 using Microsoft.CodeAnalysis.Razor.Protocol.Folding;
 using Microsoft.CodeAnalysis.Razor.Remote;
-using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 using ExternalHandlers = Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost.Handlers;
@@ -25,7 +24,6 @@ protected override IRemoteFoldingRangeService CreateService(in ServiceArgs args)
     }
 
     private readonly IFoldingRangeService _foldingRangeService = args.ExportProvider.GetExportedValue<IFoldingRangeService>();
-    private readonly IFilePathService _filePathService = args.ExportProvider.GetExportedValue<IFilePathService>();
 
     public ValueTask<ImmutableArray<RemoteFoldingRange>> GetFoldingRangesAsync(
         RazorPinnedSolutionInfoWrapper solutionInfo,
@@ -43,7 +41,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().ConfigureAwait(false);
 
         var csharpRanges = await ExternalHandlers.FoldingRanges.GetFoldingRangesAsync(generatedDocument, cancellationToken).ConfigureAwait(false);
 

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

@@ -10,7 +10,6 @@
 using Microsoft.CodeAnalysis.Razor.GoToDefinition;
 using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Remote;
-using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Remote.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Text;
@@ -33,7 +32,6 @@ protected override IRemoteGoToDefinitionService CreateService(in ServiceArgs arg
     }
 
     private readonly IRazorComponentDefinitionService _componentDefinitionService = args.ExportProvider.GetExportedValue<IRazorComponentDefinitionService>();
-    private readonly IFilePathService _filePathService = args.ExportProvider.GetExportedValue<IFilePathService>();
 
     protected override IDocumentPositionInfoStrategy DocumentPositionInfoStrategy => PreferAttributeNameDocumentPositionInfoStrategy.Instance;
 
@@ -62,14 +60,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 +73,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 +94,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().ConfigureAwait(false);
 
         var locations = await ExternalHandlers.GoToDefinition
             .GetDefinitionsAsync(
@@ -121,7 +119,7 @@ protected override IRemoteGoToDefinitionService CreateService(in ServiceArgs arg
             var (uri, range) = location;
 
             var (mappedDocumentUri, mappedRange) = await DocumentMappingService
-                .MapToHostDocumentUriAndRangeAsync((RemoteDocumentSnapshot)context.Snapshot, uri, range.ToLinePositionSpan(), cancellationToken)
+                .MapToHostDocumentUriAndRangeAsync(context.Snapshot, uri, range.ToLinePositionSpan(), cancellationToken)
                 .ConfigureAwait(false);
 
             var mappedLocation = RoslynLspFactory.CreateLocation(mappedDocumentUri, mappedRange);

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

@@ -26,8 +26,6 @@ protected override IRemoteInlayHintService CreateService(in ServiceArgs args)
             => new RemoteInlayHintService(in args);
     }
 
-    private readonly IFilePathService _filePathService = args.ExportProvider.GetExportedValue<IFilePathService>();
-
     public ValueTask<InlayHint[]?> GetInlayHintsAsync(JsonSerializableRazorPinnedSolutionInfoWrapper solutionInfo, JsonSerializableDocumentId razorDocumentId, InlayHintParams inlayHintParams, bool displayAllOverride, CancellationToken cancellationToken)
        => RunServiceAsync(
             solutionInfo,
@@ -52,7 +50,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().ConfigureAwait(false);
 
         var textDocument = inlayHintParams.TextDocument.WithUri(generatedDocument.CreateUri());
         var range = projectedLinePositionSpan.ToRange();
@@ -106,7 +104,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().ConfigureAwait(false);
 
         return await InlayHints.ResolveInlayHintAsync(generatedDocument, inlayHint, cancellationToken).ConfigureAwait(false);
     }

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

@@ -4,26 +4,28 @@
 using System;
 using System.Composition;
 using System.Runtime.CompilerServices;
+using Microsoft.CodeAnalysis.Razor.Workspaces;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 
 [Export(typeof(DocumentSnapshotFactory)), Shared]
 [method: ImportingConstructor]
-internal class DocumentSnapshotFactory(Lazy<ProjectSnapshotFactory> projectSnapshotFactory)
+internal class DocumentSnapshotFactory(Lazy<ProjectSnapshotFactory> projectSnapshotFactory, IFilePathService filePathService)
 {
-    private static readonly ConditionalWeakTable<TextDocument, RemoteDocumentSnapshot> _documentSnapshots = new();
+    private static readonly ConditionalWeakTable<TextDocument, RemoteDocumentSnapshot> s_documentSnapshots = new();
 
     private readonly Lazy<ProjectSnapshotFactory> _projectSnapshotFactory = projectSnapshotFactory;
+    private readonly IFilePathService _filePathService = filePathService;
 
     public RemoteDocumentSnapshot GetOrCreate(TextDocument textDocument)
     {
-        lock (_documentSnapshots)
+        lock (s_documentSnapshots)
         {
-            if (!_documentSnapshots.TryGetValue(textDocument, out var documentSnapshot))
+            if (!s_documentSnapshots.TryGetValue(textDocument, out var documentSnapshot))
             {
                 var projectSnapshot = _projectSnapshotFactory.Value.GetOrCreate(textDocument.Project);
-                documentSnapshot = new RemoteDocumentSnapshot(textDocument, projectSnapshot);
-                _documentSnapshots.Add(textDocument, documentSnapshot);
+                documentSnapshot = new RemoteDocumentSnapshot(textDocument, projectSnapshot, _filePathService);
+                s_documentSnapshots.Add(textDocument, documentSnapshot);
             }
 
             return documentSnapshot;