Merge pull request #73599 from CyrusNajmabadi/referenceArray

src/EditorFeatures/Core/ExternalAccess/VSTypeScript/VSTypeScriptFindUsagesContext.cs

@@ -26,7 +26,7 @@ public ValueTask OnDefinitionFoundAsync(VSTypeScriptDefinitionItem definition, C
         => UnderlyingObject.OnDefinitionFoundAsync(definition.UnderlyingObject, cancellationToken);
 
     public ValueTask OnReferenceFoundAsync(VSTypeScriptSourceReferenceItem reference, CancellationToken cancellationToken)
-        => UnderlyingObject.OnReferenceFoundAsync(reference.UnderlyingObject, cancellationToken);
+        => UnderlyingObject.OnReferencesFoundAsync([reference.UnderlyingObject], cancellationToken);
 
     public ValueTask OnCompletedAsync(CancellationToken cancellationToken)
         => UnderlyingObject.OnCompletedAsync(cancellationToken);

src/EditorFeatures/Core/ExternalAccess/VSTypeScript/VSTypeScriptFindUsagesService.cs

@@ -46,7 +46,7 @@ public ValueTask OnDefinitionFoundAsync(VSTypeScriptDefinitionItem definition, C
             => _context.OnDefinitionFoundAsync(definition.UnderlyingObject, cancellationToken);
 
         public ValueTask OnReferenceFoundAsync(VSTypeScriptSourceReferenceItem reference, CancellationToken cancellationToken)
-            => _context.OnReferenceFoundAsync(reference.UnderlyingObject, cancellationToken);
+            => _context.OnReferencesFoundAsync([reference.UnderlyingObject], cancellationToken);
 
         public ValueTask OnCompletedAsync(CancellationToken cancellationToken)
             => ValueTaskFactory.CompletedTask;

src/EditorFeatures/Core/FindUsages/BufferedFindUsagesContext.cs

@@ -201,7 +201,7 @@ async ValueTask IFindUsagesContext.OnDefinitionFoundAsync(DefinitionItem definit
         }
     }
 
-    ValueTask IFindUsagesContext.OnReferenceFoundAsync(SourceReferenceItem reference, CancellationToken cancellationToken)
+    ValueTask IFindUsagesContext.OnReferencesFoundAsync(ImmutableArray<SourceReferenceItem> references, CancellationToken cancellationToken)
     {
         // Entirely ignored.  These features do not show references.
         Contract.Fail("GoToImpl/Base should never report a reference.");

src/EditorFeatures/Core/Host/IStreamingFindReferencesPresenter.cs

@@ -23,7 +23,7 @@ internal interface IStreamingFindUsagesPresenter
     /// <summary>
     /// Tells the presenter that a search is starting.  The returned <see cref="FindUsagesContext"/>
     /// is used to push information about the search into.  i.e. when a reference is found
-    /// <see cref="FindUsagesContext.OnReferenceFoundAsync"/> should be called.  When the
+    /// <see cref="FindUsagesContext.OnReferencesFoundAsync"/> should be called.  When the
     /// search completes <see cref="FindUsagesContext.OnCompletedAsync"/> should be called. 
     /// etc. etc.
     /// </summary>

src/EditorFeatures/Test2/FindReferences/FindReferencesTests.vb

@@ -249,9 +249,9 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.FindReferences
                 Return Nothing
             End Function
 
-            Public Overrides Function OnReferenceFoundAsync(reference As SourceReferenceItem, cancellationToken As CancellationToken) As ValueTask
+            Public Overrides Function OnReferencesFoundAsync(references As ImmutableArray(Of SourceReferenceItem), cancellationToken As CancellationToken) As ValueTask
                 SyncLock gate
-                    References.Add(reference)
+                    Me.References.AddRange(references)
                 End SyncLock
 
                 Return Nothing

src/Features/Core/Portable/FindUsages/AbstractFindUsagesService.DefinitionTrackingContext.cs

@@ -6,9 +6,6 @@
 using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.CodeAnalysis.Classification;
-using Microsoft.CodeAnalysis.FindUsages;
-using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Notification;
 using Microsoft.CodeAnalysis.Shared.Utilities;
 
@@ -42,8 +39,8 @@ public ValueTask ReportMessageAsync(string message, NotificationSeverity severit
         public ValueTask SetSearchTitleAsync(string title, CancellationToken cancellationToken)
             => _underlyingContext.SetSearchTitleAsync(title, cancellationToken);
 
-        public ValueTask OnReferenceFoundAsync(SourceReferenceItem reference, CancellationToken cancellationToken)
-            => _underlyingContext.OnReferenceFoundAsync(reference, cancellationToken);
+        public ValueTask OnReferencesFoundAsync(ImmutableArray<SourceReferenceItem> references, CancellationToken cancellationToken)
+            => _underlyingContext.OnReferencesFoundAsync(references, cancellationToken);
 
         public ValueTask OnDefinitionFoundAsync(DefinitionItem definition, CancellationToken cancellationToken)
         {

src/Features/Core/Portable/FindUsages/AbstractFindUsagesService.ProgressAdapter.cs

@@ -13,6 +13,7 @@
 using Microsoft.CodeAnalysis.FindUsages;
 using Microsoft.CodeAnalysis.Navigation;
 using Microsoft.CodeAnalysis.Options;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Utilities;
 using Microsoft.CodeAnalysis.Text;
 using Roslyn.Utilities;
@@ -42,16 +43,17 @@ public async ValueTask OnReferenceFoundAsync(Document document, TextSpan span, C
             var classifiedSpans = await ClassifiedSpansAndHighlightSpanFactory.ClassifyAsync(
                 documentSpan, classifiedSpans: null, options, cancellationToken).ConfigureAwait(false);
 
-            await _context.OnReferenceFoundAsync(
-                new SourceReferenceItem(_definition, documentSpan, classifiedSpans, SymbolUsageInfo.None), cancellationToken).ConfigureAwait(false);
+            await _context.OnReferencesFoundAsync(
+                [new SourceReferenceItem(_definition, documentSpan, classifiedSpans, SymbolUsageInfo.None)], cancellationToken).ConfigureAwait(false);
         }
     }
 
     /// <summary>
     /// Forwards IFindReferencesProgress calls to an IFindUsagesContext instance.
     /// </summary>
     private sealed class FindReferencesProgressAdapter(
-        Solution solution, IFindUsagesContext context, FindReferencesSearchOptions searchOptions, OptionsProvider<ClassificationOptions> classificationOptions) : IStreamingFindReferencesProgress
+        Solution solution, IFindUsagesContext context, FindReferencesSearchOptions searchOptions, OptionsProvider<ClassificationOptions> classificationOptions)
+        : IStreamingFindReferencesProgress
     {
         /// <summary>
         /// We will hear about definition symbols many times while performing FAR.  We'll
@@ -107,20 +109,26 @@ public async ValueTask OnDefinitionFoundAsync(SymbolGroup group, CancellationTok
         }
 
         public async ValueTask OnReferencesFoundAsync(
-            ImmutableArray<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken)
+            IAsyncEnumerable<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken)
         {
-            foreach (var (group, _, location) in references)
-            {
-                var definitionItem = await GetDefinitionItemAsync(group, cancellationToken).ConfigureAwait(false);
-                var referenceItem = await location.TryCreateSourceReferenceItemAsync(
-                    classificationOptions,
-                    definitionItem,
-                    includeHiddenLocations: false,
-                    cancellationToken).ConfigureAwait(false);
-
-                if (referenceItem != null)
-                    await context.OnReferenceFoundAsync(referenceItem, cancellationToken).ConfigureAwait(false);
-            }
+            await ProducerConsumer<SourceReferenceItem>.RunParallelAsync(
+                source: references,
+                produceItems: static async (tuple, callback, args, cancellationToken) =>
+                {
+                    var (group, _, location) = tuple;
+                    var definitionItem = await args.@this.GetDefinitionItemAsync(group, cancellationToken).ConfigureAwait(false);
+                    var sourceReferenceItem = await location.TryCreateSourceReferenceItemAsync(
+                        args.classificationOptions,
+                        definitionItem,
+                        includeHiddenLocations: false,
+                        cancellationToken).ConfigureAwait(false);
+                    if (sourceReferenceItem != null)
+                        callback(sourceReferenceItem);
+                },
+                consumeItems: static async (items, args, cancellationToken) =>
+                    await args.context.OnReferencesFoundAsync(items, cancellationToken).ConfigureAwait(false),
+                args: (@this: this, context, classificationOptions),
+                cancellationToken).ConfigureAwait(false);
         }
     }
 }

src/Features/Core/Portable/FindUsages/FindUsagesContext.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Notification;
@@ -28,7 +29,7 @@ protected FindUsagesContext()
 
     public virtual ValueTask OnDefinitionFoundAsync(DefinitionItem definition, CancellationToken cancellationToken) => default;
 
-    public virtual ValueTask OnReferenceFoundAsync(SourceReferenceItem reference, CancellationToken cancellationToken) => default;
+    public virtual ValueTask OnReferencesFoundAsync(ImmutableArray<SourceReferenceItem> references, CancellationToken cancellationToken) => default;
 
     protected virtual ValueTask ReportProgressAsync(int current, int maximum, CancellationToken cancellationToken) => default;
 }

src/Features/Core/Portable/FindUsages/IFindUsagesContext.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Notification;
@@ -35,5 +36,5 @@ internal interface IFindUsagesContext
     ValueTask SetSearchTitleAsync(string title, CancellationToken cancellationToken);
 
     ValueTask OnDefinitionFoundAsync(DefinitionItem definition, CancellationToken cancellationToken);
-    ValueTask OnReferenceFoundAsync(SourceReferenceItem reference, CancellationToken cancellationToken);
+    ValueTask OnReferencesFoundAsync(ImmutableArray<SourceReferenceItem> references, CancellationToken cancellationToken);
 }

src/Features/Core/Portable/FindUsages/IRemoteFindUsagesService.cs

@@ -10,12 +10,10 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Classification;
-using Microsoft.CodeAnalysis.Collections;
 using Microsoft.CodeAnalysis.ErrorReporting;
 using Microsoft.CodeAnalysis.FindSymbols;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Notification;
-using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Remote;
 using Microsoft.CodeAnalysis.Text;
 using Roslyn.Utilities;
@@ -32,7 +30,7 @@ internal interface ICallback : IRemoteOptionsCallback<ClassificationOptions>
         ValueTask ReportInformationalMessageAsync(RemoteServiceCallbackId callbackId, string message, CancellationToken cancellationToken);
         ValueTask SetSearchTitleAsync(RemoteServiceCallbackId callbackId, string title, CancellationToken cancellationToken);
         ValueTask OnDefinitionFoundAsync(RemoteServiceCallbackId callbackId, SerializableDefinitionItem definition, CancellationToken cancellationToken);
-        ValueTask OnReferenceFoundAsync(RemoteServiceCallbackId callbackId, SerializableSourceReferenceItem reference, CancellationToken cancellationToken);
+        ValueTask OnReferencesFoundAsync(RemoteServiceCallbackId callbackId, ImmutableArray<SerializableSourceReferenceItem> references, CancellationToken cancellationToken);
     }
 
     ValueTask FindReferencesAsync(
@@ -73,8 +71,8 @@ public ValueTask ItemsCompletedAsync(RemoteServiceCallbackId callbackId, int cou
     public ValueTask OnDefinitionFoundAsync(RemoteServiceCallbackId callbackId, SerializableDefinitionItem definition, CancellationToken cancellationToken)
         => GetCallback(callbackId).OnDefinitionFoundAsync(definition, cancellationToken);
 
-    public ValueTask OnReferenceFoundAsync(RemoteServiceCallbackId callbackId, SerializableSourceReferenceItem reference, CancellationToken cancellationToken)
-        => GetCallback(callbackId).OnReferenceFoundAsync(reference, cancellationToken);
+    public ValueTask OnReferencesFoundAsync(RemoteServiceCallbackId callbackId, ImmutableArray<SerializableSourceReferenceItem> references, CancellationToken cancellationToken)
+        => GetCallback(callbackId).OnReferencesFoundAsync(references, cancellationToken);
 
     public ValueTask ReportMessageAsync(RemoteServiceCallbackId callbackId, string message, CancellationToken cancellationToken)
         => GetCallback(callbackId).ReportMessageAsync(message, cancellationToken);
@@ -131,12 +129,15 @@ public async ValueTask OnDefinitionFoundAsync(SerializableDefinitionItem definit
         }
     }
 
-    public async ValueTask OnReferenceFoundAsync(SerializableSourceReferenceItem reference, CancellationToken cancellationToken)
+    public async ValueTask OnReferencesFoundAsync(ImmutableArray<SerializableSourceReferenceItem> references, CancellationToken cancellationToken)
     {
         try
         {
-            var rehydrated = await reference.RehydrateAsync(_solution, GetDefinition(reference.DefinitionId), cancellationToken).ConfigureAwait(false);
-            await _context.OnReferenceFoundAsync(rehydrated, cancellationToken).ConfigureAwait(false);
+            var result = new FixedSizeArrayBuilder<SourceReferenceItem>(references.Length);
+            foreach (var reference in references)
+                result.Add(await reference.RehydrateAsync(_solution, GetDefinition(reference.DefinitionId), cancellationToken).ConfigureAwait(false));
+
+            await _context.OnReferencesFoundAsync(result.MoveToImmutable(), cancellationToken).ConfigureAwait(false);
         }
         catch (Exception ex) when (FatalError.ReportAndPropagateUnlessCanceled(ex, cancellationToken))
         {

src/Features/Core/Portable/ValueTracking/ValueTracker.FindReferencesProgress.cs

@@ -2,7 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using System.Collections.Immutable;
+using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.FindSymbols;
@@ -30,10 +30,10 @@ private class FindReferencesProgress(OperationCollector valueTrackingProgressCol
         public ValueTask OnDefinitionFoundAsync(SymbolGroup symbolGroup, CancellationToken _) => new();
 
         public async ValueTask OnReferencesFoundAsync(
-            ImmutableArray<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references,
+            IAsyncEnumerable<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references,
             CancellationToken cancellationToken)
         {
-            foreach (var (_, symbol, location) in references)
+            await foreach (var (_, symbol, location) in references)
                 await OnReferenceFoundAsync(symbol, location, cancellationToken).ConfigureAwait(false);
         }
 

src/Features/LanguageServer/Protocol/Features/FindUsages/SimpleFindUsagesContext.cs

@@ -65,11 +65,11 @@ public override ValueTask OnDefinitionFoundAsync(DefinitionItem definition, Canc
             return default;
         }
 
-        public override ValueTask OnReferenceFoundAsync(SourceReferenceItem reference, CancellationToken cancellationToken)
+        public override ValueTask OnReferencesFoundAsync(ImmutableArray<SourceReferenceItem> references, CancellationToken cancellationToken)
         {
             lock (_gate)
             {
-                _referenceItems.Add(reference);
+                _referenceItems.AddRange(references);
             }
 
             return default;

src/Features/LanguageServer/Protocol/Handler/References/FindUsagesLSPContext.cs

@@ -129,41 +129,44 @@ public override async ValueTask OnDefinitionFoundAsync(DefinitionItem definition
             }
         }
 
-        public override async ValueTask OnReferenceFoundAsync(SourceReferenceItem reference, CancellationToken cancellationToken)
+        public override async ValueTask OnReferencesFoundAsync(ImmutableArray<SourceReferenceItem> references, CancellationToken cancellationToken)
         {
             using (await _semaphore.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
             {
-                // Each reference should be associated with a definition. If this somehow isn't the
-                // case, we bail out early.
-                if (!_definitionToId.TryGetValue(reference.Definition, out var definitionId))
-                    return;
+                foreach (var reference in references)
+                {
+                    // Each reference should be associated with a definition. If this somehow isn't the
+                    // case, we bail out early.
+                    if (!_definitionToId.TryGetValue(reference.Definition, out var definitionId))
+                        return;
 
-                var documentSpan = reference.SourceSpan;
-                var document = documentSpan.Document;
+                    var documentSpan = reference.SourceSpan;
+                    var document = documentSpan.Document;
 
-                // If this is reference to the same physical location we've already reported, just
-                // filter this out.  it will clutter the UI to show the same places.
-                if (!_referenceLocations.Add((document.FilePath, reference.SourceSpan.SourceSpan)))
-                    return;
+                    // If this is reference to the same physical location we've already reported, just
+                    // filter this out.  it will clutter the UI to show the same places.
+                    if (!_referenceLocations.Add((document.FilePath, reference.SourceSpan.SourceSpan)))
+                        return;
 
-                // If the definition hasn't been reported yet, add it to our list of references to report.
-                if (_definitionsWithoutReference.TryGetValue(definitionId, out var definition))
-                {
-                    _workQueue.AddWork(definition);
-                    _definitionsWithoutReference.Remove(definitionId);
-                }
+                    // If the definition hasn't been reported yet, add it to our list of references to report.
+                    if (_definitionsWithoutReference.TryGetValue(definitionId, out var definition))
+                    {
+                        _workQueue.AddWork(definition);
+                        _definitionsWithoutReference.Remove(definitionId);
+                    }
 
-                // give this reference a fresh id.
-                _id++;
+                    // give this reference a fresh id.
+                    _id++;
 
-                // Creating a new VSReferenceItem for the reference
-                var referenceItem = await GenerateVSReferenceItemAsync(
-                    definitionId, _id, reference.SourceSpan,
-                    reference.AdditionalProperties, definitionText: null,
-                    definitionGlyph: Glyph.None, reference.SymbolUsageInfo, reference.IsWrittenTo, cancellationToken).ConfigureAwait(false);
+                    // Creating a new VSReferenceItem for the reference
+                    var referenceItem = await GenerateVSReferenceItemAsync(
+                        definitionId, _id, reference.SourceSpan,
+                        reference.AdditionalProperties, definitionText: null,
+                        definitionGlyph: Glyph.None, reference.SymbolUsageInfo, reference.IsWrittenTo, cancellationToken).ConfigureAwait(false);
 
-                if (referenceItem != null)
-                    _workQueue.AddWork(referenceItem.Value);
+                    if (referenceItem != null)
+                        _workQueue.AddWork(referenceItem.Value);
+                }
             }
         }
 

src/Tools/ExternalAccess/FSharp/Internal/Editor/FindUsages/FSharpFindUsagesContext.cs

@@ -29,7 +29,7 @@ public Task OnDefinitionFoundAsync(FSharp.FindUsages.FSharpDefinitionItem defini
 
         public Task OnReferenceFoundAsync(FSharp.FindUsages.FSharpSourceReferenceItem reference)
         {
-            return _context.OnReferenceFoundAsync(reference.RoslynSourceReferenceItem, _cancellationToken).AsTask();
+            return _context.OnReferencesFoundAsync([reference.RoslynSourceReferenceItem], _cancellationToken).AsTask();
         }
 
         public Task ReportMessageAsync(string message)