Lower allocations in tagging

src/Compilers/Core/Portable/Collections/TemporaryArrayExtensions.cs

@@ -107,6 +107,17 @@ private static void ThrowSequenceContainsMoreThanOneElement()
             return result;
         }
 
+        public static T? FirstOrDefault<T, TArg>(this in TemporaryArray<T> array, Func<T, TArg, bool> predicate, TArg arg)
+        {
+            foreach (var item in array)
+            {
+                if (predicate(item, arg))
+                    return item;
+            }
+
+            return default;
+        }
+
         public static void AddIfNotNull<T>(this ref TemporaryArray<T> array, T? value)
             where T : struct
         {

src/EditorFeatures/Core/Copilot/CopilotTaggerProvider.cs

@@ -14,6 +14,7 @@
 using Microsoft.CodeAnalysis.Editor.Tagging;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Options;
+using Microsoft.CodeAnalysis.Shared.Collections;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.Workspaces;
@@ -58,13 +59,14 @@ protected override ITaggerEventSource CreateEventSource(ITextView textView, ITex
             TaggerEventSources.OnTextChanged(subjectBuffer));
     }
 
-    protected override IEnumerable<SnapshotSpan> GetSpansToTag(ITextView? textView, ITextBuffer subjectBuffer)
+    protected override void AddSpansToTag(ITextView? textView, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
     {
         this.ThreadingContext.ThrowIfNotOnUIThread();
         Contract.ThrowIfNull(textView);
 
         // We only care about the cases where we have caret.
-        return textView.GetCaretPoint(subjectBuffer) is { } caret ? [new SnapshotSpan(caret, 0)] : [];
+        if (textView.GetCaretPoint(subjectBuffer) is { } caret)
+            result.Add(new SnapshotSpan(caret, 0));
     }
 
     protected override async Task ProduceTagsAsync(TaggerContext<ITextMarkerTag> context, DocumentSnapshotSpan spanToTag, int? caretPosition, CancellationToken cancellationToken)

src/EditorFeatures/Core/InlineHints/InlineHintsDataTaggerProvider.cs

@@ -14,6 +14,7 @@
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.InlineHints;
 using Microsoft.CodeAnalysis.Options;
+using Microsoft.CodeAnalysis.Shared.Collections;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.Text.Shared.Extensions;
@@ -77,7 +78,7 @@ protected override ITaggerEventSource CreateEventSource(ITextView textView, ITex
             TaggerEventSources.OnGlobalOptionChanged(GlobalOptions, InlineHintsOptionsStorage.ForImplicitObjectCreation));
     }
 
-    protected override IEnumerable<SnapshotSpan> GetSpansToTag(ITextView? textView, ITextBuffer subjectBuffer)
+    protected override void AddSpansToTag(ITextView? textView, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
     {
         this.ThreadingContext.ThrowIfNotOnUIThread();
         Contract.ThrowIfNull(textView);
@@ -88,10 +89,11 @@ protected override IEnumerable<SnapshotSpan> GetSpansToTag(ITextView? textView,
         if (visibleSpanOpt == null)
         {
             // Couldn't find anything visible, just fall back to tagging all hint locations
-            return base.GetSpansToTag(textView, subjectBuffer);
+            base.AddSpansToTag(textView, subjectBuffer, ref result);
+            return;
         }
 
-        return [visibleSpanOpt.Value];
+        result.Add(visibleSpanOpt.Value);
     }
 
     protected override async Task ProduceTagsAsync(

src/EditorFeatures/Core/ReferenceHighlighting/ReferenceHighlightingViewTaggerProvider.cs

@@ -20,6 +20,7 @@
 using Microsoft.CodeAnalysis.Internal.Log;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.ReferenceHighlighting;
+using Microsoft.CodeAnalysis.Shared.Collections;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.Text.Shared.Extensions;
@@ -86,13 +87,12 @@ protected override ITaggerEventSource CreateEventSource(ITextView textView, ITex
         return textViewOpt.BufferGraph.MapDownToFirstMatch(textViewOpt.Selection.Start.Position, PointTrackingMode.Positive, b => IsSupportedContentType(b.ContentType), PositionAffinity.Successor);
     }
 
-    protected override IEnumerable<SnapshotSpan> GetSpansToTag(ITextView textViewOpt, ITextBuffer subjectBuffer)
+    protected override void AddSpansToTag(ITextView textViewOpt, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
     {
         // Note: this may return no snapshot spans.  We have to be resilient to that
         // when processing the TaggerContext<>.SpansToTag below.
-        return textViewOpt.BufferGraph.GetTextBuffers(b => IsSupportedContentType(b.ContentType))
-                          .Select(b => b.CurrentSnapshot.GetFullSpan())
-                          .ToList();
+        foreach (var buffer in textViewOpt.BufferGraph.GetTextBuffers(b => IsSupportedContentType(b.ContentType)))
+            result.Add(buffer.CurrentSnapshot.GetFullSpan());
     }
 
     protected override Task ProduceTagsAsync(

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

@@ -15,6 +15,7 @@
 using Microsoft.CodeAnalysis.Internal.Log;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.PooledObjects;
+using Microsoft.CodeAnalysis.Shared.Collections;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.CodeAnalysis.Utilities;
 using Microsoft.CodeAnalysis.Workspaces;
@@ -291,7 +292,7 @@ await _visibilityTracker.DelayWhileNonVisibleAsync(
                 // Make a copy of all the data we need while we're on the foreground.  Then switch to a threadpool
                 // thread to do the computation. Finally, once new tags have been computed, then we update our state
                 // again on the foreground.
-                var spansToTag = GetSpansAndDocumentsToTag();
+                var snapshotSpansToTag = GetSnapshotSpansToTag();
                 var caretPosition = _dataSource.GetCaretPoint(_textView, _subjectBuffer);
 
                 // If we're being called from within a blocking JTF.Run call, we don't want to switch to the background
@@ -302,12 +303,9 @@ await _visibilityTracker.DelayWhileNonVisibleAsync(
                 if (cancellationToken.IsCancellationRequested)
                     return null;
 
-                if (frozenPartialSemantics)
-                {
-                    spansToTag = spansToTag.SelectAsArray(ds => new DocumentSnapshotSpan(
-                        ds.Document?.WithFrozenPartialSemantics(cancellationToken),
-                        ds.SnapshotSpan));
-                }
+                // Now that we're on the threadpool, figure out what documents we need to tag corresponding to those
+                // SnapshotSpan the underlying data source asked us to tag.
+                var spansToTag = GetDocumentSnapshotSpansToTag(snapshotSpansToTag, frozenPartialSemantics, cancellationToken);
 
                 // Now spin, trying to compute the updated tags.  We only need to do this as the tag state is also
                 // allowed to change on the UI thread (for example, taggers can say they want tags to be immediately
@@ -352,32 +350,51 @@ await _visibilityTracker.DelayWhileNonVisibleAsync(
 
                 return newTagTrees;
             }
-        }
 
-        private ImmutableArray<DocumentSnapshotSpan> GetSpansAndDocumentsToTag()
-        {
-            _dataSource.ThreadingContext.ThrowIfNotOnUIThread();
-
-            // TODO: Update to tag spans from all related documents.
+            ImmutableArray<SnapshotSpan> GetSnapshotSpansToTag()
+            {
+                _dataSource.ThreadingContext.ThrowIfNotOnUIThread();
 
-            using var _ = PooledDictionary<ITextSnapshot, Document?>.GetInstance(out var snapshotToDocumentMap);
-            var spansToTag = _dataSource.GetSpansToTag(_textView, _subjectBuffer);
+                using var spansToTag = TemporaryArray<SnapshotSpan>.Empty;
+                _dataSource.AddSpansToTag(_textView, _subjectBuffer, ref spansToTag.AsRef());
+                return spansToTag.ToImmutableAndClear();
+            }
 
-            var spansAndDocumentsToTag = spansToTag.SelectAsArray(span =>
+            static ImmutableArray<DocumentSnapshotSpan> GetDocumentSnapshotSpansToTag(
+                ImmutableArray<SnapshotSpan> snapshotSpansToTag,
+                bool frozenPartialSemantics,
+                CancellationToken cancellationToken)
             {
-                if (!snapshotToDocumentMap.TryGetValue(span.Snapshot, out var document))
+                // We only ever have a tiny number of snapshots we're classifying.  So it's easier and faster to just store
+                // the mapping from it to a particular document in an on-stack array.
+                //
+                // document can be null if the buffer the given span is part of is not part of our workspace.
+                using var snapshotToDocument = TemporaryArray<(ITextSnapshot snapshot, Document? document)>.Empty;
+
+                var result = new FixedSizeArrayBuilder<DocumentSnapshotSpan>(snapshotSpansToTag.Length);
+
+                foreach (var spanToTag in snapshotSpansToTag)
                 {
-                    CheckSnapshot(span.Snapshot);
+                    var snapshot = spanToTag.Snapshot;
+                    var document = snapshotToDocument.FirstOrDefault(
+                        static (t, snapshot) => t.snapshot == snapshot, snapshot).document;
 
-                    document = span.Snapshot.GetOpenDocumentInCurrentContextWithChanges();
-                    snapshotToDocumentMap[span.Snapshot] = document;
-                }
+                    if (document is null)
+                    {
+                        CheckSnapshot(snapshot);
 
-                // document can be null if the buffer the given span is part of is not part of our workspace.
-                return new DocumentSnapshotSpan(document, span);
-            });
+                        document = snapshot.GetOpenDocumentInCurrentContextWithChanges();
+                        if (frozenPartialSemantics)
+                            document = document?.WithFrozenPartialSemantics(cancellationToken);
 
-            return spansAndDocumentsToTag;
+                        snapshotToDocument.Add((snapshot, document));
+                    }
+
+                    result.Add(new DocumentSnapshotSpan(document, spanToTag));
+                }
+
+                return result.MoveToImmutable();
+            }
         }
 
         [Conditional("DEBUG")]
@@ -495,12 +512,24 @@ private bool ShouldSkipTagProduction()
             return _dataSource.Options.OfType<PerLanguageOption2<bool>>().Any(option => languageName == null || !_dataSource.GlobalOptions.GetOption(option, languageName));
         }
 
-        private Task ProduceTagsAsync(TaggerContext<TTag> context, CancellationToken cancellationToken)
+        private async ValueTask ProduceTagsAsync(TaggerContext<TTag> context, CancellationToken cancellationToken)
         {
             // If the feature is disabled, then just produce no tags.
-            return ShouldSkipTagProduction()
-                ? Task.CompletedTask
-                : _dataSource.ProduceTagsAsync(context, cancellationToken);
+            if (_dataSource.Options.OfType<Option2<bool>>().Any(option => !_dataSource.GlobalOptions.GetOption(option)))
+                return;
+
+            var languageName = _subjectBuffer.GetLanguageName();
+            if (_dataSource.Options.OfType<PerLanguageOption2<bool>>().Any(
+                    option => languageName == null || !_dataSource.GlobalOptions.GetOption(option, languageName)))
+            {
+                return;
+            }
+
+            // If we have no spans to tag, there's no point in continuing.
+            if (context.SpansToTag.IsEmpty)
+                return;
+
+            await _dataSource.ProduceTagsAsync(context, cancellationToken).ConfigureAwait(false);
         }
 
         private Dictionary<ITextBuffer, DiffResult> ProcessNewTagTrees(

src/EditorFeatures/Core/Tagging/AbstractAsynchronousTaggerProvider.cs

@@ -15,6 +15,7 @@
 using Microsoft.CodeAnalysis.Editor.Shared.Options;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Options;
+using Microsoft.CodeAnalysis.Shared.Collections;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.Tagging;
 using Microsoft.CodeAnalysis.Text;
@@ -214,10 +215,11 @@ private void StoreTagSource(ITextView? textView, ITextBuffer subjectBuffer, TagS
     /// and will asynchronously call into <see cref="ProduceTagsAsync(TaggerContext{TTag}, CancellationToken)"/> at some point in
     /// the future to produce tags for these spans.
     /// </summary>
-    protected virtual IEnumerable<SnapshotSpan> GetSpansToTag(ITextView? textView, ITextBuffer subjectBuffer)
+    protected virtual void AddSpansToTag(
+        ITextView? textView, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
     {
         // For a standard tagger, the spans to tag is the span of the entire snapshot.
-        return [subjectBuffer.CurrentSnapshot.GetFullSpan()];
+        result.Add(subjectBuffer.CurrentSnapshot.GetFullSpan());
     }
 
     /// <summary>

src/EditorFeatures/Core/Tagging/AsynchronousViewportTaggerProvider.SingleViewportTaggerProvider.cs

@@ -59,7 +59,7 @@ protected override TaggerDelay EventChangeDelay
             // looked at.
             => _viewPortToTag == ViewPortToTag.InView ? _callback.EventChangeDelay : TaggerDelay.NonFocus;
 
-        protected override IEnumerable<SnapshotSpan> GetSpansToTag(ITextView? textView, ITextBuffer subjectBuffer)
+        protected override void AddSpansToTag(ITextView? textView, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
         {
             this.ThreadingContext.ThrowIfNotOnUIThread();
             Contract.ThrowIfNull(textView);
@@ -71,16 +71,20 @@ protected override IEnumerable<SnapshotSpan> GetSpansToTag(ITextView? textView,
             {
                 // couldn't figure out the visible span.  So the InView tagger will need to tag everything, and the
                 // above/below tagger should tag nothing.
-                return _viewPortToTag == ViewPortToTag.InView
-                    ? base.GetSpansToTag(textView, subjectBuffer)
-                    : [];
+                if (_viewPortToTag == ViewPortToTag.InView)
+                    base.AddSpansToTag(textView, subjectBuffer, ref result);
+
+                return;
             }
 
             var visibleSpan = visibleSpanOpt.Value;
 
             // If we're the 'InView' tagger, tag what was visible. 
             if (_viewPortToTag is ViewPortToTag.InView)
-                return [visibleSpan];
+            {
+                result.Add(visibleSpan);
+                return;
+            }
 
             // For the above/below tagger, broaden the span to to the requested portion above/below what's visible, then
             // subtract out the visible range.
@@ -90,8 +94,6 @@ protected override IEnumerable<SnapshotSpan> GetSpansToTag(ITextView? textView,
             var widenedSpan = widenedSpanOpt.Value;
             Contract.ThrowIfFalse(widenedSpan.Span.Contains(visibleSpan.Span), "The widened span must be at least as large as the visible one.");
 
-            using var result = TemporaryArray<SnapshotSpan>.Empty;
-
             if (_viewPortToTag is ViewPortToTag.Above)
             {
                 var aboveSpan = new SnapshotSpan(visibleSpan.Snapshot, Span.FromBounds(widenedSpan.Span.Start, visibleSpan.Span.Start));
@@ -101,12 +103,9 @@ protected override IEnumerable<SnapshotSpan> GetSpansToTag(ITextView? textView,
             else if (_viewPortToTag is ViewPortToTag.Below)
             {
                 var belowSpan = new SnapshotSpan(visibleSpan.Snapshot, Span.FromBounds(visibleSpan.Span.End, widenedSpan.Span.End));
-
                 if (!belowSpan.IsEmpty)
                     result.Add(belowSpan);
             }
-
-            return result.ToImmutableAndClear();
         }
 
         protected override async Task ProduceTagsAsync(