Merge pull request #73777 from CyrusNajmabadi/farChunking

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

@@ -5,6 +5,7 @@
 #nullable disable
 
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Classification;
@@ -106,7 +107,7 @@ public async ValueTask OnDefinitionFoundAsync(SymbolGroup group, CancellationTok
         }
 
         public async ValueTask OnReferencesFoundAsync(
-            IAsyncEnumerable<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken)
+            ImmutableArray<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken)
         {
             await ProducerConsumer<SourceReferenceItem>.RunParallelAsync(
                 source: references,

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

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

src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs

@@ -134,26 +134,30 @@ await ProducerConsumer<FinderLocation>.RunAsync(
                 },
                 consumeItems: static async (values, args, cancellationToken) =>
                 {
-                    var (@this, symbol, _) = args;
-                    await @this._progress.OnReferencesFoundAsync(
-                        ReadAllAsync(@this, values, symbol, cancellationToken), cancellationToken).ConfigureAwait(false);
+                    var (@this, symbol, state) = args;
+                    var converted = await ConvertLocationsAndReportGroupsAsync(@this, values, symbol, cancellationToken).ConfigureAwait(false);
+                    await @this._progress.OnReferencesFoundAsync(converted, cancellationToken).ConfigureAwait(false);
                 },
                 args: (@this: this, symbol, state),
                 cancellationToken).ConfigureAwait(false);
         }
 
-        static async IAsyncEnumerable<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> ReadAllAsync(
-            FindReferencesSearchEngine @this, IAsyncEnumerable<FinderLocation> locations, ISymbol symbol, [EnumeratorCancellation] CancellationToken cancellationToken)
+        static async Task<ImmutableArray<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)>> ConvertLocationsAndReportGroupsAsync(
+            FindReferencesSearchEngine @this, IAsyncEnumerable<FinderLocation> locations, ISymbol symbol, CancellationToken cancellationToken)
         {
             SymbolGroup? group = null;
 
+            using var _ = ArrayBuilder<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)>.GetInstance(out var result);
+
             // Transform the individual finder-location objects to "group/symbol/location" tuples.
             await foreach (var location in locations)
             {
                 // The first time we see the location for a symbol, report its group.
                 group ??= await @this.ReportGroupAsync(symbol, cancellationToken).ConfigureAwait(false);
-                yield return (group, symbol, location.Location);
+                result.Add((group, symbol, location.Location));
             }
+
+            return result.ToImmutableAndClear();
         }
 
         async ValueTask InheritanceSymbolSearchAsync(ISymbol symbol, FindReferencesDocumentState state)
@@ -174,8 +178,7 @@ async ValueTask InheritanceSymbolSearchAsync(ISymbol symbol, FindReferencesDocum
                         var candidateGroup = await ReportGroupAsync(candidate, cancellationToken).ConfigureAwait(false);
 
                         var location = AbstractReferenceFinder.CreateReferenceLocation(state, token, candidateReason, cancellationToken);
-                        await _progress.OnReferencesFoundAsync(
-                            IAsyncEnumerableExtensions.AsAsyncEnumerable([(candidateGroup, candidate, location)]), cancellationToken).ConfigureAwait(false);
+                        await _progress.OnReferencesFoundAsync([(candidateGroup, candidate, location)], cancellationToken).ConfigureAwait(false);
                     }
                 }
             }

src/Workspaces/Core/Portable/FindSymbols/FindReferences/NoOpStreamingFindReferencesProgress.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Shared.Utilities;
@@ -27,7 +28,7 @@ private NoOpStreamingFindReferencesProgress()
     public ValueTask OnCompletedAsync(CancellationToken cancellationToken) => default;
     public ValueTask OnStartedAsync(CancellationToken cancellationToken) => default;
     public ValueTask OnDefinitionFoundAsync(SymbolGroup group, CancellationToken cancellationToken) => default;
-    public ValueTask OnReferencesFoundAsync(IAsyncEnumerable<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken) => default;
+    public ValueTask OnReferencesFoundAsync(ImmutableArray<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken) => default;
 
     private class NoOpProgressTracker : IStreamingProgressTracker
     {

src/Workspaces/Core/Portable/FindSymbols/FindReferences/StreamingFindReferencesProgress.cs

@@ -4,6 +4,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.ErrorReporting;
@@ -53,10 +54,12 @@ public ValueTask OnDefinitionFoundAsync(SymbolGroup group, CancellationToken can
         }
     }
 
-    public async ValueTask OnReferencesFoundAsync(IAsyncEnumerable<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken)
+    public ValueTask OnReferencesFoundAsync(ImmutableArray<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken)
     {
-        await foreach (var (_, symbol, location) in references)
+        foreach (var (_, symbol, location) in references)
             _progress.OnReferenceFound(symbol, location);
+
+        return ValueTaskFactory.CompletedTask;
     }
 
     public ValueTask OnStartedAsync(CancellationToken cancellationToken)

src/Workspaces/Core/Portable/FindSymbols/IStreamingFindReferencesProgress.cs

@@ -74,7 +74,7 @@ internal interface IStreamingFindReferencesProgress
     ValueTask OnCompletedAsync(CancellationToken cancellationToken);
 
     ValueTask OnDefinitionFoundAsync(SymbolGroup group, CancellationToken cancellationToken);
-    ValueTask OnReferencesFoundAsync(IAsyncEnumerable<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken);
+    ValueTask OnReferencesFoundAsync(ImmutableArray<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken);
 }
 
 internal interface IStreamingFindLiteralReferencesProgress

src/Workspaces/Core/Portable/FindSymbols/StreamingProgressCollector.cs

@@ -10,8 +10,6 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.ErrorReporting;
-using Microsoft.CodeAnalysis.PooledObjects;
-using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Shared.Utilities;
 using Roslyn.Utilities;
 
@@ -71,18 +69,14 @@ public ValueTask OnDefinitionFoundAsync(SymbolGroup group, CancellationToken can
     }
 
     public async ValueTask OnReferencesFoundAsync(
-        IAsyncEnumerable<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken)
+        ImmutableArray<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken)
     {
-        // Reading the references from the stream will cause them to be processed.  We'll make a copy here so we can
-        // defer to the underlying progress object with the same data.
-        using var _ = ArrayBuilder<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)>.GetInstance(out var copy);
-        await foreach (var tuple in references)
+        lock (_gate)
         {
-            copy.Add(tuple);
-            lock (_gate)
+            foreach (var tuple in references)
                 _symbolToLocations[tuple.symbol].Add(tuple.location);
         }
 
-        await underlyingProgress.OnReferencesFoundAsync(copy.AsAsyncEnumerable(), cancellationToken).ConfigureAwait(false);
+        await underlyingProgress.OnReferencesFoundAsync(references, cancellationToken).ConfigureAwait(false);
     }
 }

src/Workspaces/Core/Portable/FindSymbols/SymbolFinder.FindReferencesServerCallback.cs

@@ -65,40 +65,38 @@ public async ValueTask OnDefinitionFoundAsync(SerializableSymbolGroup dehydrated
             await progress.OnDefinitionFoundAsync(symbolGroup, cancellationToken).ConfigureAwait(false);
         }
 
-        public ValueTask OnReferencesFoundAsync(
+        public async ValueTask OnReferencesFoundAsync(
             ImmutableArray<(SerializableSymbolGroup serializableSymbolGroup, SerializableSymbolAndProjectId serializableSymbol, SerializableReferenceLocation reference)> references,
             CancellationToken cancellationToken)
         {
-            return progress.OnReferencesFoundAsync(ConvertAsync(cancellationToken), cancellationToken);
+            using var _ = ArrayBuilder<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)>.GetInstance(references.Length, out var rehydrated);
 
-            async IAsyncEnumerable<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> ConvertAsync([EnumeratorCancellation] CancellationToken cancellationToken)
+            foreach (var (serializableSymbolGroup, serializableSymbol, reference) in references)
             {
-                foreach (var (serializableSymbolGroup, serializableSymbol, reference) in references)
+                SymbolGroup? symbolGroup;
+                ISymbol? symbol;
+                lock (_gate)
                 {
-                    SymbolGroup? symbolGroup;
-                    ISymbol? symbol;
-                    lock (_gate)
+                    // The definition may not be in the map if we failed to map it over using TryRehydrateAsync in OnDefinitionFoundAsync.
+                    // Just ignore this reference.  Note: while this is a degraded experience:
+                    //
+                    // 1. TryRehydrateAsync logs an NFE so we can track down while we're failing to roundtrip the
+                    //    definition so we can track down that issue.
+                    // 2. NFE'ing and failing to show a result, is much better than NFE'ing and then crashing
+                    //    immediately afterwards.
+                    if (!_groupMap.TryGetValue(serializableSymbolGroup, out symbolGroup) ||
+                        !_definitionMap.TryGetValue(serializableSymbol, out symbol))
                     {
-                        // The definition may not be in the map if we failed to map it over using TryRehydrateAsync in OnDefinitionFoundAsync.
-                        // Just ignore this reference.  Note: while this is a degraded experience:
-                        //
-                        // 1. TryRehydrateAsync logs an NFE so we can track down while we're failing to roundtrip the
-                        //    definition so we can track down that issue.
-                        // 2. NFE'ing and failing to show a result, is much better than NFE'ing and then crashing
-                        //    immediately afterwards.
-                        if (!_groupMap.TryGetValue(serializableSymbolGroup, out symbolGroup) ||
-                            !_definitionMap.TryGetValue(serializableSymbol, out symbol))
-                        {
-                            continue;
-                        }
+                        continue;
                     }
-
-                    var referenceLocation = await reference.RehydrateAsync(
-                        solution, cancellationToken).ConfigureAwait(false);
-
-                    yield return (symbolGroup, symbol, referenceLocation);
                 }
+
+                var referenceLocation = await reference.RehydrateAsync(
+                    solution, cancellationToken).ConfigureAwait(false);
+                rehydrated.Add((symbolGroup, symbol, referenceLocation));
             }
+
+            await progress.OnReferencesFoundAsync(rehydrated.ToImmutableAndClear(), cancellationToken).ConfigureAwait(false);
         }
     }
 }

src/Workspaces/Remote/ServiceHub/Services/SymbolFinder/RemoteSymbolFinderService.cs

@@ -227,19 +227,14 @@ public ValueTask OnDefinitionFoundAsync(SymbolGroup group, CancellationToken can
             }
 
             public async ValueTask OnReferencesFoundAsync(
-                IAsyncEnumerable<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references,
+                ImmutableArray<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references,
                 CancellationToken cancellationToken)
             {
-                using var _ = ArrayBuilder<(SerializableSymbolGroup, SerializableSymbolAndProjectId, SerializableReferenceLocation)>.GetInstance(out var result);
-                await foreach (var (group, symbol, location) in references)
-                {
-                    result.Add((
-                        SerializableSymbolGroup.Dehydrate(_solution, group, cancellationToken),
-                         SerializableSymbolAndProjectId.Dehydrate(_solution, symbol, cancellationToken),
-                         SerializableReferenceLocation.Dehydrate(location, cancellationToken)));
-                }
+                var dehydrated = references.SelectAsArray(
+                    r => (SerializableSymbolGroup.Dehydrate(_solution, r.group, cancellationToken),
+                          SerializableSymbolAndProjectId.Dehydrate(_solution, r.symbol, cancellationToken),
+                          SerializableReferenceLocation.Dehydrate(r.location, cancellationToken)));
 
-                var dehydrated = result.ToImmutableAndClear();
                 await _callback.InvokeAsync(
                     (callback, cancellationToken) => callback.OnReferencesFoundAsync(
                         _callbackId, dehydrated, cancellationToken), cancellationToken).ConfigureAwait(false);