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/Compilers/Core/Portable/Collections/TemporaryArray`1.cs

@@ -265,11 +265,29 @@ public readonly Enumerator GetEnumerator()
             return new Enumerator(in this);
         }
 
+        /// <summary>
+        /// Create an <see cref="OneOrMany{T}"/> with the elements currently held in the temporary array, and clear the
+        /// array.
+        /// </summary>
+        public OneOrMany<T> ToOneOrManyAndClear()
+        {
+            switch (this.Count)
+            {
+                case 0:
+                    return OneOrMany<T>.Empty;
+                case 1:
+                    var result = OneOrMany.Create(this[0]);
+                    this.Clear();
+                    return result;
+                default:
+                    return new(this.ToImmutableAndClear());
+            }
+        }
+
         /// <summary>
         /// Create an <see cref="ImmutableArray{T}"/> with the elements currently held in the temporary array, and clear
         /// the array.
         /// </summary>
-        /// <returns></returns>
         public ImmutableArray<T> ToImmutableAndClear()
         {
             if (_builder is not null)

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/EditAndContinue/ActiveStatementTaggerProvider.cs

@@ -55,9 +55,9 @@ protected override ITaggerEventSource CreateEventSource(ITextView? textView, ITe
     protected override async Task ProduceTagsAsync(
         TaggerContext<ITextMarkerTag> context, CancellationToken cancellationToken)
     {
-        Debug.Assert(context.SpansToTag.IsSingle());
+        Debug.Assert(context.SpansToTag.Count == 1);
 
-        var spanToTag = context.SpansToTag.Single();
+        var spanToTag = context.SpansToTag.First();
 
         var document = spanToTag.Document;
         if (document == null)

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(
@@ -145,8 +145,6 @@ private static async Task ProduceTagsAsync(
         HighlightingOptions options,
         CancellationToken cancellationToken)
     {
-        var solution = document.Project.Solution;
-
         using (Logger.LogBlock(FunctionId.Tagger_ReferenceHighlighting_TagProducer_ProduceTags, cancellationToken))
         {
             if (document != null)
@@ -160,15 +158,17 @@ private static async Task ProduceTagsAsync(
 
                     // We only want to search inside documents that correspond to the snapshots
                     // we're looking at
-                    var documentsToSearch = ImmutableHashSet.CreateRange(context.SpansToTag.Select(vt => vt.Document).WhereNotNull());
+                    var documentsToSearchBuilder = ImmutableHashSet.CreateBuilder<Document>();
+                    foreach (var snapshotSpan in context.SpansToTag)
+                        documentsToSearchBuilder.AddIfNotNull(snapshotSpan.Document);
+
+                    var documentsToSearch = documentsToSearchBuilder.ToImmutable();
                     var documentHighlightsList = await service.GetDocumentHighlightsAsync(
                         document, position, documentsToSearch, options, cancellationToken).ConfigureAwait(false);
                     if (documentHighlightsList != null)
                     {
                         foreach (var documentHighlights in documentHighlightsList)
-                        {
                             AddTagSpans(context, documentHighlights, cancellationToken);
-                        }
                     }
                 }
             }

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.
+            OneOrMany<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.ToOneOrManyAndClear();
+            }
 
-            var spansAndDocumentsToTag = spansToTag.SelectAsArray(span =>
+            static OneOrMany<DocumentSnapshotSpan> GetDocumentSnapshotSpansToTag(
+                OneOrMany<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;
+
+                using var result = TemporaryArray<DocumentSnapshotSpan>.Empty;
+
+                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);
+
+                        snapshotToDocument.Add((snapshot, document));
+                    }
 
-            return spansAndDocumentsToTag;
+                    result.Add(new DocumentSnapshotSpan(document, spanToTag));
+                }
+
+                return result.ToOneOrManyAndClear();
+            }
         }
 
         [Conditional("DEBUG")]
@@ -486,25 +503,28 @@ private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTa
             }
         }
 
-        private bool ShouldSkipTagProduction()
+        private async ValueTask ProduceTagsAsync(TaggerContext<TTag> context, CancellationToken cancellationToken)
         {
+            // If the feature is disabled, then just produce no tags.
             if (_dataSource.Options.OfType<Option2<bool>>().Any(option => !_dataSource.GlobalOptions.GetOption(option)))
-                return true;
+                return;
 
             var languageName = _subjectBuffer.GetLanguageName();
-            return _dataSource.Options.OfType<PerLanguageOption2<bool>>().Any(option => languageName == null || !_dataSource.GlobalOptions.GetOption(option, languageName));
-        }
+            if (_dataSource.Options.OfType<PerLanguageOption2<bool>>().Any(
+                    option => languageName == null || !_dataSource.GlobalOptions.GetOption(option, languageName)))
+            {
+                return;
+            }
 
-        private Task 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 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(
-            ImmutableArray<DocumentSnapshotSpan> spansToTag,
+            OneOrMany<DocumentSnapshotSpan> spansToTag,
             ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> oldTagTrees,
             ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> newTagTrees)
         {
@@ -514,7 +534,7 @@ private Dictionary<ITextBuffer, DiffResult> ProcessNewTagTrees(
 
                 foreach (var (latestBuffer, latestSpans) in newTagTrees)
                 {
-                    var snapshot = spansToTag.First(s => s.SnapshotSpan.Snapshot.TextBuffer == latestBuffer).SnapshotSpan.Snapshot;
+                    var snapshot = GetSpanForBuffer(spansToTag, latestBuffer).SnapshotSpan.Snapshot;
 
                     if (oldTagTrees.TryGetValue(latestBuffer, out var previousSpans))
                     {
@@ -539,6 +559,19 @@ private Dictionary<ITextBuffer, DiffResult> ProcessNewTagTrees(
 
                 return bufferToChanges;
             }
+
+            static DocumentSnapshotSpan GetSpanForBuffer(
+                OneOrMany<DocumentSnapshotSpan> spansToTag,
+                ITextBuffer latestBuffer)
+            {
+                foreach (var span in spansToTag)
+                {
+                    if (span.SnapshotSpan.Snapshot.TextBuffer == latestBuffer)
+                        return span;
+                }
+
+                throw ExceptionUtilities.Unreachable();
+            }
         }
 
         /// <summary>

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>