Build interval tree from sorted nodes

src/EditorFeatures/Core/Shared/Tagging/Utilities/TagSpanIntervalTree.cs

@@ -32,13 +32,15 @@ internal sealed partial class TagSpanIntervalTree<TTag>(SpanTrackingMode spanTra
     public TagSpanIntervalTree(
         ITextSnapshot textSnapshot,
         SpanTrackingMode trackingMode,
-        IEnumerable<TagSpan<TTag>>? values1 = null,
-        IEnumerable<TagSpan<TTag>>? values2 = null)
+        SegmentedList<TagSpan<TTag>> values)
         : this(trackingMode)
     {
-        _tree = IntervalTree.Create(
-            new IntervalIntrospector(textSnapshot, trackingMode),
-            values1, values2);
+        // Sort the values by their start position.  This is extremely fast (defer'ing to the runtime's sorting
+        // routines), and allows us to build the balanced tree directly without having to do any additional work.
+        values.Sort(static (t1, t2) => t1.Span.Start.Position - t2.Span.Start.Position);
+
+        _tree = IntervalTree<TagSpan<TTag>>.CreateFromSorted(
+            new IntervalIntrospector(textSnapshot, trackingMode), values);
     }
 
     private static SnapshotSpan GetTranslatedSpan(TagSpan<TTag> originalTagSpan, ITextSnapshot textSnapshot, SpanTrackingMode trackingMode)

src/EditorFeatures/Core/Tagging/AbstractAsynchronousTaggerProvider.TagSource.cs

@@ -49,7 +49,6 @@ private sealed partial class TagSource
         /// </summary>
         private const int CoalesceDifferenceCount = 10;
 
-        private static readonly ObjectPool<SegmentedList<TagSpan<TTag>>> s_tagSpanListPool = new(() => new(), trimOnFree: false);
         private readonly ObjectPool<HashSet<TagSpan<TTag>>> _tagSpanSetPool;
 
         #region Fields that can be accessed from either thread

src/EditorFeatures/Core/Tagging/AbstractAsynchronousTaggerProvider.TagSource_ProduceTags.cs

@@ -95,16 +95,21 @@ private void RemoveTagsThatIntersectEdit(TextContentChangedEventArgs e)
                 return;
 
             using var _1 = SegmentedListPool.GetPooledList<TagSpan<TTag>>(out var tagsToRemove);
-            using var _2 = _tagSpanSetPool.GetPooledObject(out var allTags);
+            using var _2 = SegmentedListPool.GetPooledList<TagSpan<TTag>>(out var allTagsList);
+            using var _3 = _tagSpanSetPool.GetPooledObject(out var allTagsSet);
 
             // Everything we're passing in here is synchronous.  So we can assert that this must complete synchronously
             // as well.
             var (oldTagTrees, newTagTrees, _) = CompareAndSwapTagTreesAsync(
                 static (oldTagTrees, args, _) =>
                 {
-                    var (@this, e, tagsToRemove, allTags) = args;
-                    args.tagsToRemove.Clear();
-                    args.allTags.Clear();
+                    var (@this, e, tagsToRemove, allTagsList, allTagsSet) = args;
+
+                    // Compre-and-swap loops until we can successfully update the tag trees.  Clear out the collections
+                    // so we're back in an initial state before performing any work in this lambda.
+                    tagsToRemove.Clear();
+                    allTagsList.Clear();
+                    allTagsSet.Clear();
 
                     var snapshot = e.After;
                     var buffer = snapshot.TextBuffer;
@@ -116,22 +121,27 @@ private void RemoveTagsThatIntersectEdit(TextContentChangedEventArgs e)
 
                         if (tagsToRemove.Count > 0)
                         {
-                            treeForBuffer.AddAllSpans(snapshot, allTags);
+                            // Determine the final tags for the interval tree, using a set so that we can efficiently
+                            // remove the intersecting tags.
+                            treeForBuffer.AddAllSpans(snapshot, allTagsSet);
+                            allTagsSet.RemoveAll(tagsToRemove);
 
-                            allTags.RemoveAll(tagsToRemove);
+                            // Then, copy into a list so we can efficiently sort them and create the interval tree from
+                            // those sorted items.
+                            allTagsList.AddRange(allTagsSet);
 
                             var newTagTree = new TagSpanIntervalTree<TTag>(
                                 snapshot,
                                 @this._dataSource.SpanTrackingMode,
-                                allTags);
+                                allTagsList);
                             return ValueTaskFactory.FromResult((oldTagTrees.SetItem(buffer, newTagTree), default(VoidResult)));
                         }
                     }
 
                     // return oldTagTrees to indicate nothing changed.
                     return ValueTaskFactory.FromResult((oldTagTrees, default(VoidResult)));
                 },
-                args: (this, e, tagsToRemove, allTags),
+                args: (this, e, tagsToRemove, allTagsList, allTagsSet),
                 _disposalTokenSource.Token).VerifyCompleted();
 
             // Can happen if we were canceled.  Just bail out immediate.
@@ -450,21 +460,21 @@ private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTa
             foreach (var spanToTag in context.SpansToTag)
                 buffersToTag.Add(spanToTag.SnapshotSpan.Snapshot.TextBuffer);
 
-            using var _2 = s_tagSpanListPool.GetPooledObject(out var newTagsInBuffer);
+            using var _2 = SegmentedListPool.GetPooledList<TagSpan<TTag>>(out var newTagsInBuffer_safeToMutate);
             using var _3 = ArrayBuilder<SnapshotSpan>.GetInstance(out var spansToInvalidateInBuffer);
 
             var newTagTrees = ImmutableDictionary.CreateBuilder<ITextBuffer, TagSpanIntervalTree<TTag>>();
             foreach (var buffer in buffersToTag)
             {
-                newTagsInBuffer.Clear();
+                newTagsInBuffer_safeToMutate.Clear();
                 spansToInvalidateInBuffer.Clear();
 
                 // Ignore any tag spans reported for any buffers we weren't interested in.
 
                 foreach (var tagSpan in context.TagSpans)
                 {
                     if (tagSpan.Span.Snapshot.TextBuffer == buffer)
-                        newTagsInBuffer.Add(tagSpan);
+                        newTagsInBuffer_safeToMutate.Add(tagSpan);
                 }
 
                 // Invalidate all the spans that were actually tagged.  If the context doesn't have any recorded spans
@@ -475,7 +485,10 @@ private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTa
                         spansToInvalidateInBuffer.Add(span);
                 }
 
-                var newTagTree = ComputeNewTagTree(oldTagTrees, buffer, newTagsInBuffer, spansToInvalidateInBuffer);
+                // Note: newTagsInBuffer_safeToMutate will be mutated by ComputeNewTagTree.  This is fine as we don't
+                // use it after this and immediately clear it on the next iteration of the loop (or dispose of it once
+                // the loop finishes).
+                var newTagTree = ComputeNewTagTree(oldTagTrees, buffer, newTagsInBuffer_safeToMutate, spansToInvalidateInBuffer);
                 if (newTagTree != null)
                     newTagTrees.Add(buffer, newTagTree);
             }
@@ -486,10 +499,10 @@ private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTa
         private TagSpanIntervalTree<TTag>? ComputeNewTagTree(
             ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> oldTagTrees,
             ITextBuffer textBuffer,
-            SegmentedList<TagSpan<TTag>> newTags,
+            SegmentedList<TagSpan<TTag>> newTags_safeToMutate,
             ArrayBuilder<SnapshotSpan> spansToInvalidate)
         {
-            var noNewTags = newTags.Count == 0;
+            var noNewTags = newTags_safeToMutate.Count == 0;
             var noSpansToInvalidate = spansToInvalidate.IsEmpty;
             oldTagTrees.TryGetValue(textBuffer, out var oldTagTree);
 
@@ -500,7 +513,7 @@ private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTa
                     return null;
 
                 // If we don't have any old tags then we just need to return the new tags.
-                return new TagSpanIntervalTree<TTag>(newTags[0].Span.Snapshot, _dataSource.SpanTrackingMode, newTags);
+                return new TagSpanIntervalTree<TTag>(newTags_safeToMutate[0].Span.Snapshot, _dataSource.SpanTrackingMode, newTags_safeToMutate);
             }
 
             // If we don't have any new tags, and there was nothing to invalidate, then we can 
@@ -511,14 +524,14 @@ private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTa
             if (noSpansToInvalidate)
             {
                 // If we have no spans to invalidate, then we can just keep the old tags and add the new tags.
-                var snapshot = newTags.First().Span.Snapshot;
+                var snapshot = newTags_safeToMutate.First().Span.Snapshot;
 
                 // For efficiency, just grab the old tags, remap them to the current snapshot, and place them in the
                 // newTags buffer.  This is a safe mutation of this buffer as the caller doesn't use it after this point
                 // and instead immediately clears it.
-                oldTagTree.AddAllSpans(snapshot, newTags);
+                oldTagTree.AddAllSpans(snapshot, newTags_safeToMutate);
                 return new TagSpanIntervalTree<TTag>(
-                    snapshot, _dataSource.SpanTrackingMode, newTags);
+                    snapshot, _dataSource.SpanTrackingMode, newTags_safeToMutate);
             }
             else
             {
@@ -536,8 +549,11 @@ private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTa
                     oldTagTree.RemoveIntersectingTagSpans(spansToInvalidate, nonIntersectingOldTags);
                 }
 
+                // For efficiency, add the non-intersecting old tags to the new tags buffer.  This is a safe mutation of
+                // of that buffer as it is not used by us or our caller after this point.
+                newTags_safeToMutate.AddRange(nonIntersectingOldTags);
                 return new TagSpanIntervalTree<TTag>(
-                    snapshot, _dataSource.SpanTrackingMode, nonIntersectingOldTags, newTags);
+                    snapshot, _dataSource.SpanTrackingMode, newTags_safeToMutate);
             }
         }
 
@@ -613,8 +629,8 @@ private DiffResult ComputeDifference(
             TagSpanIntervalTree<TTag> latestTree,
             TagSpanIntervalTree<TTag> previousTree)
         {
-            using var _1 = s_tagSpanListPool.GetPooledObject(out var latestSpans);
-            using var _2 = s_tagSpanListPool.GetPooledObject(out var previousSpans);
+            using var _1 = SegmentedListPool.GetPooledList<TagSpan<TTag>>(out var latestSpans);
+            using var _2 = SegmentedListPool.GetPooledList<TagSpan<TTag>>(out var previousSpans);
 
             using var _3 = ArrayBuilder<SnapshotSpan>.GetInstance(out var added);
             using var _4 = ArrayBuilder<SnapshotSpan>.GetInstance(out var removed);

src/EditorFeatures/Test/Tagging/TagSpanIntervalTreeTests.cs

@@ -24,7 +24,8 @@ private static (TagSpanIntervalTree<ITextMarkerTag>, ITextBuffer) CreateTree(str
     {
         var exportProvider = EditorTestCompositions.Editor.ExportProviderFactory.CreateExportProvider();
         var buffer = EditorFactory.CreateBuffer(exportProvider, text);
-        var tags = spans.Select(s => new TagSpan<ITextMarkerTag>(new SnapshotSpan(buffer.CurrentSnapshot, s), new TextMarkerTag(string.Empty)));
+        var tags = new SegmentedList<TagSpan<ITextMarkerTag>>(
+            spans.Select(s => new TagSpan<ITextMarkerTag>(new SnapshotSpan(buffer.CurrentSnapshot, s), new TextMarkerTag(string.Empty))));
         return (new TagSpanIntervalTree<ITextMarkerTag>(buffer.CurrentSnapshot, SpanTrackingMode.EdgeInclusive, tags), buffer);
     }
 

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/Collections/IntervalTree`1.cs

@@ -8,6 +8,7 @@
 using System.Collections.Immutable;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
+using Microsoft.CodeAnalysis.Collections;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Roslyn.Utilities;
@@ -35,24 +36,69 @@ private delegate bool TestInterval<TIntrospector>(T value, int start, int length
 
     protected Node? root;
 
-    public static IntervalTree<T> Create<TIntrospector>(in TIntrospector introspector, IEnumerable<T>? values1 = null, IEnumerable<T>? values2 = null)
+    public static IntervalTree<T> Create<TIntrospector>(in TIntrospector introspector, IEnumerable<T>? values = null)
         where TIntrospector : struct, IIntervalIntrospector<T>
     {
         var result = new IntervalTree<T>();
 
-        AddAll(in introspector, values1);
-        AddAll(in introspector, values2);
+        if (values != null)
+        {
+            foreach (var value in values)
+                result.root = Insert(result.root, new Node(value), in introspector);
+        }
 
         return result;
+    }
+
+    /// <summary>
+    /// Creates an interval tree from a sorted list of values.  This is more efficient than creating from an unsorted
+    /// list as building doesn't need to figure out where the nodes need to go n-log(n) and doesn't have to rebalance
+    /// anything (again, another n-log(n) operation).  Rebalancing is particularly expensive as it involves tons of
+    /// pointer chasing operations, which is both slow, and which impacts the GC which has to track all those writes.
+    /// </summary>
+    /// <remarks>
+    /// The values must be sorted such that given any two elements 'a' and 'b' in the list, if 'a' comes before 'b' in
+    /// the list, then it's "start position" (as determined by the introspector) must be less than or equal to 'b's
+    /// start position.  This is a requirement for the algorithm to work correctly.
+    /// </remarks>
+    public static IntervalTree<T> CreateFromSorted<TIntrospector>(in TIntrospector introspector, SegmentedList<T> values)
+        where TIntrospector : struct, IIntervalIntrospector<T>
+    {
+#if DEBUG
+        var localIntrospector = introspector;
+        Debug.Assert(values.IsSorted(Comparer<T>.Create((t1, t2) => localIntrospector.GetSpan(t1).Start - localIntrospector.GetSpan(t2).Start)));
+#endif
+
+        if (values.Count == 0)
+            return Empty;
 
-        void AddAll(in TIntrospector introspector, IEnumerable<T>? values)
+        return new IntervalTree<T>
         {
-            if (values != null)
-            {
-                foreach (var value in values)
-                    result.root = Insert(result.root, new Node(value), in introspector);
-            }
-        }
+            root = CreateFromSortedWorker(values, 0, values.Count, in introspector),
+        };
+    }
+
+    private static Node? CreateFromSortedWorker<TIntrospector>(
+        SegmentedList<T> values, int startInclusive, int endExclusive, in TIntrospector introspector) where TIntrospector : struct, IIntervalIntrospector<T>
+    {
+        var length = endExclusive - startInclusive;
+        if (length <= 0)
+            return null;
+
+        var mid = startInclusive + (length >> 1);
+        var node = new Node(values[mid]);
+        node.SetLeftRight(
+            CreateFromSortedWorker(values, startInclusive, mid, in introspector),
+            CreateFromSortedWorker(values, mid + 1, endExclusive, in introspector),
+            in introspector);
+
+        // Everything is sorted, and we're always building a node up from equal subtrees.  So we're never unbalanced
+        // enough to require balancing here.
+        var balanceFactor = BalanceFactor(node);
+        Debug.Assert(balanceFactor >= -1, "balanceFactor >= -1");
+        Debug.Assert(balanceFactor <= 1, "balanceFactor <= 1");
+
+        return node;
     }
 
     protected static bool Contains<TIntrospector>(T value, int start, int length, in TIntrospector introspector)