Further reduce allocations in SemanticToken calculations

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

@@ -4,13 +4,14 @@
 #nullable disable
 
 using System;
-using System.Collections.Immutable;
+using System.Collections.Generic;
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
 using BenchmarkDotNet.Attributes;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.LanguageServer.Semantic;
+using Microsoft.AspNetCore.Razor.Threading;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
@@ -101,16 +102,17 @@ public TestRazorSemanticTokensInfoService(
         {
         }
 
-        // We can't get C# responses without significant amounts of extra work, so let's just shim it for now, any non-Null result is fine.
-        protected override Task<ImmutableArray<SemanticRange>?> GetCSharpSemanticRangesAsync(
+        // We can't get C# responses without significant amounts of extra work, so let's just shim it for now and not append any ranges.
+        protected override Task<bool> AddCSharpSemanticRangesAsync(
+            List<SemanticRange> ranges,
             DocumentContext documentContext,
             RazorCodeDocument codeDocument,
             LinePositionSpan razorRange,
             bool colorBackground,
             Guid correlationId,
             CancellationToken cancellationToken)
         {
-            return Task.FromResult<ImmutableArray<SemanticRange>?>(ImmutableArray<SemanticRange>.Empty);
+            return SpecializedTasks.True;
         }
     }
 }

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

@@ -4,6 +4,7 @@
 #nullable disable
 
 using System;
+using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.IO;
 using System.Threading;
@@ -13,6 +14,7 @@
 using Microsoft.AspNetCore.Razor.LanguageServer;
 using Microsoft.AspNetCore.Razor.LanguageServer.EndpointContracts;
 using Microsoft.AspNetCore.Razor.LanguageServer.Semantic;
+using Microsoft.AspNetCore.Razor.Threading;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
@@ -137,16 +139,17 @@ public TestCustomizableRazorSemanticTokensInfoService(
         {
         }
 
-        // We can't get C# responses without significant amounts of extra work, so let's just shim it for now, any non-Null result is fine.
-        protected override Task<ImmutableArray<SemanticRange>?> GetCSharpSemanticRangesAsync(
+        // We can't get C# responses without significant amounts of extra work, so let's just shim it for now and not append any ranges.
+        protected override Task<bool> AddCSharpSemanticRangesAsync(
+            List<SemanticRange> ranges,
             DocumentContext documentContext,
             RazorCodeDocument codeDocument,
             LinePositionSpan razorSpan,
             bool colorBackground,
             Guid correlationId,
             CancellationToken cancellationToken)
         {
-            return Task.FromResult<ImmutableArray<SemanticRange>?>(PregeneratedRandomSemanticRanges);
+            return SpecializedTasks.True;
         }
     }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/SemanticTokens/AbstractRazorSemanticTokensInfoService.Policy.cs

@@ -0,0 +1,40 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Collections.Generic;
+using Microsoft.Extensions.ObjectPool;
+
+namespace Microsoft.CodeAnalysis.Razor.SemanticTokens;
+
+internal abstract partial class AbstractRazorSemanticTokensInfoService
+{
+    private sealed class Policy : IPooledObjectPolicy<List<SemanticRange>>
+    {
+        public static readonly Policy Instance = new();
+
+        // Significantly larger than DefaultPool.MaximumObjectSize as these arrays are commonly large.
+        // The 2048 limit should be large enough for nearly all semantic token requests, while still
+        // keeping the backing arrays off the LOH.
+        public const int MaximumObjectSize = 2048;
+
+        private Policy()
+        {
+        }
+
+        public List<SemanticRange> Create() => [];
+
+        public bool Return(List<SemanticRange> list)
+        {
+            var count = list.Count;
+
+            list.Clear();
+
+            if (count > MaximumObjectSize)
+            {
+                list.TrimExcess();
+            }
+
+            return true;
+        }
+    }
+}

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/SemanticTokens/AbstractRazorSemanticTokensInfoService.cs

@@ -15,10 +15,11 @@
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Text;
+using Microsoft.Extensions.ObjectPool;
 
 namespace Microsoft.CodeAnalysis.Razor.SemanticTokens;
 
-internal abstract class AbstractRazorSemanticTokensInfoService(
+internal abstract partial class AbstractRazorSemanticTokensInfoService(
     IDocumentMappingService documentMappingService,
     ISemanticTokensLegendService semanticTokensLegendService,
     ICSharpSemanticTokensProvider csharpSemanticTokensProvider,
@@ -27,6 +28,9 @@ internal abstract class AbstractRazorSemanticTokensInfoService(
 {
     private const int TokenSize = 5;
 
+    // Use a custom pool as these lists commonly exceed the size threshold for returning into the default ListPool.
+    private static readonly ObjectPool<List<SemanticRange>> s_pool = DefaultPool.Create(Policy.Instance, size: 8);
+
     private readonly IDocumentMappingService _documentMappingService = documentMappingService;
     private readonly ISemanticTokensLegendService _semanticTokensLegendService = semanticTokensLegendService;
     private readonly ICSharpSemanticTokensProvider _csharpSemanticTokensProvider = csharpSemanticTokensProvider;
@@ -66,12 +70,16 @@ internal abstract class AbstractRazorSemanticTokensInfoService(
         cancellationToken.ThrowIfCancellationRequested();
 
         var textSpan = codeDocument.Source.Text.GetTextSpan(span);
-        var razorSemanticRanges = SemanticTokensVisitor.GetSemanticRanges(codeDocument, textSpan, _semanticTokensLegendService, colorBackground);
-        ImmutableArray<SemanticRange>? csharpSemanticRangesResult = null;
+        using var _ = s_pool.GetPooledObject(out var combinedSemanticRanges);
+
+        SemanticTokensVisitor.AddSemanticRanges(combinedSemanticRanges, codeDocument, textSpan, _semanticTokensLegendService, colorBackground);
+        Debug.Assert(combinedSemanticRanges.SequenceEqual(combinedSemanticRanges.OrderBy(g => g)));
+
+        var successfullyRetrievedCSharpSemanticRanges = false;
 
         try
         {
-            csharpSemanticRangesResult = await GetCSharpSemanticRangesAsync(documentContext, codeDocument, span, colorBackground, correlationId, cancellationToken).ConfigureAwait(false);
+            successfullyRetrievedCSharpSemanticRanges = await AddCSharpSemanticRangesAsync(combinedSemanticRanges, documentContext, codeDocument, span, colorBackground, correlationId, cancellationToken).ConfigureAwait(false);
         }
         catch (OperationCanceledException)
         {
@@ -84,53 +92,24 @@ internal abstract class AbstractRazorSemanticTokensInfoService(
 
         // Didn't get any C# tokens, likely because the user kept typing and a future semantic tokens request will occur.
         // We return null (which to the LSP is a no-op) to prevent flashing of CSharp elements.
-        if (csharpSemanticRangesResult is not { } csharpSemanticRanges)
+        if (!successfullyRetrievedCSharpSemanticRanges)
         {
             _logger.LogDebug($"Couldn't get C# tokens for version {documentContext.Snapshot.Version} of {documentContext.Uri}. Returning null");
             return null;
         }
 
-        var combinedSemanticRanges = CombineSemanticRanges(razorSemanticRanges, csharpSemanticRanges);
-
-        return ConvertSemanticRangesToSemanticTokensData(combinedSemanticRanges, codeDocument);
-    }
-
-    private static ImmutableArray<SemanticRange> CombineSemanticRanges(ImmutableArray<SemanticRange> razorRanges, ImmutableArray<SemanticRange> csharpRanges)
-    {
-        Debug.Assert(razorRanges.SequenceEqual(razorRanges.OrderBy(g => g)));
-
-        // If there are no C# in what we're trying to classify we don't need to do anything special since we know the razor ranges will be sorted
-        // because we use a visitor to create them, and the above Assert will validate it in our tests.
-        if (csharpRanges.Length == 0)
-        {
-            return razorRanges;
-        }
-
-        // If there are no Razor ranges then we can't just return the C# ranges, as they ranges are not necessarily sorted. They would have been
-        // in order when the C# server gave them to us, but the data we have here is after re-mapping to the Razor document, which can result in
-        // things being moved around. We need to sort before we return them.
-        if (razorRanges.Length == 0)
-        {
-            return csharpRanges.Sort();
-        }
-
         // If we have both types of tokens then we need to sort them all together, even though we know the Razor ranges will be sorted already,
         // because they can arbitrarily interleave. The SemanticRange.CompareTo method also has some logic to ensure that if Razor and C# ranges
         // are equivalent, the Razor range will be ordered first, so we can later drop the C# range, and prefer our classification over C#s.
         // Additionally, as mentioned above, the C# ranges are not guaranteed to be in order
-        using var _ = ArrayBuilderPool<SemanticRange>.GetPooledObject(out var newList);
-        newList.SetCapacityIfLarger(razorRanges.Length + csharpRanges.Length);
+        combinedSemanticRanges.Sort();
 
-        newList.AddRange(razorRanges);
-        newList.AddRange(csharpRanges);
-
-        newList.Sort();
-
-        return newList.ToImmutableAndClear();
+        return ConvertSemanticRangesToSemanticTokensData(combinedSemanticRanges, codeDocument);
     }
 
     // Virtual for benchmarks
-    protected virtual async Task<ImmutableArray<SemanticRange>?> GetCSharpSemanticRangesAsync(
+    protected virtual async Task<bool> AddCSharpSemanticRangesAsync(
+        List<SemanticRange> ranges,
         DocumentContext documentContext,
         RazorCodeDocument codeDocument,
         LinePositionSpan razorSpan,
@@ -144,7 +123,7 @@ private static ImmutableArray<SemanticRange> CombineSemanticRanges(ImmutableArra
         if (!TryGetSortedCSharpRanges(codeDocument, razorSpan, out var csharpRanges))
         {
             // There's no C# in the range.
-            return ImmutableArray<SemanticRange>.Empty;
+            return true;
         }
 
         _logger.LogDebug($"Requesting C# semantic tokens for host version {documentContext.Snapshot.Version}, correlation ID {correlationId}, and the server thinks there are {codeDocument.GetCSharpSourceText().Lines.Count} lines of C#");
@@ -156,11 +135,10 @@ private static ImmutableArray<SemanticRange> CombineSemanticRanges(ImmutableArra
         // the server call that will cause us to retry in a bit.
         if (csharpResponse is null)
         {
-            return null;
+            return false;
         }
 
-        using var _ = ArrayBuilderPool<SemanticRange>.GetPooledObject(out var razorRanges);
-        razorRanges.SetCapacityIfLarger(csharpResponse.Length / TokenSize);
+        ranges.SetCapacityIfLarger(csharpResponse.Length / TokenSize);
 
         var textClassification = _semanticTokensLegendService.TokenTypes.MarkupTextLiteral;
         var razorSource = codeDocument.Source.Text;
@@ -184,10 +162,10 @@ private static ImmutableArray<SemanticRange> CombineSemanticRanges(ImmutableArra
                     if (colorBackground)
                     {
                         tokenModifiers |= _semanticTokensLegendService.TokenModifiers.RazorCodeModifier;
-                        AddAdditionalCSharpWhitespaceRanges(razorRanges, textClassification, razorSource, previousRazorSemanticRange, originalRange);
+                        AddAdditionalCSharpWhitespaceRanges(ranges, textClassification, razorSource, previousRazorSemanticRange, originalRange);
                     }
 
-                    razorRanges.Add(new SemanticRange(semanticRange.Kind, originalRange.Start.Line, originalRange.Start.Character, originalRange.End.Line, originalRange.End.Character, tokenModifiers, fromRazor: false));
+                    ranges.Add(new SemanticRange(semanticRange.Kind, originalRange.Start.Line, originalRange.Start.Character, originalRange.End.Line, originalRange.End.Character, tokenModifiers, fromRazor: false));
                 }
 
                 previousRazorSemanticRange = originalRange;
@@ -196,10 +174,10 @@ private static ImmutableArray<SemanticRange> CombineSemanticRanges(ImmutableArra
             previousSemanticRange = semanticRange;
         }
 
-        return razorRanges.ToImmutableAndClear();
+        return true;
     }
 
-    private void AddAdditionalCSharpWhitespaceRanges(ImmutableArray<SemanticRange>.Builder razorRanges, int textClassification, SourceText razorSource, LinePositionSpan? previousRazorSemanticRange, LinePositionSpan originalRange)
+    private void AddAdditionalCSharpWhitespaceRanges(List<SemanticRange> razorRanges, int textClassification, SourceText razorSource, LinePositionSpan? previousRazorSemanticRange, LinePositionSpan originalRange)
     {
         var startLine = originalRange.Start.Line;
         var startChar = originalRange.Start.Character;
@@ -301,38 +279,49 @@ private static SemanticRange CSharpDataToSemanticRange(
     }
 
     private static int[] ConvertSemanticRangesToSemanticTokensData(
-        ImmutableArray<SemanticRange> semanticRanges,
+        List<SemanticRange> semanticRanges,
         RazorCodeDocument razorCodeDocument)
     {
-        SemanticRange previousResult = default;
-
         var sourceText = razorCodeDocument.Source.Text;
 
         // We don't bother filtering out duplicate ranges (eg, where C# and Razor both have opinions), but instead take advantage of
         // our sort algorithm to be correct, so we can skip duplicates here. That means our final array may end up smaller than the
-        // expected size, so we have to use a list to build it.
-        using var _ = ListPool<int>.GetPooledObject(out var data);
-        data.SetCapacityIfLarger(semanticRanges.Length * TokenSize);
+        // expected size.
+        var tokens = new int[semanticRanges.Count * TokenSize];
 
-        var firstRange = true;
-        foreach (var result in semanticRanges)
+        var isFirstRange = true;
+        var index = 0;
+        SemanticRange previousRange = default;
+        foreach (var range in semanticRanges)
         {
-            AppendData(result, previousResult, firstRange, sourceText, data);
-            firstRange = false;
+            if (TryWriteToken(range, previousRange, isFirstRange, sourceText, tokens.AsSpan(index, TokenSize)))
+            {
+                index += TokenSize;
+            }
 
-            previousResult = result;
+            isFirstRange = false;
+            previousRange = range;
         }
 
-        return [.. data];
+        // The common case is that the ConvertIntoDataArray calls didn't find any overlap, and we can just directly use the
+        // data array we allocated. If there was overlap, then we need to allocate a smaller array and copy the data over.
+        if (index < tokens.Length)
+        {
+            Array.Resize(ref tokens, newSize: index);
+        }
 
-        // We purposely capture and manipulate the "data" array here to avoid allocation
-        static void AppendData(
+        return tokens;
+
+        // We purposely capture and manipulate the destination array here to avoid allocation
+        static bool TryWriteToken(
             SemanticRange currentRange,
             SemanticRange previousRange,
-            bool firstRange,
+            bool isFirstRange,
             SourceText sourceText,
-            List<int> data)
+            Span<int> destination)
         {
+            Debug.Assert(destination.Length == TokenSize);
+
             /*
              * In short, each token takes 5 integers to represent, so a specific token `i` in the file consists of the following array indices:
              *  - at index `5*i`   - `deltaLine`: token line number, relative to the previous token
@@ -347,24 +336,24 @@ static void AppendData(
             var deltaLine = currentRange.StartLine - previousLineIndex;
 
             int deltaStart;
-            if (!firstRange && previousRange.StartLine == currentRange.StartLine)
+            if (!isFirstRange && previousRange.StartLine == currentRange.StartLine)
             {
                 deltaStart = currentRange.StartCharacter - previousRange.StartCharacter;
 
                 // If there is no line delta, no char delta, and this isn't the first range
                 // then it means this range overlaps the previous, so we skip it.
                 if (deltaStart == 0)
                 {
-                    return;
+                    return false;
                 }
             }
             else
             {
                 deltaStart = currentRange.StartCharacter;
             }
 
-            data.Add(deltaLine);
-            data.Add(deltaStart);
+            destination[0] = deltaLine;
+            destination[1] = deltaStart;
 
             // length
 
@@ -376,13 +365,15 @@ static void AppendData(
 
             var length = endPosition - startPosition;
             Debug.Assert(length > 0);
-            data.Add(length);
+            destination[2] = length;
 
             // tokenType
-            data.Add(currentRange.Kind);
+            destination[3] = currentRange.Kind;
 
             // tokenModifiers
-            data.Add(currentRange.Modifier);
+            destination[4] = currentRange.Modifier;
+
+            return true;
         }
     }
 }