Use precise ranges when requesting for semantic tokens with feature flags

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/LanguageServer/RazorSemanticTokensBenchmark.cs

@@ -15,6 +15,7 @@
 using Microsoft.AspNetCore.Razor.LanguageServer.Semantic;
 using Microsoft.CodeAnalysis.Razor;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
+using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
@@ -131,10 +132,11 @@ internal class TestRazorSemanticTokensInfoService : RazorSemanticTokensInfoServi
     {
         public TestRazorSemanticTokensInfoService(
             ClientNotifierServiceBase languageServer,
+            LanguageServerFeatureOptions languageServerFeatureOptions,
             IRazorDocumentMappingService documentMappingService,
             RazorLSPOptionsMonitor razorLSPOptionsMonitor,
             ILoggerFactory loggerFactory)
-            : base(languageServer, documentMappingService, razorLSPOptionsMonitor, loggerFactory)
+            : base(languageServer, documentMappingService, razorLSPOptionsMonitor, languageServerFeatureOptions, loggerFactory)
         {
         }
 

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/LanguageServer/RazorSemanticTokensRangeEndpointBenchmark.cs

@@ -16,6 +16,7 @@
 using Microsoft.AspNetCore.Razor.LanguageServer.Semantic;
 using Microsoft.CodeAnalysis.Razor;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
+using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
@@ -151,10 +152,11 @@ internal class TestCustomizableRazorSemanticTokensInfoService : RazorSemanticTok
     {
         public TestCustomizableRazorSemanticTokensInfoService(
             ClientNotifierServiceBase languageServer,
+            LanguageServerFeatureOptions languageServerFeatureOptions,
             IRazorDocumentMappingService documentMappingService,
             RazorLSPOptionsMonitor razorLSPOptionsMonitor,
             ILoggerFactory loggerFactory)
-            : base(languageServer, documentMappingService, razorLSPOptionsMonitor, loggerFactory)
+            : base(languageServer, documentMappingService, razorLSPOptionsMonitor, languageServerFeatureOptions, loggerFactory)
         {
         }
 

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

@@ -19,6 +19,7 @@ internal class ConfigurableLanguageServerFeatureOptions : LanguageServerFeatureO
     private readonly bool? _delegateToCSharpOnDiagnosticPublish;
     private readonly bool? _returnCodeActionAndRenamePathsWithPrefixedSlash;
     private readonly bool? _showAllCSharpCodeActions;
+    private readonly bool? _usePreciseSemanticTokenRanges;
     private readonly bool? _updateBuffersForClosedDocuments;
     private readonly bool? _includeProjectKeyInGeneratedFilePath;
 
@@ -31,6 +32,7 @@ internal class ConfigurableLanguageServerFeatureOptions : LanguageServerFeatureO
     public override bool DelegateToCSharpOnDiagnosticPublish => _delegateToCSharpOnDiagnosticPublish ?? _defaults.DelegateToCSharpOnDiagnosticPublish;
     public override bool ReturnCodeActionAndRenamePathsWithPrefixedSlash => _returnCodeActionAndRenamePathsWithPrefixedSlash ?? _defaults.ReturnCodeActionAndRenamePathsWithPrefixedSlash;
     public override bool ShowAllCSharpCodeActions => _showAllCSharpCodeActions ?? _defaults.ShowAllCSharpCodeActions;
+    public override bool UsePreciseSemanticTokenRanges => _usePreciseSemanticTokenRanges ?? _defaults.UsePreciseSemanticTokenRanges;
     public override bool UpdateBuffersForClosedDocuments => _updateBuffersForClosedDocuments ?? _defaults.UpdateBuffersForClosedDocuments;
     public override bool IncludeProjectKeyInGeneratedFilePath => _includeProjectKeyInGeneratedFilePath ?? _defaults.IncludeProjectKeyInGeneratedFilePath;
 
@@ -52,6 +54,7 @@ public ConfigurableLanguageServerFeatureOptions(string[] args)
             TryProcessBoolOption(nameof(DelegateToCSharpOnDiagnosticPublish), ref _delegateToCSharpOnDiagnosticPublish, option, args, i);
             TryProcessBoolOption(nameof(ReturnCodeActionAndRenamePathsWithPrefixedSlash), ref _returnCodeActionAndRenamePathsWithPrefixedSlash, option, args, i);
             TryProcessBoolOption(nameof(ShowAllCSharpCodeActions), ref _showAllCSharpCodeActions, option, args, i);
+            TryProcessBoolOption(nameof(UsePreciseSemanticTokenRanges), ref _usePreciseSemanticTokenRanges, option, args, i);
             TryProcessBoolOption(nameof(UpdateBuffersForClosedDocuments), ref _updateBuffersForClosedDocuments, option, args, i);
             TryProcessBoolOption(nameof(IncludeProjectKeyInGeneratedFilePath), ref _includeProjectKeyInGeneratedFilePath, option, args, i);
         }

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

@@ -36,4 +36,6 @@ public override bool ReturnCodeActionAndRenamePathsWithPrefixedSlash
     public override bool ShowAllCSharpCodeActions => false;
 
     public override bool IncludeProjectKeyInGeneratedFilePath => false;
+
+    public override bool UsePreciseSemanticTokenRanges => false;
 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Extensions/RangeExtensions.cs

@@ -176,4 +176,16 @@ public static bool IsUndefined(this Range range)
 
         return range == UndefinedRange;
     }
+
+    public static int CompareTo(this Range range1, Range range2)
+    {
+        var result = range1.Start.CompareTo(range2.Start);
+
+        if (result == 0)
+        {
+            result = range1.End.CompareTo(range2.End);
+        }
+
+        return result;
+    }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Semantic/Models/ProvideSemanticTokensRangeParams.cs → src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Semantic/Models/ProvideSemanticTokensRangesParams.cs

@@ -7,22 +7,22 @@
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer.Semantic.Models;
 
-internal class ProvideSemanticTokensRangeParams : SemanticTokensParams
+internal class ProvideSemanticTokensRangesParams : SemanticTokensParams
 {
     [DataMember(Name = "requiredHostDocumentVersion", IsRequired = true)]
     public long RequiredHostDocumentVersion { get; }
 
-    [DataMember(Name = "range", IsRequired = true)]
-    public Range Range { get; }
+    [DataMember(Name = "ranges", IsRequired = true)]
+    public Range[] Ranges { get; }
 
     [DataMember(Name = "correlationId", IsRequired = true)]
     public Guid CorrelationId { get; }
 
-    public ProvideSemanticTokensRangeParams(TextDocumentIdentifier textDocument, long requiredHostDocumentVersion, Range range, Guid correlationId)
+    public ProvideSemanticTokensRangesParams(TextDocumentIdentifier textDocument, long requiredHostDocumentVersion, Range[] ranges, Guid correlationId)
     {
         TextDocument = textDocument;
         RequiredHostDocumentVersion = requiredHostDocumentVersion;
-        Range = range;
+        Ranges = ranges;
         CorrelationId = correlationId;
     }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Semantic/Models/ProvideSemanticTokensResponse.cs

@@ -1,40 +1,20 @@
 // 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.Linq;
-
 namespace Microsoft.AspNetCore.Razor.LanguageServer.Semantic;
 
 /// <summary>
 /// Transports C# semantic token responses from the Razor LS client to the Razor LS.
 /// </summary>
 internal class ProvideSemanticTokensResponse
 {
-    public ProvideSemanticTokensResponse(int[]? tokens, long hostDocumentSyncVersion)
+    public ProvideSemanticTokensResponse(int[][]? tokens, long hostDocumentSyncVersion)
     {
         Tokens = tokens;
         HostDocumentSyncVersion = hostDocumentSyncVersion;
     }
 
-    public int[]? Tokens { get; }
+    public int[][]? Tokens { get; }
 
     public long HostDocumentSyncVersion { get; }
-
-    public override bool Equals(object? obj)
-    {
-        if (obj is not ProvideSemanticTokensResponse other ||
-            HostDocumentSyncVersion != other.HostDocumentSyncVersion)
-        {
-            return false;
-        }
-
-        return Tokens == other.Tokens ||
-              (Tokens is not null && other.Tokens is not null && Tokens.SequenceEqual(other.Tokens));
-    }
-
-    public override int GetHashCode()
-    {
-        throw new NotImplementedException();
-    }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Semantic/Services/RazorSemanticTokensInfoService.cs

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
@@ -13,6 +14,7 @@
 using Microsoft.AspNetCore.Razor.LanguageServer.Semantic.Models;
 using Microsoft.AspNetCore.Razor.PooledObjects;
 using Microsoft.CodeAnalysis.Razor;
+using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Razor.Workspaces.Extensions;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.Extensions.Logging;
@@ -25,6 +27,7 @@ internal class RazorSemanticTokensInfoService : IRazorSemanticTokensInfoService
     private const int TokenSize = 5;
 
     private readonly IRazorDocumentMappingService _documentMappingService;
+    private readonly LanguageServerFeatureOptions _languageServerFeatureOptions;
     private readonly RazorLSPOptionsMonitor _razorLSPOptionsMonitor;
     private readonly ClientNotifierServiceBase _languageServer;
     private readonly ILogger _logger;
@@ -33,11 +36,13 @@ public RazorSemanticTokensInfoService(
         ClientNotifierServiceBase languageServer,
         IRazorDocumentMappingService documentMappingService,
         RazorLSPOptionsMonitor razorLSPOptionsMonitor,
+        LanguageServerFeatureOptions languageServerFeatureOptions,
         ILoggerFactory loggerFactory)
     {
         _languageServer = languageServer ?? throw new ArgumentNullException(nameof(languageServer));
         _documentMappingService = documentMappingService ?? throw new ArgumentNullException(nameof(documentMappingService));
         _razorLSPOptionsMonitor = razorLSPOptionsMonitor ?? throw new ArgumentNullException(nameof(razorLSPOptionsMonitor));
+        _languageServerFeatureOptions = languageServerFeatureOptions ?? throw new ArgumentNullException(nameof(languageServerFeatureOptions));
 
         if (loggerFactory is null)
         {
@@ -111,21 +116,40 @@ private static List<SemanticRange> CombineSemanticRanges(List<SemanticRange> ran
         CancellationToken cancellationToken,
         string? previousResultId = null)
     {
-        // We'll try to call into the mapping service to map to the projected range for us. If that doesn't work,
-        // we'll try to find the minimal range ourselves.
         var generatedDocument = codeDocument.GetCSharpDocument();
-        if (!_documentMappingService.TryMapToGeneratedDocumentRange(generatedDocument, razorRange, out var csharpRange) &&
-            !TryGetMinimalCSharpRange(codeDocument, razorRange, out csharpRange))
+        Range[]? csharpRanges;
+
+        // When the feature flag is enabled we try to get a list of precise ranges for the C# code embedded in the Razor document.
+        // The feature flag allows to make calls to Roslyn using multiple smaller and disjoint ranges of the document
+        if (_languageServerFeatureOptions.UsePreciseSemanticTokenRanges)
+        {
+            if (!TryGetSortedCSharpRanges(codeDocument, razorRange, out csharpRanges))
+            {
+                // There's no C# in the range.
+                return new List<SemanticRange>();
+            }
+        }
+        else
         {
-            // There's no C# in the range.
-            return new List<SemanticRange>();
+            // When the feature flag is disabled, we fallback to computing a single range for the entire document.
+            // This single range is the minimal range that contains all of the C# code in the document.
+            // We'll try to call into the mapping service to map to the projected range for us. If that doesn't work,
+            // we'll try to find the minimal range ourselves.
+            if (!_documentMappingService.TryMapToGeneratedDocumentRange(generatedDocument, razorRange, out var csharpRange) &&
+                !TryGetMinimalCSharpRange(codeDocument, razorRange, out csharpRange))
+            {
+                // There's no C# in the range.
+                return new List<SemanticRange>();
+            }
+
+            csharpRanges = new Range[] { csharpRange };
         }
 
-        var csharpResponse = await GetMatchingCSharpResponseAsync(textDocumentIdentifier, documentVersion, csharpRange, correlationId, cancellationToken).ConfigureAwait(false);
+        var csharpResponse = await GetMatchingCSharpResponseAsync(textDocumentIdentifier, documentVersion, csharpRanges, correlationId, cancellationToken).ConfigureAwait(false);
 
         // Indicates an issue with retrieving the C# response (e.g. no response or C# is out of sync with us).
         // Unrecoverable, return default to indicate no change. We've already queued up a refresh request in
-        // `GetMatchingCSharpResponseAsync` that will cause us to retry in a bit.
+        // the server call that will cause us to retry in a bit.
         if (csharpResponse is null)
         {
             return null;
@@ -265,13 +289,13 @@ internal static bool TryGetMinimalCSharpRange(RazorCodeDocument codeDocument, Ra
     private async Task<int[]?> GetMatchingCSharpResponseAsync(
         TextDocumentIdentifier textDocumentIdentifier,
         long documentVersion,
-        Range csharpRange,
+        Range[] csharpRanges,
         Guid correlationId,
         CancellationToken cancellationToken)
     {
-        var parameter = new ProvideSemanticTokensRangeParams(textDocumentIdentifier, documentVersion, csharpRange, correlationId);
+        var parameter = new ProvideSemanticTokensRangesParams(textDocumentIdentifier, documentVersion, csharpRanges, correlationId);
 
-        var csharpResponse = await _languageServer.SendRequestAsync<ProvideSemanticTokensRangeParams, ProvideSemanticTokensResponse>(
+        var csharpResponse = await _languageServer.SendRequestAsync<ProvideSemanticTokensRangesParams, ProvideSemanticTokensResponse>(
             CustomMessageNames.RazorProvideSemanticTokensRangeEndpoint,
             parameter,
             cancellationToken).ConfigureAwait(false);
@@ -303,8 +327,113 @@ internal static bool TryGetMinimalCSharpRange(RazorCodeDocument codeDocument, Ra
             return null;
         }
 
-        var response = csharpResponse.Tokens ?? Array.Empty<int>();
-        return response;
+        return StitchSemanticTokenResponsesTogether(csharpResponse.Tokens);
+    }
+
+    // Internal for testing
+    internal static int[] StitchSemanticTokenResponsesTogether(int[][]? responseData)
+    {
+        // Each inner array in `responseData` represents a single C# document that is broken down into a list of tokens.
+        // This method stitches these lists of tokens together into a single, coherent list of semantic tokens.
+        // The resulting array is a flattened version of the input array, and is in the precise format expected by the Microsoft Language Server Protocol.
+        if (responseData is null || responseData.Length == 0)
+        {
+            return Array.Empty<int>();
+        }
+
+        if (responseData.Length == 1)
+        {
+            return responseData[0];
+        }
+
+        var count = responseData.Sum(r => r.Length);
+        var data = new int[count];
+        var dataIndex = 0;
+        var lastTokenLine = 0;
+
+        for (var i = 0; i < responseData.Length; i++)
+        {
+            var curData = responseData[i];
+
+            if (curData.Length == 0)
+            {
+                continue;
+            }
+
+            Array.Copy(curData, 0, data, dataIndex, curData.Length);
+            if (i != 0)
+            {
+                // The first two items in result.Data will potentially need it's line/col offset modified
+                var lineDelta = data[dataIndex] - lastTokenLine;
+                Debug.Assert(lineDelta >= 0);
+
+                // Update the first line copied over from curData
+                data[dataIndex] = lineDelta;
+
+                // Update the first column copied over from curData if on the same line as the previous token
+                if (lineDelta == 0)
+                {
+                    var lastTokenCol = 0;
+
+                    // Walk back accumulating column deltas until we find a start column (indicated by it's line offset being non-zero)
+                    for (var j = dataIndex - RazorSemanticTokensInfoService.TokenSize; j >= 0; j -= RazorSemanticTokensInfoService.TokenSize)
+                    {
+                        lastTokenCol += data[j + 1];
+                        if (data[j] != 0)
+                        {
+                            break;
+                        }
+                    }
+
+                    Debug.Assert(lastTokenCol >= 0);
+                    data[dataIndex + 1] -= lastTokenCol;
+                    Debug.Assert(data[dataIndex + 1] >= 0);
+                }
+            }
+
+            lastTokenLine = 0;
+            for (var j = 0; j < curData.Length; j += RazorSemanticTokensInfoService.TokenSize)
+            {
+                lastTokenLine += curData[j];
+            }
+
+            dataIndex += curData.Length;
+        }
+
+        return data;
+    }
+
+    // Internal for testing only
+    internal static bool TryGetSortedCSharpRanges(RazorCodeDocument codeDocument, Range razorRange, [NotNullWhen(true)] out Range[]? ranges)
+    {
+        using var _ = ListPool<Range>.GetPooledObject(out var csharpRanges);
+        var csharpSourceText = codeDocument.GetCSharpSourceText();
+        var sourceText = codeDocument.GetSourceText();
+        var textSpan = razorRange.AsTextSpan(sourceText);
+        var csharpDoc = codeDocument.GetCSharpDocument();
+
+        // We want to find the min and max C# source mapping that corresponds with our Razor range.
+        foreach (var mapping in csharpDoc.SourceMappings)
+        {
+            var mappedTextSpan = mapping.OriginalSpan.AsTextSpan();
+
+            if (textSpan.OverlapsWith(mappedTextSpan))
+            {
+                var mappedRange = mapping.GeneratedSpan.AsRange(csharpSourceText);
+                csharpRanges.Add(mappedRange);
+            }
+        }
+
+        if (csharpRanges.Count == 0)
+        {
+            ranges = null;
+            return false;
+        }
+
+        ranges = csharpRanges.ToArray();
+        // Ensure the C# ranges are sorted
+        Array.Sort(ranges, static (r1, r2) => r1.CompareTo(r2));
+        return true;
     }
 
     private static SemanticRange CSharpDataToSemanticRange(

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/LanguageServerFeatureOptions.cs

@@ -19,6 +19,8 @@ internal abstract class LanguageServerFeatureOptions
 
     public abstract bool DelegateToCSharpOnDiagnosticPublish { get; }
 
+    public abstract bool UsePreciseSemanticTokenRanges { get; }
+
     public abstract bool ShowAllCSharpCodeActions { get; }
 
     public abstract bool UpdateBuffersForClosedDocuments { get; }