Don't use a pool in SemanticTokensHelper.ComputeTokens

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

@@ -12,7 +12,6 @@
 using Microsoft.CodeAnalysis.Collections;
 using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Options;
-using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Text;
 using LSP = Roslyn.LanguageServer.Protocol;
@@ -21,8 +20,6 @@ namespace Microsoft.CodeAnalysis.LanguageServer.Handler.SemanticTokens;
 
 internal static class SemanticTokensHelpers
 {
-    private static readonly ObjectPool<List<int>> s_tokenListPool = new(() => new List<int>(capacity: 1000));
-
     /// <param name="ranges">The ranges to get semantic tokens for.  If <c>null</c> then the entire document will be
     /// processed.</param>
     internal static async Task<int[]> HandleRequestHelperAsync(
@@ -241,13 +238,9 @@ private static int[] ComputeTokens(
         var lastStartCharacter = 0;
 
         var tokenTypeMap = SemanticTokensSchema.GetSchema(supportsVisualStudioExtensions).TokenTypeMap;
+        var data = AllocateTokenArray(classifiedSpans);
 
-        using var pooledData = s_tokenListPool.GetPooledObject();
-        var data = pooledData.Object;
-
-        // Items in the pool may not have been cleared
-        data.Clear();
-
+        var i = 0;
         for (var currentClassifiedSpanIndex = 0; currentClassifiedSpanIndex < classifiedSpans.Count; currentClassifiedSpanIndex++)
         {
             currentClassifiedSpanIndex = ComputeNextToken(
@@ -256,14 +249,40 @@ private static int[] ComputeTokens(
                 out var deltaLine, out var startCharacterDelta, out var tokenLength,
                 out var tokenType, out var tokenModifiers);
 
-            data.Add(deltaLine);
-            data.Add(startCharacterDelta);
-            data.Add(tokenLength);
-            data.Add(tokenType);
-            data.Add(tokenModifiers);
+            data[i++] = deltaLine;
+            data[i++] = startCharacterDelta;
+            data[i++] = tokenLength;
+            data[i++] = tokenType;
+            data[i++] = tokenModifiers;
+        }
+
+        Contract.ThrowIfFalse(i == data.Length, "The number of computed tokens does not match the expected size.");
+
+        return data;
+    }
+
+    // This method allocates an array of integers to hold the semantic tokens data.
+    // NOTE: The number of items in the array is based on the number of unique classified spans
+    // in the provided list and is closely tied with how ComputeNextToken's loop works
+    private static int[] AllocateTokenArray(SegmentedList<ClassifiedSpan> classifiedSpans)
+    {
+        if (classifiedSpans.Count == 0)
+            return Array.Empty<int>();
+
+        var uniqueSpanCount = 1;
+        var lastSpan = classifiedSpans[0].TextSpan;
+
+        for (var index = 1; index < classifiedSpans.Count; index++)
+        {
+            var currentSpan = classifiedSpans[index].TextSpan;
+            if (currentSpan != lastSpan)
+            {
+                uniqueSpanCount++;
+                lastSpan = currentSpan;
+            }
         }
 
-        return [.. data];
+        return new int[5 * uniqueSpanCount];
     }
 
     private static int ComputeNextToken(
@@ -315,6 +334,8 @@ private static int ComputeNextToken(
         var tokenTypeIndex = 0;
 
         // Classified spans with the same text span should be combined into one token.
+        // NOTE: The update of currentClassifiedSpanIndex is closely tied to the allocation
+        // of the data array in AllocateTokenArray.
         while (classifiedSpans[currentClassifiedSpanIndex].TextSpan == originalTextSpan)
         {
             var classificationType = classifiedSpans[currentClassifiedSpanIndex].ClassificationType;