Code review feedback - reduced allocations refactored logic

src/Features/LanguageServer/Protocol/Handler/SemanticTokens/SemanticTokensCache.cs

@@ -112,14 +112,9 @@ public SemanticTokensCache()
         /// </summary>
         public async Task<int[]?> GetCachedTokensDataAsync(
             Uri uri,
-            string? resultId,
+            string resultId,
             CancellationToken cancellationToken)
         {
-            if (resultId == null)
-            {
-                return null;
-            }
-
             using (await _semaphore.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
             {
                 if (!_tokens.TryGetValue(uri, out var tokenSets))

src/Features/LanguageServer/Protocol/Handler/SemanticTokens/SemanticTokensEditsHandler.cs

@@ -44,6 +44,7 @@ internal class SemanticTokensEditsHandler : AbstractRequestHandler<LSP.SemanticT
             CancellationToken cancellationToken)
         {
             Contract.ThrowIfNull(request.TextDocument);
+            Contract.ThrowIfNull(request.PreviousResultId);
 
             // Even though we want to ultimately pass edits back to LSP, we still need to compute all semantic tokens,
             // both for caching purposes and in order to have a baseline comparison when computing the edits.
@@ -117,10 +118,12 @@ private static SemanticTokensEdit[] ConvertToSemanticTokenEdits(SemanticToken[]
                 {
                     case EditKind.Insert:
                         indexToEditKinds.TryGetValue(edit.NewIndex, out var editKindWithoutInsert);
+                        Contract.ThrowIfTrue(editKindWithoutInsert == SemanticTokenEditKind.Insert); // We shouldn't have two inserts at the same position
                         indexToEditKinds[edit.NewIndex] = editKindWithoutInsert == default ? SemanticTokenEditKind.Insert : SemanticTokenEditKind.Update;
                         break;
                     case EditKind.Delete:
                         indexToEditKinds.TryGetValue(edit.OldIndex, out var editKindWithoutDelete);
+                        Contract.ThrowIfTrue(editKindWithoutDelete == SemanticTokenEditKind.Delete); // We shouldn't have two deletions at the same position
                         indexToEditKinds[edit.OldIndex] = editKindWithoutDelete == default ? SemanticTokenEditKind.Delete : SemanticTokenEditKind.Update;
                         break;
                 }
@@ -192,35 +195,27 @@ private static SemanticTokensEdit[] ConvertToSemanticTokenEdits(SemanticToken[]
                 Dictionary<int, SemanticTokenEditKind> indexToEditKinds,
                 ArrayBuilder<SemanticTokensEdit> semanticTokensEdits,
                 int[] editIndices,
-                int i,
+                int orderedEditNumber,
                 SemanticToken groupedSemanticToken,
                 int editStartPosition)
             {
-                var deleteCount = 5;
                 var _ = ArrayBuilder<int>.GetInstance(out var tokensToInsert);
-                tokensToInsert.AddRange(groupedSemanticToken.ConvertToArray());
+                groupedSemanticToken.AddToEnd(tokensToInsert);
 
                 // For simplicitly, we only allow an "update" (i.e. a dual insertion/deletion) to be
                 // combined with other updates.
+                var updatedOrderedEditNumber = GetUpdatedOrderedEditNumber(
+                    SemanticTokenEditKind.Update, indexToEditKinds, editIndices, orderedEditNumber);
 
-                // To continue combining edits, we need to ensure:
-                // 1) There is an edit following the current edit
-                // 2) The current edit and next edit involve tokens that are located right next to
-                // each other in the file.
-                // The two above criteria are also true for the similar loops in the local functions below,
-                // AddInsertionEdit and AddDeletionEdit.
-                while (i + 1 < editIndices.Length && indexToEditKinds[editIndices[i + 1]] == SemanticTokenEditKind.Update &&
-                    editIndices[i + 1] == editIndices[i] + 1)
+                var deleteCount = 5 + 5 * (updatedOrderedEditNumber - orderedEditNumber);
+                for (var i = 1; i <= updatedOrderedEditNumber - orderedEditNumber; i++)
                 {
-                    tokensToInsert.AddRange(newGroupedSemanticTokens[editIndices[i + 1]].ConvertToArray());
-                    deleteCount += 5;
-                    i++;
+                    newGroupedSemanticTokens[editIndices[orderedEditNumber + i]].AddToEnd(tokensToInsert);
                 }
 
                 semanticTokensEdits.Add(
                     GenerateEdit(start: editStartPosition, deleteCount: deleteCount, data: tokensToInsert.ToArray()));
-
-                return i;
+                return updatedOrderedEditNumber;
             }
 
             static int AddInsertionEdit(
@@ -228,46 +223,66 @@ private static SemanticTokensEdit[] ConvertToSemanticTokenEdits(SemanticToken[]
                 Dictionary<int, SemanticTokenEditKind> indexToEditKinds,
                 ArrayBuilder<SemanticTokensEdit> semanticTokensEdits,
                 int[] editIndices,
-                int i,
+                int orderedEditNumber,
                 SemanticToken groupedSemanticToken,
                 int editStartPosition)
             {
                 var _ = ArrayBuilder<int>.GetInstance(out var tokensToInsert);
-                tokensToInsert.AddRange(groupedSemanticToken.ConvertToArray());
+                groupedSemanticToken.AddToEnd(tokensToInsert);
 
                 // An insert can only be combined with other inserts that directly follow it.
-                while (i + 1 < editIndices.Length && indexToEditKinds[editIndices[i + 1]] == SemanticTokenEditKind.Insert &&
-                    editIndices[i + 1] == editIndices[i] + 1)
+                var updatedOrderedEditNumber = GetUpdatedOrderedEditNumber(
+                    SemanticTokenEditKind.Insert, indexToEditKinds, editIndices, orderedEditNumber);
+
+                for (var i = 1; i <= updatedOrderedEditNumber - orderedEditNumber; i++)
                 {
-                    tokensToInsert.AddRange(newGroupedSemanticTokens[editIndices[i + 1]].ConvertToArray());
-                    i++;
+                    newGroupedSemanticTokens[editIndices[orderedEditNumber + i]].AddToEnd(tokensToInsert);
                 }
 
                 semanticTokensEdits.Add(
                     GenerateEdit(start: editStartPosition, deleteCount: 0, data: tokensToInsert.ToArray()));
-                return i;
+                return updatedOrderedEditNumber;
             }
 
             static int AddDeletionEdit(
                 Dictionary<int, SemanticTokenEditKind> indexToEditKinds,
                 ArrayBuilder<SemanticTokensEdit> semanticTokensEdits,
                 int[] editIndices,
-                int i,
+                int orderedEditNumber,
                 int editStartPosition)
             {
-                var deleteCount = 5;
-
                 // A deletion can only be combined with other deletions that directly follow it.
-                while (i + 1 < editIndices.Length && indexToEditKinds[editIndices[i + 1]] == SemanticTokenEditKind.Delete &&
-                    editIndices[i + 1] == editIndices[i] + 1)
-                {
-                    deleteCount += 5;
-                    i++;
-                }
+                var updatedOrderedEditNumber = GetUpdatedOrderedEditNumber(
+                    SemanticTokenEditKind.Delete, indexToEditKinds, editIndices, orderedEditNumber);
 
+                var deleteCount = 5 + 5 * (updatedOrderedEditNumber - orderedEditNumber);
                 semanticTokensEdits.Add(
                     GenerateEdit(start: editStartPosition, deleteCount: deleteCount, data: Array.Empty<int>()));
-                return i;
+                return updatedOrderedEditNumber;
+            }
+
+            // Returns the updated ordered edit number after we know how many edits we can combine.
+            static int GetUpdatedOrderedEditNumber(
+                SemanticTokenEditKind editKind,
+                Dictionary<int, SemanticTokenEditKind> indexToEditKinds,
+                int[] editIndices,
+                int orderedEditNumber)
+            {
+                var originalOrderedEditNumber = orderedEditNumber;
+
+                // To continue combining edits, we need to ensure:
+                // 1) There is an edit following the current edit.
+                // 2) The current and next edits involve tokens that are located right next to
+                // each other in the file.
+                // 3) The next edit is the same type as the current edit.
+                while (orderedEditNumber + 1 < editIndices.Length &&
+                    indexToEditKinds[editIndices[orderedEditNumber + 1]] == editKind &&
+                    editIndices[orderedEditNumber + 1] == editIndices[orderedEditNumber] + 1)
+                {
+                    orderedEditNumber++;
+                }
+
+                return orderedEditNumber;
             }
         }
 
@@ -331,9 +346,13 @@ public SemanticToken(int deltaLine, int deltaStartCharacter, int length, int tok
                 _tokenModifiers = tokenModifiers;
             }
 
-            public int[] ConvertToArray()
+            public void AddToEnd(ArrayBuilder<int> tokensToInsert)
             {
-                return new int[] { _deltaLine, _deltaStartCharacter, _length, _tokenType, _tokenModifiers };
+                tokensToInsert.Add(_deltaLine);
+                tokensToInsert.Add(_deltaStartCharacter);
+                tokensToInsert.Add(_length);
+                tokensToInsert.Add(_tokenType);
+                tokensToInsert.Add(_tokenModifiers);
             }
 
             public bool Equals([AllowNull] SemanticToken otherToken)

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

@@ -4,6 +4,7 @@
 
 #nullable enable
 
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
@@ -122,21 +123,18 @@ internal class SemanticTokensHelpers
             ClassifiedSpan[] classifiedSpans,
             Dictionary<string, int> tokenTypesToIndex)
         {
-            // A TextSpan can be associated with multiple ClassifiedSpans (i.e. if  a token has
-            // modifiers). We perform this grouping since LSP requires that each token is
-            // reported together with all its modifiers.
-            using var _1 = PooledDictionary<TextSpan, List<string>>.GetInstance(out var textSpanToClassificationTypes);
-            GroupClassificationTypesByTextSpan(classifiedSpans, textSpanToClassificationTypes);
-
-            using var _2 = ArrayBuilder<int>.GetInstance(out var data);
+            using var _ = ArrayBuilder<int>.GetInstance(out var data);
 
+            // We keep track of the last line number and last start character since tokens are
+            // reported relative to each other.
             var lastLineNumber = 0;
             var lastStartCharacter = 0;
 
-            foreach (var textSpanToClassificationType in textSpanToClassificationTypes)
+            for (var currentClassifiedSpanIndex = 0; currentClassifiedSpanIndex < classifiedSpans.Length; currentClassifiedSpanIndex++)
             {
-                ComputeNextToken(lines, ref lastLineNumber, ref lastStartCharacter, textSpanToClassificationType.Key,
-                    textSpanToClassificationType.Value, tokenTypesToIndex,
+                currentClassifiedSpanIndex = ComputeNextToken(
+                    lines, ref lastLineNumber, ref lastStartCharacter, classifiedSpans,
+                    currentClassifiedSpanIndex, tokenTypesToIndex,
                     out var deltaLine, out var startCharacterDelta, out var tokenLength,
                     out var tokenType, out var tokenModifiers);
 
@@ -144,32 +142,14 @@ internal class SemanticTokensHelpers
             }
 
             return data.ToArray();
-
-            // Local functions
-            static void GroupClassificationTypesByTextSpan(
-                ClassifiedSpan[] classifiedSpans,
-                PooledDictionary<TextSpan, List<string>> textSpanToClassificationTypes)
-            {
-                foreach (var classifiedSpan in classifiedSpans)
-                {
-                    if (!textSpanToClassificationTypes.TryGetValue(classifiedSpan.TextSpan, out var classificationTypes))
-                    {
-                        textSpanToClassificationTypes.Add(classifiedSpan.TextSpan, new List<string> { classifiedSpan.ClassificationType });
-                    }
-                    else
-                    {
-                        classificationTypes.Add(classifiedSpan.ClassificationType);
-                    }
-                }
-            }
         }
 
-        private static void ComputeNextToken(
+        private static int ComputeNextToken(
             TextLineCollection lines,
             ref int lastLineNumber,
             ref int lastStartCharacter,
-            TextSpan textSpan,
-            List<string> classificationTypes,
+            ClassifiedSpan[] classifiedSpans,
+            int currentClassifiedSpanIndex,
             Dictionary<string, int> tokenTypesToIndex,
             // Out params
             out int deltaLineOut,
@@ -185,7 +165,8 @@ internal class SemanticTokensHelpers
             //     4. Token type (index) - looked up in SemanticTokensLegend.tokenTypes
             //     5. Token modifiers - each set bit will be looked up in SemanticTokensLegend.tokenModifiers
 
-            var linePosition = lines.GetLinePositionSpan(textSpan).Start;
+            var classifiedSpan = classifiedSpans[currentClassifiedSpanIndex];
+            var linePosition = lines.GetLinePositionSpan(classifiedSpan.TextSpan).Start;
             var lineNumber = linePosition.Line;
             var startCharacter = linePosition.Character;
 
@@ -202,14 +183,18 @@ internal class SemanticTokensHelpers
             }
 
             // 3. Token length
-            var tokenLength = textSpan.Length;
+            var tokenLength = classifiedSpan.TextSpan.Length;
 
             // We currently only have one modifier (static). The logic below will need to change in the future if other
             // modifiers are added in the future.
             var modifierBits = TokenModifiers.None;
             var tokenTypeIndex = 0;
-            foreach (var classificationType in classificationTypes)
+            var originalTextSpan = classifiedSpan.TextSpan;
+
+            // Classified spans with the same text span should be combined into one token.
+            while (classifiedSpans[currentClassifiedSpanIndex].TextSpan == originalTextSpan)
             {
+                var classificationType = classifiedSpans[currentClassifiedSpanIndex].ClassificationType;
                 if (classificationType != ClassificationTypeNames.StaticSymbol)
                 {
                     // 4. Token type - looked up in SemanticTokensLegend.tokenTypes (language server defined mapping
@@ -221,6 +206,13 @@ internal class SemanticTokensHelpers
                     // 5. Token modifiers - each set bit will be looked up in SemanticTokensLegend.tokenModifiers
                     modifierBits = TokenModifiers.Static;
                 }
+
+                if (currentClassifiedSpanIndex + 1 >= classifiedSpans.Length || classifiedSpans[currentClassifiedSpanIndex + 1].TextSpan != originalTextSpan)
+                {
+                    break;
+                }
+
+                currentClassifiedSpanIndex++;
             }
 
             lastLineNumber = lineNumber;
@@ -231,6 +223,8 @@ internal class SemanticTokensHelpers
             tokenLengthOut = tokenLength;
             tokenTypeOut = tokenTypeIndex;
             tokenModifiersOut = (int)modifierBits;
+
+            return currentClassifiedSpanIndex;
         }
 
         private static int GetTokenTypeIndex(string classificationType, Dictionary<string, int> tokenTypesToIndex)

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

@@ -28,6 +28,9 @@ public static class Classifier
             return GetClassifiedSpans(semanticModel, textSpan, document.Project.Solution.Workspace, cancellationToken);
         }
 
+        /// <summary>
+        /// Returns classified spans in order.
+        /// </summary>
         public static IEnumerable<ClassifiedSpan> GetClassifiedSpans(
             SemanticModel semanticModel,
             TextSpan textSpan,