Fix issue classifying while performing a layout

src/EditorFeatures/Core/Copilot/CopilotTaggerProvider.cs

@@ -48,14 +48,18 @@ protected override ITaggerEventSource CreateEventSource(ITextView textView, ITex
             TaggerEventSources.OnTextChanged(subjectBuffer));
     }
 
-    protected override void AddSpansToTag(ITextView? textView, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
+    protected override bool TryAddSpansToTag(ITextView? textView, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
     {
         this.ThreadingContext.ThrowIfNotOnUIThread();
         Contract.ThrowIfNull(textView);
 
         // We only care about the cases where we have caret.
         if (textView.GetCaretPoint(subjectBuffer) is { } caret)
             result.Add(new SnapshotSpan(caret, 0));
+
+        // Note: we report 'true' unconditionally here.  This is because we want to ensure that we still compute 'no tags'
+        // in the case that we don't have a caret point, as opposed to bailing out.
+        return true;
     }
 
     protected override async Task ProduceTagsAsync(TaggerContext<ITextMarkerTag> context, DocumentSnapshotSpan spanToTag, int? caretPosition, CancellationToken cancellationToken)

src/EditorFeatures/Core/InlineHints/InlineHintsDataTaggerProvider.cs

@@ -74,22 +74,25 @@ protected override ITaggerEventSource CreateEventSource(ITextView textView, ITex
                 option.Equals(InlineHintsOptionsStorage.ForCollectionExpressions)));
     }
 
-    protected override void AddSpansToTag(ITextView? textView, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
+    protected override bool TryAddSpansToTag(ITextView? textView, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
     {
         this.ThreadingContext.ThrowIfNotOnUIThread();
         Contract.ThrowIfNull(textView);
 
-        // Find the visible span some 100 lines +/- what's actually in view.  This way
-        // if the user scrolls up/down, we'll already have the results.
+        // Find the visible span some 100 lines +/- what's actually in view.  This way if the user scrolls up/down,
+        // we'll already have the results.
         var visibleSpanOpt = textView.GetVisibleLinesSpan(subjectBuffer, extraLines: 100);
+
+        // If we can't even determine what our visible span is, bail out immediately.  We may be in something like a
+        // layout pass, and we don't want to suddenly flip to tagging everything, then go back to tagging a small
+        // subset of the view afterwards.
+        //
+        // In this case we literally do not know what is visible, so we want to bail and try again later.
         if (visibleSpanOpt == null)
-        {
-            // Couldn't find anything visible, just fall back to tagging all hint locations
-            base.AddSpansToTag(textView, subjectBuffer, ref result);
-            return;
-        }
+            return false;
 
         result.Add(visibleSpanOpt.Value);
+        return true;
     }
 
     protected override async Task ProduceTagsAsync(

src/EditorFeatures/Core/ReferenceHighlighting/ReferenceHighlightingViewTaggerProvider.cs

@@ -77,12 +77,16 @@ 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 void AddSpansToTag(ITextView textViewOpt, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
+    protected override bool TryAddSpansToTag(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.
         foreach (var buffer in textViewOpt.BufferGraph.GetTextBuffers(b => IsSupportedContentType(b.ContentType)))
             result.Add(buffer.CurrentSnapshot.GetFullSpan());
+
+        // Fine to always return 'true' here.  If we get no spans to tag, we can move ourselves to the no-tags state and
+        // update the editor.
+        return true;
     }
 
     protected override Task ProduceTagsAsync(

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

@@ -294,9 +294,19 @@ private async ValueTask<VoidResult> ProcessEventChangeAsync(
             // us avoid hammering the dispatcher queue with lots of work that causes contention.  Additionally, use
             // a no-throw awaitable so that in the common case where we cancel before, we don't throw an exception
             // that can exacerbate cross process debugging scenarios.
-            var (isVisible, caretPosition, snapshotSpansToTag) = await _dataSource.MainThreadManager.PerformWorkOnMainThreadAsync(
+            var valueOpt = await _dataSource.MainThreadManager.PerformWorkOnMainThreadAsync(
                 GetTaggerUIData, cancellationToken).ConfigureAwait(true);
 
+            if (valueOpt is null)
+            {
+                // We failed to get the UI data we need.  This can happen in cases like trying to get that during a layout pass.
+                // In that case, we just bail out and try again later.
+                this.EnqueueWork(highPriority);
+                return null;
+            }
+
+            var (isVisible, caretPosition, snapshotSpansToTag) = valueOpt.Value;
+
             // Since we don't ever throw above, check and see if the await completed due to cancellation and do not
             // proceed.
             if (cancellationToken.IsCancellationRequested)
@@ -379,7 +389,9 @@ static async (oldTagTrees, args, cancellationToken) =>
                 return newTagTrees;
             }
 
-            (bool isVisible, SnapshotPoint? caretPosition, OneOrMany<SnapshotSpan> spansToTag) GetTaggerUIData()
+            // Returns null if we we're currently in a state where we can't get the tags to span and we should bail out
+            // from producing tags on this turn of the crank.
+            (bool isVisible, SnapshotPoint? caretPosition, OneOrMany<SnapshotSpan> spansToTag)? GetTaggerUIData()
             {
                 _dataSource.ThreadingContext.ThrowIfNotOnUIThread();
 
@@ -393,7 +405,8 @@ static async (oldTagTrees, args, cancellationToken) =>
                 var caretPosition = _dataSource.GetCaretPoint(_textView, _subjectBuffer);
 
                 using var spansToTag = TemporaryArray<SnapshotSpan>.Empty;
-                _dataSource.AddSpansToTag(_textView, _subjectBuffer, ref spansToTag.AsRef());
+                if (!_dataSource.TryAddSpansToTag(_textView, _subjectBuffer, ref spansToTag.AsRef()))
+                    return null;
 
 #if DEBUG
                 foreach (var snapshotSpan in spansToTag)

src/EditorFeatures/Core/Tagging/AbstractAsynchronousTaggerProvider.cs

@@ -215,11 +215,14 @@ 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 void AddSpansToTag(
+    /// <returns><see langword="false"/> if spans could not be added and if tagging should abort and wait until its next
+    /// run.</returns>
+    protected virtual bool TryAddSpansToTag(
         ITextView? textView, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
     {
         // For a standard tagger, the spans to tag is the span of the entire snapshot.
         result.Add(subjectBuffer.CurrentSnapshot.GetFullSpan());
+        return true;
     }
 
     /// <summary>

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

@@ -62,53 +62,57 @@ protected override bool CancelOnNewWork
             // This can save a lot of CPU time for things that may never even be looked at.
             => _viewPortToTag != ViewPortToTag.InView;
 
-        protected override void AddSpansToTag(ITextView? textView, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
+        protected override bool TryAddSpansToTag(ITextView? textView, ITextBuffer subjectBuffer, ref TemporaryArray<SnapshotSpan> result)
         {
             this.ThreadingContext.ThrowIfNotOnUIThread();
             Contract.ThrowIfNull(textView);
 
             // if we're the current view, attempt to just get what's visible, plus 10 lines above and below.  This will
             // ensure that moving up/down a few lines tends to have immediate accurate results.
             var visibleSpanOpt = textView.GetVisibleLinesSpan(subjectBuffer, extraLines: s_standardLineCountAroundViewportToTag);
-            if (visibleSpanOpt is null)
-            {
-                // couldn't figure out the visible span.  So the InView tagger will need to tag everything, and the
-                // above/below tagger should tag nothing.
-                if (_viewPortToTag == ViewPortToTag.InView)
-                    base.AddSpansToTag(textView, subjectBuffer, ref result);
 
-                return;
-            }
+            // If we can't even determine what our visible span is, bail out immediately.  We may be in something like a
+            // layout pass, and we don't want to suddenly flip to tagging everything, then go back to tagging a small
+            // subset of the view afterwards.
+            //
+            // In this case we literally do not know what is visible, so we want to bail and try again later.
+            if (visibleSpanOpt is null)
+                return false;
 
             var visibleSpan = visibleSpanOpt.Value;
 
             // If we're the 'InView' tagger, tag what was visible. 
             if (_viewPortToTag is ViewPortToTag.InView)
             {
                 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.
-            var widenedSpanOpt = textView.GetVisibleLinesSpan(subjectBuffer, extraLines: _callback._extraLinesAroundViewportToTag);
-            Contract.ThrowIfNull(widenedSpanOpt, "Should not ever fail getting the widened span as we were able to get the normal visible span");
-
-            var widenedSpan = widenedSpanOpt.Value;
-            Contract.ThrowIfFalse(widenedSpan.Span.Contains(visibleSpan.Span), "The widened span must be at least as large as the visible one.");
-
-            if (_viewPortToTag is ViewPortToTag.Above)
-            {
-                var aboveSpan = new SnapshotSpan(visibleSpan.Snapshot, Span.FromBounds(widenedSpan.Span.Start, visibleSpan.Span.Start));
-                if (!aboveSpan.IsEmpty)
-                    result.Add(aboveSpan);
             }
-            else if (_viewPortToTag is ViewPortToTag.Below)
+            else
             {
-                var belowSpan = new SnapshotSpan(visibleSpan.Snapshot, Span.FromBounds(visibleSpan.Span.End, widenedSpan.Span.End));
-                if (!belowSpan.IsEmpty)
-                    result.Add(belowSpan);
+                // For the above/below tagger, broaden the span to to the requested portion above/below what's visible, then
+                // subtract out the visible range.
+                var widenedSpanOpt = textView.GetVisibleLinesSpan(subjectBuffer, extraLines: _callback._extraLinesAroundViewportToTag);
+                Contract.ThrowIfNull(widenedSpanOpt, "Should not ever fail getting the widened span as we were able to get the normal visible span");
+
+                var widenedSpan = widenedSpanOpt.Value;
+                Contract.ThrowIfFalse(widenedSpan.Span.Contains(visibleSpan.Span), "The widened span must be at least as large as the visible one.");
+
+                if (_viewPortToTag is ViewPortToTag.Above)
+                {
+                    var aboveSpan = new SnapshotSpan(visibleSpan.Snapshot, Span.FromBounds(widenedSpan.Span.Start, visibleSpan.Span.Start));
+                    if (!aboveSpan.IsEmpty)
+                        result.Add(aboveSpan);
+                }
+                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);
+                }
             }
+
+            // Unilaterally return true here, even if we determine we don't have a span to tag.  In this case, we've
+            // computed that there really is nothing visible, in which case we *do* want to move to having no tags.
+            return true;
         }
 
         protected override async Task ProduceTagsAsync(

src/EditorFeatures/Core/Tagging/TaggerMainThreadManager.cs

@@ -14,7 +14,7 @@
 
 namespace Microsoft.CodeAnalysis.Editor.Tagging;
 
-using QueueData = (Func<TaggerUIData> action, TaskCompletionSource<TaggerUIData> taskCompletionSource, CancellationToken cancellationToken);
+using QueueData = (Func<TaggerUIData?> action, TaskCompletionSource<TaggerUIData?> taskCompletionSource, CancellationToken cancellationToken);
 
 internal sealed class TaggerMainThreadManager
 {
@@ -36,8 +36,8 @@ public TaggerMainThreadManager(
 
     /// <remarks>This will not ever throw.</remarks>
     private static void RunActionAndUpdateCompletionSource_NoThrow(
-        Func<TaggerUIData> action,
-        TaskCompletionSource<TaggerUIData> taskCompletionSource,
+        Func<TaggerUIData?> action,
+        TaskCompletionSource<TaggerUIData?> taskCompletionSource,
         CancellationToken cancellationToken)
     {
         try
@@ -60,7 +60,7 @@ private static void RunActionAndUpdateCompletionSource_NoThrow(
         }
         finally
         {
-            taskCompletionSource.TrySetResult(default);
+            taskCompletionSource.TrySetResult(null);
             Contract.ThrowIfFalse(taskCompletionSource.Task.IsCompleted);
         }
     }
@@ -69,9 +69,9 @@ private static void RunActionAndUpdateCompletionSource_NoThrow(
     /// Adds the provided action to a queue that will run on the UI thread in the near future (batched with other
     /// registered actions).  If the cancellation token is triggered before the action runs, it will not be run.
     /// </summary>
-    public async ValueTask<TaggerUIData> PerformWorkOnMainThreadAsync(Func<TaggerUIData> action, CancellationToken cancellationToken)
+    public async ValueTask<TaggerUIData?> PerformWorkOnMainThreadAsync(Func<TaggerUIData?> action, CancellationToken cancellationToken)
     {
-        var taskSource = new TaskCompletionSource<TaggerUIData>();
+        var taskSource = new TaskCompletionSource<TaggerUIData?>();
 
         // If we're already on the main thread, just run the action directly without any delay.  This is important
         // for cases where the tagger is performing a blocking call to get tags synchronously on the UI thread (for
@@ -84,7 +84,7 @@ public async ValueTask<TaggerUIData> PerformWorkOnMainThreadAsync(Func<TaggerUID
         {
             // Ensure that if the host is closing and hte queue stops running that we transition this task to the canceled state.
             var registration = _threadingContext.DisposalToken.Register(
-                static taskSourceObj => ((TaskCompletionSource<TaggerUIData>)taskSourceObj!).TrySetCanceled(), taskSource);
+                static taskSourceObj => ((TaskCompletionSource<TaggerUIData?>)taskSourceObj!).TrySetCanceled(), taskSource);
 
             _workQueue.AddWork((action, taskSource, cancellationToken));