Merge pull request #73708 from CyrusNajmabadi/treeEfficiency

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

@@ -24,13 +24,14 @@ namespace Microsoft.CodeAnalysis.Editor.Shared.Tagging;
 internal sealed partial class TagSpanIntervalTree<TTag>(
     ITextBuffer textBuffer,
     SpanTrackingMode trackingMode,
-    IEnumerable<ITagSpan<TTag>>? values = null) where TTag : ITag
+    IEnumerable<ITagSpan<TTag>>? values1 = null,
+    IEnumerable<ITagSpan<TTag>>? values2 = null) where TTag : ITag
 {
     private readonly ITextBuffer _textBuffer = textBuffer;
     private readonly SpanTrackingMode _spanTrackingMode = trackingMode;
     private readonly IntervalTree<ITagSpan<TTag>> _tree = IntervalTree.Create(
         new IntervalIntrospector(textBuffer.CurrentSnapshot, trackingMode),
-        values);
+        values1, values2);
 
     private static SnapshotSpan GetTranslatedSpan(
         ITagSpan<TTag> originalTagSpan, ITextSnapshot textSnapshot, SpanTrackingMode trackingMode)
@@ -42,8 +43,17 @@ private static SnapshotSpan GetTranslatedSpan(
             : localSpan.TranslateTo(textSnapshot, trackingMode);
     }
 
+    private ITagSpan<TTag> GetTranslatedITagSpan(ITagSpan<TTag> originalTagSpan, ITextSnapshot textSnapshot)
+        // Avoid reallocating in the case where we're on the same snapshot.
+        => originalTagSpan.Span.Snapshot == textSnapshot
+            ? originalTagSpan
+            : GetTranslatedTagSpan(originalTagSpan, textSnapshot, _spanTrackingMode);
+
     private static TagSpan<TTag> GetTranslatedTagSpan(ITagSpan<TTag> originalTagSpan, ITextSnapshot textSnapshot, SpanTrackingMode trackingMode)
-        => new(GetTranslatedSpan(originalTagSpan, textSnapshot, trackingMode), originalTagSpan.Tag);
+        // Avoid reallocating in the case where we're on the same snapshot.
+        => originalTagSpan is TagSpan<TTag> tagSpan && tagSpan.Span.Snapshot == textSnapshot
+            ? tagSpan
+            : new(GetTranslatedSpan(originalTagSpan, textSnapshot, trackingMode), originalTagSpan.Tag);
 
     public ITextBuffer Buffer => _textBuffer;
 
@@ -83,7 +93,42 @@ ref intersectingIntervals.AsRef(),
     }
 
     public IEnumerable<ITagSpan<TTag>> GetSpans(ITextSnapshot snapshot)
-        => _tree.Select(tn => GetTranslatedTagSpan(tn, snapshot, _spanTrackingMode));
+        => _tree.Select(tn => GetTranslatedITagSpan(tn, snapshot));
+
+    /// <summary>
+    /// Adds all the tag spans in <see langword="this"/> to <paramref name="tagSpans"/>, translating them to the given
+    /// location <paramref name="textSnapshot"/> based on <see cref="_spanTrackingMode"/>.
+    /// </summary>
+    public void AddAllSpans(ITextSnapshot textSnapshot, HashSet<ITagSpan<TTag>> tagSpans)
+    {
+        foreach (var tagSpan in _tree)
+            tagSpans.Add(GetTranslatedITagSpan(tagSpan, textSnapshot));
+    }
+
+    /// <summary>
+    /// Removes from <paramref name="tagSpans"/> all the tags spans in <see langword="this"/> that intersect with any of
+    /// the spans in <paramref name="snapshotSpansToRemove"/>.
+    /// </summary>
+    public void RemoveIntersectingTagSpans(
+        ArrayBuilder<SnapshotSpan> snapshotSpansToRemove, HashSet<ITagSpan<TTag>> tagSpans)
+    {
+        using var buffer = TemporaryArray<ITagSpan<TTag>>.Empty;
+
+        foreach (var snapshotSpan in snapshotSpansToRemove)
+        {
+            buffer.Clear();
+
+            var textSnapshot = snapshotSpan.Snapshot;
+            _tree.FillWithIntervalsThatIntersectWith(
+                snapshotSpan.Span.Start,
+                snapshotSpan.Span.Length,
+                ref buffer.AsRef(),
+                new IntervalIntrospector(textSnapshot, _spanTrackingMode));
+
+            foreach (var tagSpan in buffer)
+                tagSpans.Remove(GetTranslatedITagSpan(tagSpan, textSnapshot));
+        }
+    }
 
     public bool IsEmpty()
         => _tree.IsEmpty();

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

@@ -12,6 +12,7 @@
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.Editor.Shared.Tagging;
 using Microsoft.CodeAnalysis.Options;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.CodeAnalysis.Workspaces;
@@ -171,6 +172,7 @@ public TagSource(
             _dataSource = dataSource;
             _asyncListener = asyncListener;
             _nonFrozenComputationCancellationSeries = new(_disposalTokenSource.Token);
+            _tagSpanSetPool = new ObjectPool<HashSet<ITagSpan<TTag>>>(() => new HashSet<ITagSpan<TTag>>(this), trimOnFree: false);
 
             _workspaceRegistration = Workspace.GetWorkspaceRegistration(subjectBuffer.AsTextContainer());
 

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

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System.Collections.Generic;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.VisualStudio.Text.Tagging;
 
 namespace Microsoft.CodeAnalysis.Editor.Tagging;
@@ -11,8 +12,18 @@ internal abstract partial class AbstractAsynchronousTaggerProvider<TTag>
 {
     private partial class TagSource : IEqualityComparer<ITagSpan<TTag>>
     {
+        private readonly ObjectPool<HashSet<ITagSpan<TTag>>> _tagSpanSetPool;
+
         public bool Equals(ITagSpan<TTag>? x, ITagSpan<TTag>? y)
-            => x != null && y != null && x.Span == y.Span && _dataSource.TagEquals(x.Tag, y.Tag);
+        {
+            if (x == y)
+                return true;
+
+            if (x is null || y is null)
+                return false;
+
+            return x.Span == y.Span && _dataSource.TagEquals(x.Tag, y.Tag);
+        }
 
         /// <summary>
         /// For the purposes of hashing, just hash spans.  This will prevent most collisions.  And the rare

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

@@ -389,52 +389,56 @@ private static void CheckSnapshot(ITextSnapshot snapshot)
 
         private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTagTrees(ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> oldTagTrees, TaggerContext<TTag> context)
         {
-            // Ignore any tag spans reported for any buffers we weren't interested in.
-
-            var spansToTag = context.SpansToTag;
-            var buffersToTag = spansToTag.Select(dss => dss.SnapshotSpan.Snapshot.TextBuffer).ToSet();
-            var newTagsByBuffer =
-                context.TagSpans.Where(ts => buffersToTag.Contains(ts.Span.Snapshot.TextBuffer))
-                                .ToLookup(t => t.Span.Snapshot.TextBuffer);
-            var spansTagged = context._spansTagged;
-
-            var spansToInvalidateByBuffer = spansTagged.ToLookup(
-                keySelector: span => span.Snapshot.TextBuffer,
-                elementSelector: span => span);
-
-            // Walk through each relevant buffer and decide what the interval tree should be
-            // for that buffer.  In general this will work by keeping around old tags that
-            // weren't in the range that was re-tagged, and merging them with the new tags
-            // produced for the range that was re-tagged.
-            var newTagTrees = ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>>.Empty;
+            using var _1 = PooledHashSet<ITextBuffer>.GetInstance(out var buffersToTag);
+            foreach (var spanToTag in context.SpansToTag)
+                buffersToTag.Add(spanToTag.SnapshotSpan.Snapshot.TextBuffer);
+
+            using var _2 = ArrayBuilder<ITagSpan<TTag>>.GetInstance(out var newTagsInBuffer);
+            using var _3 = ArrayBuilder<SnapshotSpan>.GetInstance(out var spansToInvalidateInBuffer);
+
+            var newTagTrees = ImmutableDictionary.CreateBuilder<ITextBuffer, TagSpanIntervalTree<TTag>>();
             foreach (var buffer in buffersToTag)
             {
-                var newTagTree = ComputeNewTagTree(oldTagTrees, buffer, newTagsByBuffer[buffer], spansToInvalidateByBuffer[buffer]);
+                newTagsInBuffer.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);
+                }
+
+                foreach (var span in context._spansTagged)
+                {
+                    if (span.Snapshot.TextBuffer == buffer)
+                        spansToInvalidateInBuffer.Add(span);
+                }
+
+                var newTagTree = ComputeNewTagTree(oldTagTrees, buffer, newTagsInBuffer, spansToInvalidateInBuffer);
                 if (newTagTree != null)
-                    newTagTrees = newTagTrees.Add(buffer, newTagTree);
+                    newTagTrees.Add(buffer, newTagTree);
             }
 
-            return newTagTrees;
+            return newTagTrees.ToImmutable();
         }
 
         private TagSpanIntervalTree<TTag>? ComputeNewTagTree(
             ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> oldTagTrees,
             ITextBuffer textBuffer,
-            IEnumerable<ITagSpan<TTag>> newTags,
-            IEnumerable<SnapshotSpan> spansToInvalidate)
+            ArrayBuilder<ITagSpan<TTag>> newTags,
+            ArrayBuilder<SnapshotSpan> spansToInvalidate)
         {
-            var noNewTags = newTags.IsEmpty();
-            var noSpansToInvalidate = spansToInvalidate.IsEmpty();
+            var noNewTags = newTags.IsEmpty;
+            var noSpansToInvalidate = spansToInvalidate.IsEmpty;
             oldTagTrees.TryGetValue(textBuffer, out var oldTagTree);
 
             if (oldTagTree == null)
             {
+                // If we have no new tags, and no old tags either.  No need to store anything for this buffer.
                 if (noNewTags)
-                {
-                    // We have no new tags, and no old tags either.  No need to store anything
-                    // for this buffer.
                     return null;
-                }
 
                 // If we don't have any old tags then we just need to return the new tags.
                 return new TagSpanIntervalTree<TTag>(textBuffer, _dataSource.SpanTrackingMode, newTags);
@@ -443,35 +447,40 @@ private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTa
             // If we don't have any new tags, and there was nothing to invalidate, then we can 
             // keep whatever old tags we have without doing any additional work.
             if (noNewTags && noSpansToInvalidate)
-            {
                 return oldTagTree;
-            }
 
-            // We either have some new tags, or we have some tags to invalidate.
-            // First, determine which of the old tags we want to keep around.
-            var snapshot = noNewTags ? spansToInvalidate.First().Snapshot : newTags.First().Span.Snapshot;
-            var oldTagsToKeep = noSpansToInvalidate
-                ? oldTagTree.GetSpans(snapshot)
-                : GetNonIntersectingTagSpans(spansToInvalidate, oldTagTree);
-
-            // Then union those with the new tags to produce the final tag tree.
-            var finalTags = oldTagsToKeep.Concat(newTags);
-            return new TagSpanIntervalTree<TTag>(textBuffer, _dataSource.SpanTrackingMode, finalTags);
+            if (noSpansToInvalidate)
+            {
+                // If we have no spans to invalidate, then we can just keep the old tags and add the new tags.
+                var oldTagsToKeep = oldTagTree.GetSpans(newTags.First().Span.Snapshot);
+                return new TagSpanIntervalTree<TTag>(
+                    textBuffer, _dataSource.SpanTrackingMode, oldTagsToKeep, newTags);
+            }
+            else
+            {
+                // We do have spans to invalidate. Get the set of old tags that don't intersect with those and add the new tags.
+                using var _1 = _tagSpanSetPool.GetPooledObject(out var nonIntersectingTagSpans);
+                AddNonIntersectingTagSpans(spansToInvalidate, oldTagTree, nonIntersectingTagSpans);
+                return new TagSpanIntervalTree<TTag>(
+                    textBuffer, _dataSource.SpanTrackingMode, nonIntersectingTagSpans, newTags);
+            }
         }
 
-        private IEnumerable<ITagSpan<TTag>> GetNonIntersectingTagSpans(IEnumerable<SnapshotSpan> spansToInvalidate, TagSpanIntervalTree<TTag> oldTagTree)
+        private static void AddNonIntersectingTagSpans(
+            ArrayBuilder<SnapshotSpan> spansToInvalidate,
+            TagSpanIntervalTree<TTag> oldTagTree,
+            HashSet<ITagSpan<TTag>> nonIntersectingTagSpans)
         {
             var firstSpanToInvalidate = spansToInvalidate.First();
             var snapshot = firstSpanToInvalidate.Snapshot;
 
             // Performance: No need to fully realize spansToInvalidate or do any of the calculations below if the
             //   full snapshot is being invalidated.
             if (firstSpanToInvalidate.Length == snapshot.Length)
-                return [];
+                return;
 
-            return oldTagTree.GetSpans(snapshot).Except(
-                spansToInvalidate.SelectMany(oldTagTree.GetIntersectingSpans),
-                comparer: this);
+            oldTagTree.AddAllSpans(snapshot, nonIntersectingTagSpans);
+            oldTagTree.RemoveIntersectingTagSpans(spansToInvalidate, nonIntersectingTagSpans);
         }
 
         private bool ShouldSkipTagProduction()

src/Workspaces/Core/Portable/Shared/Collections/IntervalTree.cs

@@ -2,8 +2,6 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-#nullable disable
-
 using System.Collections.Generic;
 
 namespace Microsoft.CodeAnalysis.Shared.Collections;
@@ -16,9 +14,9 @@ public static IntervalTree<T> Create<T, TIntrospector>(in TIntrospector introspe
         return Create(in introspector, (IEnumerable<T>)values);
     }
 
-    public static IntervalTree<T> Create<T, TIntrospector>(in TIntrospector introspector, IEnumerable<T> values = null)
+    public static IntervalTree<T> Create<T, TIntrospector>(in TIntrospector introspector, IEnumerable<T>? values1 = null, IEnumerable<T>? values2 = null)
         where TIntrospector : struct, IIntervalIntrospector<T>
     {
-        return IntervalTree<T>.Create(in introspector, values ?? []);
+        return IntervalTree<T>.Create(in introspector, values1, values2);
     }
 }

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

@@ -37,16 +37,24 @@ 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> values)
+    public static IntervalTree<T> Create<TIntrospector>(in TIntrospector introspector, IEnumerable<T>? values1 = null, IEnumerable<T>? values2 = null)
         where TIntrospector : struct, IIntervalIntrospector<T>
     {
         var result = new IntervalTree<T>();
-        foreach (var value in values)
-        {
-            result.root = Insert(result.root, new Node(value), in introspector);
-        }
+
+        AddAll(in introspector, values1);
+        AddAll(in introspector, values2);
 
         return result;
+
+        void AddAll(in TIntrospector introspector, IEnumerable<T>? values)
+        {
+            if (values != null)
+            {
+                foreach (var value in values)
+                    result.root = Insert(result.root, new Node(value), in introspector);
+            }
+        }
     }
 
     protected static bool Contains<TIntrospector>(T value, int start, int length, in TIntrospector introspector)

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/ObjectPools/Extensions.cs

@@ -152,7 +152,7 @@ public static void ClearAndFree<T>(this ObjectPool<HashSet<T>> pool, HashSet<T>
         var count = set.Count;
         set.Clear();
 
-        if (count > Threshold)
+        if (count > Threshold && pool.TrimOnFree)
         {
             set.TrimExcess();
         }