Merge pull request #73729 from CyrusNajmabadi/taggerAllocs

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());
+        Contract.ThrowIfTrue(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,9 +292,14 @@ 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 DEBUG
+                foreach (var snapshotSpan in snapshotSpansToTag)
+                    CheckSnapshot(snapshotSpan.Snapshot);
+#endif
+
                 // If we're being called from within a blocking JTF.Run call, we don't want to switch to the background
                 // if we can avoid it.
                 if (!calledFromJtfRun)
@@ -302,12 +308,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
@@ -321,7 +324,7 @@ await _visibilityTracker.DelayWhileNonVisibleAsync(
                     {
                         // Create a context to store pass the information along and collect the results.
                         context = new TaggerContext<TTag>(
-                            oldState, frozenPartialSemantics, spansToTag, caretPosition, oldTagTrees);
+                            oldState, frozenPartialSemantics, spansToTag, snapshotSpansToTag, caretPosition, oldTagTrees);
                         await ProduceTagsAsync(context, cancellationToken).ConfigureAwait(false);
 
                         return ComputeNewTagTrees(oldTagTrees, context);
@@ -352,46 +355,68 @@ 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 (foundSnapshot, document) = snapshotToDocument.FirstOrDefault(
+                        static (t, snapshot) => t.snapshot == snapshot, snapshot);
 
-                    document = span.Snapshot.GetOpenDocumentInCurrentContextWithChanges();
-                    snapshotToDocumentMap[span.Snapshot] = document;
-                }
+                    // If this is the first time looking at this snapshot, then go fetch the document (which we may or
+                    // may not have), and freeze it if necessary..
+                    if (foundSnapshot is null)
+                    {
+                        document = snapshot.GetOpenDocumentInCurrentContextWithChanges();
+                        if (frozenPartialSemantics)
+                            document = document?.WithFrozenPartialSemantics(cancellationToken);
 
-                // document can be null if the buffer the given span is part of is not part of our workspace.
-                return new DocumentSnapshotSpan(document, span);
-            });
+                        snapshotToDocument.Add((snapshot, document));
+                    }
 
-            return spansAndDocumentsToTag;
-        }
+                    result.Add(new DocumentSnapshotSpan(document, spanToTag));
+                }
 
-        [Conditional("DEBUG")]
-        private static void CheckSnapshot(ITextSnapshot snapshot)
-        {
-            var container = snapshot.TextBuffer.AsTextContainer();
-            if (Workspace.TryGetWorkspace(container, out _))
+                return result.ToOneOrManyAndClear();
+            }
+
+#if DEBUG
+            static void CheckSnapshot(ITextSnapshot snapshot)
             {
-                // if the buffer is part of our workspace, it must be the latest.
-                Debug.Assert(snapshot.Version.Next == null, "should be on latest snapshot");
+                var container = snapshot.TextBuffer.AsTextContainer();
+                if (Workspace.TryGetWorkspace(container, out _))
+                {
+                    // if the buffer is part of our workspace, it must be the latest.
+                    Debug.Assert(snapshot.Version.Next == null, "should be on latest snapshot");
+                }
             }
+#endif
         }
 
-        private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTagTrees(ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> oldTagTrees, TaggerContext<TTag> context)
+        private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTagTrees(
+            ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> oldTagTrees,
+            TaggerContext<TTag> context)
         {
             using var _1 = PooledHashSet<ITextBuffer>.GetInstance(out var buffersToTag);
             foreach (var spanToTag in context.SpansToTag)
@@ -414,6 +439,8 @@ private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTa
                         newTagsInBuffer.Add(tagSpan);
                 }
 
+                // Invalidate all the spans that were actually tagged.  If the context doesn't have any recorded spans
+                // that were tagged, then assume we tagged everything we were asked to tag.
                 foreach (var span in context._spansTagged)
                 {
                     if (span.Snapshot.TextBuffer == buffer)
@@ -486,25 +513,31 @@ private ImmutableDictionary<ITextBuffer, TagSpanIntervalTree<TTag>> ComputeNewTa
             }
         }
 
-        private bool ShouldSkipTagProduction()
+        private async ValueTask ProduceTagsAsync(TaggerContext<TTag> context, CancellationToken cancellationToken)
         {
-            if (_dataSource.Options.OfType<Option2<bool>>().Any(option => !_dataSource.GlobalOptions.GetOption(option)))
-                return true;
+            // If we have no spans to tag, there's no point in continuing.
+            if (context.SpansToTag.IsEmpty)
+                return;
 
+            // If the feature is disabled, then just produce no tags.
             var languageName = _subjectBuffer.GetLanguageName();
-            return _dataSource.Options.OfType<PerLanguageOption2<bool>>().Any(option => languageName == null || !_dataSource.GlobalOptions.GetOption(option, languageName));
-        }
+            foreach (var option in _dataSource.Options)
+            {
+                if (option is Option2<bool> option2 && !_dataSource.GlobalOptions.GetOption(option2))
+                    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 (option is PerLanguageOption2<bool> perLanguageOption &&
+                    (languageName == null || !_dataSource.GlobalOptions.GetOption(perLanguageOption, languageName)))
+                {
+                    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 +547,10 @@ 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 = spansToTag.FirstOrDefault(
+                        static (span, latestBuffer) => span.SnapshotSpan.Snapshot.TextBuffer == latestBuffer,
+                        latestBuffer).SnapshotSpan.Snapshot;
+                    Contract.ThrowIfNull(snapshot);
 
                     if (oldTagTrees.TryGetValue(latestBuffer, out var previousSpans))
                     {