Reduce allocations obtaining classified spans in ClassifierHelper (#79856)

ClassifiedSpan[] shows up as 0.7% of total allocations (in LOH) in the completion scenarios in the razor cohosting speedometer test.

Surrounding code already uses a custom pool for ClassifiedSpan[] via Classifier.GetPooledList. This PR changes ClassifierHelper.GetClassifiedSpansAsync to take in a SegmentedList<ClassifiedSpan> to populate rather than returning an ImmutableArray<ClassifiedSpan>.

Author: ToddGrun

src/LanguageServer/Protocol/Handler/SemanticTokens/SemanticTokensHelpers.cs

@@ -91,7 +91,7 @@ public static async Task<int[]> ComputeSemanticTokensDataAsync(
         using var _2 = Classifier.GetPooledList(out var updatedClassifiedSpans);
 
         var textSpans = spans.SelectAsArray(static (span, text) => text.Lines.GetTextSpan(span), text);
-        await GetClassifiedSpansForDocumentAsync(
+        await AddClassifiedSpansForDocumentAsync(
             classifiedSpans, document, textSpans, options, cancellationToken).ConfigureAwait(false);
 
         // Multi-line tokens are not supported by VS (tracked by https://devdiv.visualstudio.com/DevDiv/_workitems/edit/1265495).
@@ -108,7 +108,7 @@ await GetClassifiedSpansForDocumentAsync(
         return ComputeTokens(text.Lines, updatedClassifiedSpans, supportsVisualStudioExtensions, tokenTypesToIndex);
     }
 
-    private static async Task GetClassifiedSpansForDocumentAsync(
+    private static async Task AddClassifiedSpansForDocumentAsync(
         SegmentedList<ClassifiedSpan> classifiedSpans,
         Document document,
         ImmutableArray<TextSpan> textSpans,
@@ -121,13 +121,11 @@ private static async Task GetClassifiedSpansForDocumentAsync(
         // then the semantic token classifications will override them.
 
         // `includeAdditiveSpans` will add token modifiers such as 'static', which we want to include in LSP.
-        var spans = await ClassifierHelper.GetClassifiedSpansAsync(
-            document, textSpans, options, includeAdditiveSpans: true, cancellationToken).ConfigureAwait(false);
+        await ClassifierHelper.AddClassifiedSpansAsync(
+            classifiedSpans, document, textSpans, options, includeAdditiveSpans: true, cancellationToken).ConfigureAwait(false);
 
         // Some classified spans may not be relevant and should be filtered out before we convert to tokens.
-        var filteredSpans = spans.Where(s => !ShouldFilterClassification(s));
-
-        classifiedSpans.AddRange(filteredSpans);
+        classifiedSpans.RemoveAll(static s => ShouldFilterClassification(s));
     }
 
     private static bool ShouldFilterClassification(ClassifiedSpan s)

src/Workspaces/Core/Portable/Classification/ClassifierHelper.cs

@@ -33,8 +33,12 @@ public static async Task<ImmutableArray<ClassifiedSpan>> GetClassifiedSpansAsync
         bool includeAdditiveSpans,
         CancellationToken cancellationToken)
     {
-        return await GetClassifiedSpansAsync(document, [span], options, includeAdditiveSpans, cancellationToken)
+        using var _ = Classifier.GetPooledList(out var classifiedSpans);
+
+        await AddClassifiedSpansAsync(classifiedSpans, document, [span], options, includeAdditiveSpans, cancellationToken)
             .ConfigureAwait(false);
+
+        return [.. classifiedSpans];
     }
 
     /// <summary>
@@ -49,7 +53,8 @@ public static async Task<ImmutableArray<ClassifiedSpan>> GetClassifiedSpansAsync
     /// results or not.  'Additive' spans are things like 'this variable is static' or 'this variable is
     /// overwritten'.  i.e. they add additional information to a previous classification.</param>
     /// <param name="cancellationToken">A cancellation token.</param>
-    public static async Task<ImmutableArray<ClassifiedSpan>> GetClassifiedSpansAsync(
+    public static async Task AddClassifiedSpansAsync(
+        SegmentedList<ClassifiedSpan> classifiedSpans,
         Document document,
         ImmutableArray<TextSpan> spans,
         ClassificationOptions options,
@@ -58,7 +63,7 @@ public static async Task<ImmutableArray<ClassifiedSpan>> GetClassifiedSpansAsync
     {
         var classificationService = document.GetLanguageService<IClassificationService>();
         if (classificationService == null)
-            return default;
+            return;
 
         // Call out to the individual language to classify the chunk of text around the
         // reference. We'll get both the syntactic and semantic spans for this region.
@@ -90,8 +95,7 @@ public static async Task<ImmutableArray<ClassifiedSpan>> GetClassifiedSpansAsync
         }
 
         var widenedSpan = new TextSpan(spans[0].Start, spans[^1].End);
-        var classifiedSpans = MergeClassifiedSpans(syntaxSpans, semanticSpans, widenedSpan);
-        return classifiedSpans;
+        MergeClassifiedSpans(syntaxSpans, semanticSpans, widenedSpan, classifiedSpans);
     }
 
     private static void RemoveAdditiveSpans(SegmentedList<ClassifiedSpan> spans)
@@ -104,10 +108,11 @@ private static void RemoveAdditiveSpans(SegmentedList<ClassifiedSpan> spans)
         }
     }
 
-    private static ImmutableArray<ClassifiedSpan> MergeClassifiedSpans(
+    private static void MergeClassifiedSpans(
         SegmentedList<ClassifiedSpan> syntaxSpans,
         SegmentedList<ClassifiedSpan> semanticSpans,
-        TextSpan widenedSpan)
+        TextSpan widenedSpan,
+        SegmentedList<ClassifiedSpan> classifiedSpans)
     {
         // The spans produced by the language services may not be ordered
         // (indeed, this happens with semantic classification as different
@@ -130,16 +135,14 @@ private static ImmutableArray<ClassifiedSpan> MergeClassifiedSpans(
         AdjustSpans(syntaxSpans, widenedSpan);
         AdjustSpans(semanticSpans, widenedSpan);
 
-        using var _1 = Classifier.GetPooledList(out var mergedSpans);
+        using var _ = Classifier.GetPooledList(out var mergedSpans);
 
         MergeParts(syntaxSpans, semanticSpans, mergedSpans);
         Order(mergedSpans);
 
         // The classification service will only produce classifications for things it knows about.  i.e. there will
         // be gaps in what it produces. Fill in those gaps so we have *all* parts of the span classified properly.
-        using var _2 = Classifier.GetPooledList(out var filledInSpans);
-        FillInClassifiedSpanGaps(widenedSpan.Start, mergedSpans, filledInSpans);
-        return [.. filledInSpans];
+        FillInClassifiedSpanGaps(widenedSpan.Start, mergedSpans, classifiedSpans);
     }
 
     private static readonly Comparison<ClassifiedSpan> s_spanComparison = static (s1, s2) => s1.TextSpan.Start - s2.TextSpan.Start;