Merge pull request #30821 from sharwell/optimize-tokenstream

Remove TokenStream._tokenToIndexMap and replace with binary search

src/Workspaces/Core/Portable/Formatting/Engine/TokenData.cs

@@ -17,9 +17,9 @@ namespace Microsoft.CodeAnalysis.Formatting
     /// this object is supposed to be live very short but created a lot of time. that is why it is struct. 
     /// (same reason why SyntaxToken is struct - to reduce heap allocation)
     /// </summary>
-    internal struct TokenData : IEqualityComparer<TokenData>, IEquatable<TokenData>, IComparable<TokenData>, IComparer<TokenData>
+    internal readonly struct TokenData : IEqualityComparer<TokenData>, IEquatable<TokenData>, IComparable<TokenData>, IComparer<TokenData>
     {
-        public TokenData(TokenStream tokenStream, int indexInStream, SyntaxToken token) : this()
+        public TokenData(TokenStream tokenStream, int indexInStream, SyntaxToken token)
         {
             Contract.ThrowIfNull(tokenStream);
             Contract.ThrowIfFalse((indexInStream == -1) || (0 <= indexInStream && indexInStream < tokenStream.TokenCount));

src/Workspaces/Core/Portable/Formatting/Engine/TokenStream.cs

@@ -25,7 +25,6 @@ internal partial class TokenStream
 
         // caches token information within given formatting span to improve perf
         private readonly List<SyntaxToken> _tokens;
-        private readonly Dictionary<SyntaxToken, int> _tokenToIndexMap;
 
         // caches original trivia info to improve perf
         private readonly TriviaData[] _cachedOriginalTriviaInfo;
@@ -65,12 +64,6 @@ public TokenStream(TreeData treeData, OptionSet optionSet, TextSpan spanToFormat
                 // initialize trivia related info
                 _cachedOriginalTriviaInfo = new TriviaData[this.TokenCount - 1];
 
-                _tokenToIndexMap = new Dictionary<SyntaxToken, int>(this.TokenCount);
-                for (int i = 0; i < this.TokenCount; i++)
-                {
-                    _tokenToIndexMap.Add(_tokens[i], i);
-                }
-
                 // Func Cache
                 _getTriviaData = this.GetTriviaData;
                 _getOriginalTriviaData = this.GetOriginalTriviaData;
@@ -510,9 +503,51 @@ private bool IsFirstTokenOnLine(TokenData tokenData1, TokenData tokenData2)
 
         private int GetTokenIndexInStream(SyntaxToken token)
         {
-            if (_tokenToIndexMap.TryGetValue(token, out var value))
+            var tokenIndex = _tokens.BinarySearch(token, TokenOrderComparer.Instance);
+            if (tokenIndex < 0)
+            {
+                return -1;
+            }
+
+            // Source characters cannot be assigned to multiple tokens. If the token has non-zero width, it will be an
+            // exact match for at most one token.
+            if (!token.FullSpan.IsEmpty)
             {
-                return value;
+                return _tokens[tokenIndex] == token ? tokenIndex : -1;
+            }
+
+            // Multiple tokens can have empty spans. The binary search operation will return one of them; look forward
+            // and then backward to locate the desired token within the set of one or more zero-width tokens located at
+            // the same position.
+            Debug.Assert(token.FullSpan.IsEmpty);
+            for (var i = tokenIndex; i < _tokens.Count; i++)
+            {
+                if (!_tokens[i].FullSpan.IsEmpty)
+                {
+                    // Current token can't match because the span is different, and future tokens won't match because
+                    // they are lexicographically after the token we are interested in.
+                    break;
+                }
+
+                if (_tokens[i] == token)
+                {
+                    return i;
+                }
+            }
+
+            for (var i = tokenIndex - 1; i >= 0; i--)
+            {
+                if (!_tokens[i].FullSpan.IsEmpty)
+                {
+                    // Current token can't match because the span is different, and future tokens won't match because
+                    // they are lexicographically before the token we are interested in.
+                    break;
+                }
+
+                if (_tokens[i] == token)
+                {
+                    return i;
+                }
             }
 
             return -1;
@@ -525,5 +560,17 @@ private int GetTokenIndexInStream(SyntaxToken token)
                 return new Iterator(_tokens);
             }
         }
+
+        private sealed class TokenOrderComparer : IComparer<SyntaxToken>
+        {
+            public static readonly TokenOrderComparer Instance = new TokenOrderComparer();
+
+            private TokenOrderComparer() { }
+
+            public int Compare(SyntaxToken x, SyntaxToken y)
+            {
+                return x.FullSpan.CompareTo(y.FullSpan);
+            }
+        }
     }
 }