Reduce lots of allocations in TagSpanIntervalTree

src/EditorFeatures/Core/Classification/CopyPasteAndPrintingClassificationBufferTaggerProvider.Tagger.cs

@@ -159,8 +159,7 @@ await TotalClassificationAggregateTagger.AddTagsAsync(
                     arg: default).ConfigureAwait(false);
             });
 
-            var cachedTags = new TagSpanIntervalTree<IClassificationTag>(snapshot.TextBuffer, SpanTrackingMode.EdgeExclusive, mergedTags);
-
+            var cachedTags = new TagSpanIntervalTree<IClassificationTag>(snapshot, SpanTrackingMode.EdgeExclusive, mergedTags);
             lock (_gate)
             {
                 _cachedTaggedSpan = spanToTag;

src/EditorFeatures/Core/KeywordHighlighting/HighlighterViewTaggerProvider.cs

@@ -93,10 +93,10 @@ protected override async Task ProduceTagsAsync(
         var position = caretPosition.Value;
         var snapshot = snapshotSpan.Snapshot;
 
-        // See if the user is just moving their caret around in an existing tag.  If so, we don't
-        // want to actually go recompute things.  Note: this only works for containment.  If the
-        // user moves their caret to the end of a highlighted reference, we do want to recompute
-        // as they may now be at the start of some other reference that should be highlighted instead.
+        // See if the user is just moving their caret around in an existing tag.  If so, we don't want to actually go
+        // recompute things.  Note: this only works for containment.  If the user moves their caret to the end of a
+        // highlighted reference, we do want to recompute as they may now be at the start of some other reference that
+        // should be highlighted instead.
         var onExistingTags = context.HasExistingContainingTags(new SnapshotPoint(snapshot, position));
         if (onExistingTags)
         {

src/EditorFeatures/Core/ReferenceHighlighting/ReferenceHighlightingViewTaggerProvider.cs

@@ -122,10 +122,10 @@ protected override Task ProduceTagsAsync(
             return Task.CompletedTask;
         }
 
-        // See if the user is just moving their caret around in an existing tag.  If so, we don't
-        // want to actually go recompute things.  Note: this only works for containment.  If the
-        // user moves their caret to the end of a highlighted reference, we do want to recompute
-        // as they may now be at the start of some other reference that should be highlighted instead.
+        // See if the user is just moving their caret around in an existing tag.  If so, we don't want to actually go
+        // recompute things.  Note: this only works for containment.  If the user moves their caret to the end of a
+        // highlighted reference, we do want to recompute as they may now be at the start of some other reference that
+        // should be highlighted instead.
         var onExistingTags = context.HasExistingContainingTags(caretPosition);
         if (onExistingTags)
         {

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

@@ -21,66 +21,53 @@ namespace Microsoft.CodeAnalysis.Editor.Shared.Tagging;
 /// tracked. That way you can query for intersecting/overlapping spans in a different snapshot
 /// than the one for the tag spans that were added.
 /// </summary>
-internal sealed partial class TagSpanIntervalTree<TTag>(
-    ITextBuffer textBuffer,
-    SpanTrackingMode trackingMode,
-    IEnumerable<TagSpan<TTag>>? values1 = null,
-    IEnumerable<TagSpan<TTag>>? values2 = null) where TTag : ITag
+internal sealed partial class TagSpanIntervalTree<TTag>(SpanTrackingMode spanTrackingMode) where TTag : ITag
 {
-    private readonly ITextBuffer _textBuffer = textBuffer;
-    private readonly SpanTrackingMode _spanTrackingMode = trackingMode;
-    private readonly IntervalTree<TagSpan<TTag>> _tree = IntervalTree.Create(
-        new IntervalIntrospector(textBuffer.CurrentSnapshot, trackingMode),
-        values1, values2);
-
-    private static SnapshotSpan GetTranslatedSpan(
-        TagSpan<TTag> originalTagSpan, ITextSnapshot textSnapshot, SpanTrackingMode trackingMode)
+    // Tracking mode passed in here doesn't matter (since the tree is empty).
+    public static readonly TagSpanIntervalTree<TTag> Empty = new(SpanTrackingMode.EdgeInclusive);
+
+    private readonly SpanTrackingMode _spanTrackingMode = spanTrackingMode;
+    private readonly IntervalTree<TagSpan<TTag>> _tree = IntervalTree<TagSpan<TTag>>.Empty;
+
+    public TagSpanIntervalTree(
+        ITextSnapshot textSnapshot,
+        SpanTrackingMode trackingMode,
+        IEnumerable<TagSpan<TTag>>? values1 = null,
+        IEnumerable<TagSpan<TTag>>? values2 = null)
+        : this(trackingMode)
     {
-        var localSpan = originalTagSpan.Span;
-
-        return localSpan.Snapshot == textSnapshot
-            ? localSpan
-            : localSpan.TranslateTo(textSnapshot, trackingMode);
+        _tree = IntervalTree.Create(
+            new IntervalIntrospector(textSnapshot, trackingMode),
+            values1, values2);
     }
 
-    private TagSpan<TTag> GetTranslatedITagSpan(TagSpan<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 SnapshotSpan GetTranslatedSpan(TagSpan<TTag> originalTagSpan, ITextSnapshot textSnapshot, SpanTrackingMode trackingMode)
+        // SnapshotSpan no-ops if you pass it the same snapshot that it is holding onto.
+        => originalTagSpan.Span.TranslateTo(textSnapshot, trackingMode);
+
+    private TagSpan<TTag> GetTranslatedTagSpan(TagSpan<TTag> originalTagSpan, ITextSnapshot textSnapshot)
+        => GetTranslatedTagSpan(originalTagSpan, textSnapshot, _spanTrackingMode);
 
     private static TagSpan<TTag> GetTranslatedTagSpan(TagSpan<TTag> originalTagSpan, ITextSnapshot textSnapshot, SpanTrackingMode trackingMode)
         // Avoid reallocating in the case where we're on the same snapshot.
-        => originalTagSpan is TagSpan<TTag> tagSpan && tagSpan.Span.Snapshot == textSnapshot
-            ? tagSpan
+        => originalTagSpan.Span.Snapshot == textSnapshot
+            ? originalTagSpan
             : new(GetTranslatedSpan(originalTagSpan, textSnapshot, trackingMode), originalTagSpan.Tag);
 
-    public ITextBuffer Buffer => _textBuffer;
-
-    public SpanTrackingMode SpanTrackingMode => _spanTrackingMode;
-
     public bool HasSpanThatContains(SnapshotPoint point)
-    {
-        var snapshot = point.Snapshot;
-        Debug.Assert(snapshot.TextBuffer == _textBuffer);
-
-        return _tree.HasIntervalThatContains(point.Position, length: 0, new IntervalIntrospector(snapshot, _spanTrackingMode));
-    }
+        => _tree.HasIntervalThatContains(point.Position, length: 0, new IntervalIntrospector(point.Snapshot, _spanTrackingMode));
 
-    public IReadOnlyList<TagSpan<TTag>> GetIntersectingSpans(SnapshotSpan snapshotSpan)
-        => SegmentedListPool<TagSpan<TTag>>.ComputeList(
-            static (args, tags) => [email protected](args.snapshotSpan, tags),
-            (@this: this, snapshotSpan));
+    public bool HasSpanThatIntersects(SnapshotPoint point)
+        => _tree.HasIntervalThatIntersectsWith(point.Position, new IntervalIntrospector(point.Snapshot, _spanTrackingMode));
 
     /// <summary>
     /// Gets all the spans that intersect with <paramref name="snapshotSpan"/> in sorted order and adds them to
     /// <paramref name="result"/>.  Note the sorted chunk of items are appended to <paramref name="result"/>.  This
     /// means that <paramref name="result"/> may not be sorted if there were already items in them.
     /// </summary>
-    private void AppendIntersectingSpansInSortedOrder(SnapshotSpan snapshotSpan, SegmentedList<TagSpan<TTag>> result)
+    public void AddIntersectingTagSpans(SnapshotSpan snapshotSpan, SegmentedList<TagSpan<TTag>> result)
     {
         var snapshot = snapshotSpan.Snapshot;
-        Debug.Assert(snapshot.TextBuffer == _textBuffer);
 
         using var intersectingIntervals = TemporaryArray<TagSpan<TTag>>.Empty;
         _tree.FillWithIntervalsThatIntersectWith(
@@ -92,8 +79,24 @@ ref intersectingIntervals.AsRef(),
             result.Add(GetTranslatedTagSpan(tagSpan, snapshot, _spanTrackingMode));
     }
 
-    public IEnumerable<TagSpan<TTag>> GetSpans(ITextSnapshot snapshot)
-        => _tree.Select(tn => GetTranslatedITagSpan(tn, snapshot));
+    /// <summary>
+    /// Gets all the tag spans in this tree, remapped to <paramref name="snapshot"/>, and returns them as a <see
+    /// cref="NormalizedSnapshotSpanCollection"/>.
+    /// </summary>
+    public NormalizedSnapshotSpanCollection GetSnapshotSpanCollection(ITextSnapshot snapshot)
+    {
+        if (this == Empty)
+            return NormalizedSnapshotSpanCollection.Empty;
+
+        using var _ = ArrayBuilder<SnapshotSpan>.GetInstance(out var spans);
+
+        foreach (var tagSpan in _tree)
+            spans.Add(GetTranslatedSpan(tagSpan, snapshot, _spanTrackingMode));
+
+        return spans.Count == 0
+            ? NormalizedSnapshotSpanCollection.Empty
+            : new(spans);
+    }
 
     /// <summary>
     /// Adds all the tag spans in <see langword="this"/> to <paramref name="tagSpans"/>, translating them to the given
@@ -102,7 +105,15 @@ public IEnumerable<TagSpan<TTag>> GetSpans(ITextSnapshot snapshot)
     public void AddAllSpans(ITextSnapshot textSnapshot, HashSet<TagSpan<TTag>> tagSpans)
     {
         foreach (var tagSpan in _tree)
-            tagSpans.Add(GetTranslatedITagSpan(tagSpan, textSnapshot));
+            tagSpans.Add(GetTranslatedTagSpan(tagSpan, textSnapshot));
+    }
+
+    /// <inheritdoc cref="AddAllSpans(ITextSnapshot, HashSet{TagSpan{TTag}})"/>
+    /// <remarks>Spans will be added in sorted order</remarks>
+    public void AddAllSpans(ITextSnapshot textSnapshot, SegmentedList<TagSpan<TTag>> tagSpans)
+    {
+        foreach (var tagSpan in _tree)
+            tagSpans.Add(GetTranslatedTagSpan(tagSpan, textSnapshot));
     }
 
     /// <summary>
@@ -126,127 +137,118 @@ ref buffer.AsRef(),
                 new IntervalIntrospector(textSnapshot, _spanTrackingMode));
 
             foreach (var tagSpan in buffer)
-                tagSpans.Remove(GetTranslatedITagSpan(tagSpan, textSnapshot));
+                tagSpans.Remove(GetTranslatedTagSpan(tagSpan, textSnapshot));
         }
     }
 
-    public bool IsEmpty()
-        => _tree.IsEmpty();
-
     public void AddIntersectingTagSpans(NormalizedSnapshotSpanCollection requestedSpans, SegmentedList<TagSpan<TTag>> tags)
-    {
-        AddIntersectingTagSpansWorker(requestedSpans, tags);
-        DebugVerifyTags(requestedSpans, tags);
-    }
-
-    [Conditional("DEBUG")]
-    private static void DebugVerifyTags(NormalizedSnapshotSpanCollection requestedSpans, SegmentedList<TagSpan<TTag>> tags)
-    {
-        if (tags == null)
-        {
-            return;
-        }
-
-        foreach (var tag in tags)
-        {
-            var span = tag.Span;
-
-            if (!requestedSpans.Any(s => s.IntersectsWith(span)))
-            {
-                Contract.Fail(tag + " doesn't intersects with any requested span");
-            }
-        }
-    }
-
-    private void AddIntersectingTagSpansWorker(
-        NormalizedSnapshotSpanCollection requestedSpans,
-        SegmentedList<TagSpan<TTag>> tags)
     {
         const int MaxNumberOfRequestedSpans = 100;
 
         // Special case the case where there is only one requested span.  In that case, we don't
         // need to allocate any intermediate collections
         if (requestedSpans.Count == 1)
-            AppendIntersectingSpansInSortedOrder(requestedSpans[0], tags);
+        {
+            AddIntersectingTagSpans(requestedSpans[0], tags);
+        }
         else if (requestedSpans.Count < MaxNumberOfRequestedSpans)
-            AddTagsForSmallNumberOfSpans(requestedSpans, tags);
+        {
+            foreach (var span in requestedSpans)
+                AddIntersectingTagSpans(span, tags);
+        }
         else
+        {
             AddTagsForLargeNumberOfSpans(requestedSpans, tags);
-    }
+        }
 
-    private void AddTagsForSmallNumberOfSpans(
-        NormalizedSnapshotSpanCollection requestedSpans,
-        SegmentedList<TagSpan<TTag>> tags)
-    {
-        foreach (var span in requestedSpans)
-            AppendIntersectingSpansInSortedOrder(span, tags);
-    }
+        DebugVerifyTags(requestedSpans, tags);
+        return;
 
-    private void AddTagsForLargeNumberOfSpans(NormalizedSnapshotSpanCollection requestedSpans, SegmentedList<TagSpan<TTag>> tags)
-    {
-        // we are asked with bunch of spans. rather than asking same question again and again, ask once with big span
-        // which will return superset of what we want. and then filter them out in O(m+n) cost. 
-        // m == number of requested spans, n = number of returned spans
-        var mergedSpan = new SnapshotSpan(requestedSpans[0].Start, requestedSpans[^1].End);
+        void AddTagsForLargeNumberOfSpans(NormalizedSnapshotSpanCollection requestedSpans, SegmentedList<TagSpan<TTag>> tags)
+        {
+            // we are asked with bunch of spans. rather than asking same question again and again, ask once with big span
+            // which will return superset of what we want. and then filter them out in O(m+n) cost. 
+            // m == number of requested spans, n = number of returned spans
+            var mergedSpan = new SnapshotSpan(requestedSpans[0].Start, requestedSpans[^1].End);
 
-        using var _1 = SegmentedListPool.GetPooledList<TagSpan<TTag>>(out var tempList);
+            using var _1 = SegmentedListPool.GetPooledList<TagSpan<TTag>>(out var tempList);
 
-        AppendIntersectingSpansInSortedOrder(mergedSpan, tempList);
-        if (tempList.Count == 0)
-            return;
+            AddIntersectingTagSpans(mergedSpan, tempList);
+            if (tempList.Count == 0)
+                return;
 
-        // Note: both 'requstedSpans' and 'tempList' are in sorted order.
+            // Note: both 'requestedSpans' and 'tempList' are in sorted order.
 
-        using var enumerator = tempList.GetEnumerator();
+            using var enumerator = tempList.GetEnumerator();
 
-        if (!enumerator.MoveNext())
-            return;
+            if (!enumerator.MoveNext())
+                return;
 
-        using var _2 = PooledHashSet<TagSpan<TTag>>.GetInstance(out var hashSet);
+            using var _2 = PooledHashSet<TagSpan<TTag>>.GetInstance(out var hashSet);
 
-        var requestIndex = 0;
-        while (true)
-        {
-            var currentTag = enumerator.Current;
+            var requestIndex = 0;
+            while (true)
+            {
+                var currentTag = enumerator.Current;
 
-            var currentRequestSpan = requestedSpans[requestIndex];
-            var currentTagSpan = currentTag.Span;
+                var currentRequestSpan = requestedSpans[requestIndex];
+                var currentTagSpan = currentTag.Span;
 
-            // The current tag is *before* the current span we're trying to intersect with.  Move to the next tag to
-            // see if it intersects with the current span.
-            if (currentTagSpan.End < currentRequestSpan.Start)
-            {
-                // If there are no more tags, then we're done.
-                if (!enumerator.MoveNext())
-                    return;
+                // The current tag is *before* the current span we're trying to intersect with.  Move to the next tag to
+                // see if it intersects with the current span.
+                if (currentTagSpan.End < currentRequestSpan.Start)
+                {
+                    // If there are no more tags, then we're done.
+                    if (!enumerator.MoveNext())
+                        return;
 
-                continue;
-            }
+                    continue;
+                }
 
-            // The current tag is *after* teh current span we're trying to intersect with.  Move to the next span to
-            // see if it intersects with the current tag.
-            if (currentTagSpan.Start > currentRequestSpan.End)
-            {
-                requestIndex++;
+                // The current tag is *after* teh current span we're trying to intersect with.  Move to the next span to
+                // see if it intersects with the current tag.
+                if (currentTagSpan.Start > currentRequestSpan.End)
+                {
+                    requestIndex++;
+
+                    // If there are no more spans to intersect with, then we're done.
+                    if (requestIndex >= requestedSpans.Count)
+                        return;
+
+                    continue;
+                }
 
-                // If there are no more spans to intersect with, then we're done.
-                if (requestIndex >= requestedSpans.Count)
-                    return;
+                // This tag intersects the current span we're trying to intersect with.  Ensure we only see and add a
+                // particular tag once. 
 
-                continue;
+                if (currentTagSpan.Length > 0 &&
+                    hashSet.Add(currentTag))
+                {
+                    tags.Add(currentTag);
+                }
+
+                if (!enumerator.MoveNext())
+                    break;
             }
+        }
+    }
+
+    [Conditional("DEBUG")]
+    private static void DebugVerifyTags(NormalizedSnapshotSpanCollection requestedSpans, SegmentedList<TagSpan<TTag>> tags)
+    {
+        if (tags == null)
+        {
+            return;
+        }
 
-            // This tag intersects the current span we're trying to intersect with.  Ensure we only see and add a
-            // particular tag once. 
+        foreach (var tag in tags)
+        {
+            var span = tag.Span;
 
-            if (currentTagSpan.Length > 0 &&
-                hashSet.Add(currentTag))
+            if (!requestedSpans.Any(s => s.IntersectsWith(span)))
             {
-                tags.Add(currentTag);
+                Contract.Fail(tag + " doesn't intersects with any requested span");
             }
-
-            if (!enumerator.MoveNext())
-                break;
         }
     }
 }