Refactor searching and reporting results in navto.

src/EditorFeatures/Core.Wpf/NavigateTo/NavigateToItemProvider.Callback.cs

@@ -3,10 +3,12 @@
 // 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.ErrorReporting;
 using Microsoft.CodeAnalysis.NavigateTo;
+using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Text.Shared.Extensions;
 using Microsoft.VisualStudio.Language.NavigateTo.Interfaces;
 using Microsoft.VisualStudio.Text.PatternMatching;
@@ -18,11 +20,13 @@ internal partial class NavigateToItemProvider
     {
         private class NavigateToItemProviderCallback : INavigateToSearchCallback
         {
+            private readonly Solution _solution;
             private readonly INavigateToItemDisplayFactory _displayFactory;
             private readonly INavigateToCallback _callback;
 
-            public NavigateToItemProviderCallback(INavigateToItemDisplayFactory displayFactory, INavigateToCallback callback)
+            public NavigateToItemProviderCallback(Solution solution, INavigateToItemDisplayFactory displayFactory, INavigateToCallback callback)
             {
+                _solution = solution;
                 _displayFactory = displayFactory;
                 _callback = callback;
             }
@@ -39,9 +43,41 @@ public void Done(bool isFullyLoaded)
                 }
             }
 
-            public Task AddItemAsync(Project project, INavigateToSearchResult result, CancellationToken cancellationToken)
+            public Task AddResultsAsync(ImmutableArray<INavigateToSearchResult> results, CancellationToken cancellationToken)
             {
-                ReportMatchResult(project, result);
+                foreach (var result in results)
+                {
+                    var matchedSpans = result.NameMatchSpans.SelectAsArray(t => t.ToSpan());
+
+                    var patternMatch = new PatternMatch(
+                        GetPatternMatchKind(result.MatchKind),
+                        punctuationStripped: false,
+                        result.IsCaseSensitive,
+                        matchedSpans);
+
+                    var project = _solution.GetRequiredProject(result.NavigableItem.Document.Project.Id);
+                    var navigateToItem = new NavigateToItem(
+                        result.Name,
+                        result.Kind,
+                        GetNavigateToLanguage(project.Language),
+                        result.SecondarySort,
+                        result,
+                        patternMatch,
+                        _displayFactory);
+
+                    try
+                    {
+                        _callback.AddItem(navigateToItem);
+                    }
+                    catch (InvalidOperationException ex) when (FatalError.ReportAndCatch(ex, ErrorSeverity.Critical))
+                    {
+                        // Mitigation for race condition in platform https://devdiv.visualstudio.com/DevDiv/_workitems/edit/1534364
+                        //
+                        // Catch this so that don't tear down OOP, but still report the exception so that we ensure this issue
+                        // gets attention and is fixed.
+                    }
+                }
+
                 return Task.CompletedTask;
             }
 
@@ -54,38 +90,6 @@ public void ReportIncomplete()
             {
             }
 
-            private void ReportMatchResult(Project project, INavigateToSearchResult result)
-            {
-                var matchedSpans = result.NameMatchSpans.SelectAsArray(t => t.ToSpan());
-
-                var patternMatch = new PatternMatch(
-                    GetPatternMatchKind(result.MatchKind),
-                    punctuationStripped: false,
-                    result.IsCaseSensitive,
-                    matchedSpans);
-
-                var navigateToItem = new NavigateToItem(
-                    result.Name,
-                    result.Kind,
-                    GetNavigateToLanguage(project.Language),
-                    result.SecondarySort,
-                    result,
-                    patternMatch,
-                    _displayFactory);
-
-                try
-                {
-                    _callback.AddItem(navigateToItem);
-                }
-                catch (InvalidOperationException ex) when (FatalError.ReportAndCatch(ex, ErrorSeverity.Critical))
-                {
-                    // Mitigation for race condition in platform https://devdiv.visualstudio.com/DevDiv/_workitems/edit/1534364
-                    //
-                    // Catch this so that don't tear down OOP, but still report the exception so that we ensure this issue
-                    // gets attention and is fixed.
-                }
-            }
-
             private static PatternMatchKind GetPatternMatchKind(NavigateToMatchKind matchKind)
                 => matchKind switch
                 {

src/EditorFeatures/Core.Wpf/NavigateTo/NavigateToItemProvider.cs

@@ -109,9 +109,10 @@ private void StartSearch(INavigateToCallback callback, string searchValue, IImmu
             ? NavigateToSearchScope.Document
             : NavigateToSearchScope.Solution;
 
-        var roslynCallback = new NavigateToItemProviderCallback(_displayFactory, callback);
+        var solution = _workspace.CurrentSolution;
+        var roslynCallback = new NavigateToItemProviderCallback(solution, _displayFactory, callback);
         var searcher = NavigateToSearcher.Create(
-            _workspace.CurrentSolution,
+            solution,
             _asyncListener,
             roslynCallback,
             searchValue,

src/Features/Core/Portable/ExternalAccess/VSTypeScript/VSTypeScriptNavigateToSearchService.cs

@@ -35,14 +35,14 @@ public async Task SearchDocumentAsync(
         Document document,
         string searchPattern,
         IImmutableSet<string> kinds,
-        Func<INavigateToSearchResult, Task> onResultFound,
+        Func<ImmutableArray<INavigateToSearchResult>, Task> onResultsFound,
         CancellationToken cancellationToken)
     {
         if (_searchService != null)
         {
             var results = await _searchService.SearchDocumentAsync(document, searchPattern, kinds, cancellationToken).ConfigureAwait(false);
-            foreach (var result in results)
-                await onResultFound(Convert(result)).ConfigureAwait(false);
+            if (results.Length > 0)
+                await onResultsFound(results.SelectAsArray(Convert)).ConfigureAwait(false);
         }
     }
 
@@ -53,7 +53,7 @@ public async Task SearchProjectsAsync(
         string searchPattern,
         IImmutableSet<string> kinds,
         Document? activeDocument,
-        Func<Project, INavigateToSearchResult, Task> onResultFound,
+        Func<ImmutableArray<INavigateToSearchResult>, Task> onResultsFound,
         Func<Task> onProjectCompleted,
         CancellationToken cancellationToken)
     {
@@ -78,8 +78,9 @@ async Task ProcessProjectAsync(Project project)
                 {
                     var results = await _searchService.SearchProjectAsync(
                         project, priorityDocuments.WhereAsArray(d => d.Project == project), searchPattern, kinds, cancellationToken).ConfigureAwait(false);
-                    foreach (var result in results)
-                        await onResultFound(project, Convert(result)).ConfigureAwait(false);
+
+                    if (results.Length > 0)
+                        await onResultsFound(results.SelectAsArray(Convert)).ConfigureAwait(false);
                 }
 
                 await onProjectCompleted().ConfigureAwait(false);

src/Features/Core/Portable/NavigateTo/AbstractNavigateToSearchService.CachedDocumentSearch.cs

@@ -4,7 +4,6 @@
 
 using System;
 using System.Collections.Concurrent;
-using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
@@ -14,7 +13,6 @@
 using Microsoft.CodeAnalysis.FindSymbols;
 using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.PatternMatching;
-using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Remote;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Storage;
@@ -63,25 +61,27 @@ public async Task SearchCachedDocumentsAsync(
         string searchPattern,
         IImmutableSet<string> kinds,
         Document? activeDocument,
-        Func<Project, INavigateToSearchResult, Task> onResultFound,
+        Func<ImmutableArray<INavigateToSearchResult>, Task> onResultsFound,
         Func<Task> onProjectCompleted,
         CancellationToken cancellationToken)
     {
-        cancellationToken.ThrowIfCancellationRequested();
+        if (cancellationToken.IsCancellationRequested)
+            return;
+
         Contract.ThrowIfTrue(projects.IsEmpty);
         Contract.ThrowIfTrue(projects.Select(p => p.Language).Distinct().Count() != 1);
 
         Debug.Assert(priorityDocuments.All(d => projects.Contains(d.Project)));
 
-        var onItemFound = GetOnItemFoundCallback(solution, activeDocument, onResultFound, cancellationToken);
+        var onItemsFound = GetOnItemsFoundCallback(solution, activeDocument, onResultsFound, cancellationToken);
 
         var documentKeys = projects.SelectManyAsArray(p => p.Documents.Select(DocumentKey.ToDocumentKey));
         var priorityDocumentKeys = priorityDocuments.SelectAsArray(DocumentKey.ToDocumentKey);
 
         var client = await RemoteHostClient.TryGetClientAsync(solution.Services, cancellationToken).ConfigureAwait(false);
         if (client != null)
         {
-            var callback = new NavigateToSearchServiceCallback(onItemFound, onProjectCompleted);
+            var callback = new NavigateToSearchServiceCallback(onItemsFound, onProjectCompleted);
             await client.TryInvokeAsync<IRemoteNavigateToSearchService>(
                 (service, callbackId, cancellationToken) =>
                     service.SearchCachedDocumentsAsync(documentKeys, priorityDocumentKeys, searchPattern, [.. kinds], callbackId, cancellationToken),
@@ -92,7 +92,7 @@ await client.TryInvokeAsync<IRemoteNavigateToSearchService>(
 
         var storageService = solution.Services.GetPersistentStorageService();
         await SearchCachedDocumentsInCurrentProcessAsync(
-            storageService, documentKeys, priorityDocumentKeys, searchPattern, kinds, onItemFound, onProjectCompleted, cancellationToken).ConfigureAwait(false);
+            storageService, documentKeys, priorityDocumentKeys, searchPattern, kinds, onItemsFound, onProjectCompleted, cancellationToken).ConfigureAwait(false);
     }
 
     public static async Task SearchCachedDocumentsInCurrentProcessAsync(
@@ -101,16 +101,14 @@ public static async Task SearchCachedDocumentsInCurrentProcessAsync(
         ImmutableArray<DocumentKey> priorityDocumentKeys,
         string searchPattern,
         IImmutableSet<string> kinds,
-        Func<RoslynNavigateToItem, Task> onItemFound,
+        Func<ImmutableArray<RoslynNavigateToItem>, Task> onItemsFound,
         Func<Task> onProjectCompleted,
         CancellationToken cancellationToken)
     {
-        cancellationToken.ThrowIfCancellationRequested();
         // Quick abort if OOP is now fully loaded.
         if (!ShouldSearchCachedDocuments(out _, out _))
             return;
 
-        // If the user created a dotted pattern then we'll grab the last part of the name
         var (patternName, patternContainer) = PatternMatcher.GetNameAndContainer(searchPattern);
         var declaredSymbolInfoKindsSet = new DeclaredSymbolInfoKindSet(kinds);
 
@@ -120,84 +118,56 @@ public static async Task SearchCachedDocumentsInCurrentProcessAsync(
 
         using var _1 = GetPooledHashSet(priorityDocumentKeys, out var priorityDocumentKeysSet);
 
-        // Sort the groups into a high pri group (projects that contain a high-pri doc), and low pri groups (those
-        // that don't).
+        // Sort the groups into a high pri group (projects that contain a high-pri doc), and low pri groups (those that
+        // don't), and process in that order.
         using var _2 = GetPooledHashSet(groups.Where(g => g.Any(priorityDocumentKeysSet.Contains)), out var highPriorityGroups);
         using var _3 = GetPooledHashSet(groups.Where(g => !highPriorityGroups.Contains(g)), out var lowPriorityGroups);
 
-        await ProcessProjectGroupsAsync(highPriorityGroups).ConfigureAwait(false);
-        await ProcessProjectGroupsAsync(lowPriorityGroups).ConfigureAwait(false);
-
+        await PerformParallelSearchAsync(
+            highPriorityGroups.Concat(lowPriorityGroups), ProcessSingleProjectGroupAsync, onItemsFound, cancellationToken).ConfigureAwait(false);
         return;
 
-        async Task ProcessProjectGroupsAsync(HashSet<IGrouping<ProjectKey, DocumentKey>> groups)
+        async ValueTask ProcessSingleProjectGroupAsync(
+            IGrouping<ProjectKey, DocumentKey> group,
+            Action<RoslynNavigateToItem> onItemFound,
+            CancellationToken cancellationToken)
         {
-            cancellationToken.ThrowIfCancellationRequested();
-            using var _ = ArrayBuilder<Task>.GetInstance(out var tasks);
+            if (cancellationToken.IsCancellationRequested)
+                return;
 
-            foreach (var group in groups)
-                tasks.Add(ProcessProjectGroupAsync(group));
-
-            await Task.WhenAll(tasks).ConfigureAwait(false);
-        }
-
-        async Task ProcessProjectGroupAsync(IGrouping<ProjectKey, DocumentKey> group)
-        {
-            cancellationToken.ThrowIfCancellationRequested();
-            await Task.Yield().ConfigureAwait(false);
             var project = group.Key;
 
-            // Break the project into high-pri docs and low pri docs.
+            // Break the project into high-pri docs and low pri docs, and process in that order.
             using var _1 = GetPooledHashSet(group.Where(priorityDocumentKeysSet.Contains), out var highPriDocs);
             using var _2 = GetPooledHashSet(group.Where(d => !highPriDocs.Contains(d)), out var lowPriDocs);
 
-            await SearchCachedDocumentsInCurrentProcessAsync(
-                storageService, patternName, patternContainer, declaredSymbolInfoKindsSet,
-                onItemFound, highPriDocs, cancellationToken).ConfigureAwait(false);
+            await ParallelForEachAsync(
+                highPriDocs.Concat(lowPriDocs),
+                cancellationToken,
+                async (documentKey, cancellationToken) =>
+                {
+                    var index = await GetIndexAsync(storageService, documentKey, cancellationToken).ConfigureAwait(false);
+                    if (index == null)
+                        return;
 
-            await SearchCachedDocumentsInCurrentProcessAsync(
-                storageService, patternName, patternContainer, declaredSymbolInfoKindsSet,
-                onItemFound, lowPriDocs, cancellationToken).ConfigureAwait(false);
+                    ProcessIndex(
+                        documentKey, document: null, patternName, patternContainer, declaredSymbolInfoKindsSet,
+                        index, linkedIndices: null, onItemFound, cancellationToken);
+                }).ConfigureAwait(false);
 
             // done with project.  Let the host know.
             await onProjectCompleted().ConfigureAwait(false);
         }
     }
 
-    private static async Task SearchCachedDocumentsInCurrentProcessAsync(
-        IChecksummedPersistentStorageService storageService,
-        string patternName,
-        string? patternContainer,
-        DeclaredSymbolInfoKindSet kinds,
-        Func<RoslynNavigateToItem, Task> onItemFound,
-        HashSet<DocumentKey> documentKeys,
-        CancellationToken cancellationToken)
-    {
-        cancellationToken.ThrowIfCancellationRequested();
-        using var _ = ArrayBuilder<Task>.GetInstance(out var tasks);
-
-        foreach (var documentKey in documentKeys)
-        {
-            tasks.Add(Task.Run(async () =>
-            {
-                var index = await GetIndexAsync(storageService, documentKey, cancellationToken).ConfigureAwait(false);
-                if (index == null)
-                    return;
-
-                await ProcessIndexAsync(
-                    documentKey, document: null, patternName, patternContainer, kinds, onItemFound, index, cancellationToken).ConfigureAwait(false);
-            }, cancellationToken));
-        }
-
-        await Task.WhenAll(tasks).ConfigureAwait(false);
-    }
-
     private static Task<TopLevelSyntaxTreeIndex?> GetIndexAsync(
         IChecksummedPersistentStorageService storageService,
         DocumentKey documentKey,
         CancellationToken cancellationToken)
     {
-        cancellationToken.ThrowIfCancellationRequested();
+        if (cancellationToken.IsCancellationRequested)
+            return SpecializedTasks.Null<TopLevelSyntaxTreeIndex>();
+
         // Retrieve the string table we use to dedupe strings.  If we can't get it, that means the solution has 
         // fully loaded and we've switched over to normal navto lookup.
         if (!ShouldSearchCachedDocuments(out var cachedIndexMap, out var stringTable))