Extract async navigation portions of https://github.com/dotnet/roslyn/pull/55635

src/EditorFeatures/Core.Cocoa/NavigationCommandHandlers/FindExtensionMethodsCommandHandler.cs

@@ -62,66 +62,64 @@ protected override bool TryExecuteCommand(int caretPosition, Document document,
         private async Task FindExtensionMethodsAsync(
             Document document, int caretPosition, IStreamingFindUsagesPresenter presenter)
         {
+            var solution = document.Project.Solution;
             try
             {
                 using var token = _asyncListener.BeginAsyncOperation(nameof(FindExtensionMethodsAsync));
 
                 var (context, cancellationToken) = presenter.StartSearch(EditorFeaturesResources.Navigating, supportsReferences: true);
 
-                using (Logger.LogBlock(
-                    FunctionId.CommandHandler_FindAllReference,
-                    KeyValueLogMessage.Create(LogType.UserAction, m => m["type"] = "streaming"),
-                    cancellationToken))
+                try
                 {
-                    var candidateSymbolProjectPair = await FindUsagesHelpers.GetRelevantSymbolAndProjectAtPositionAsync(document, caretPosition, cancellationToken).ConfigureAwait(false);
-
-                    var symbol = candidateSymbolProjectPair?.symbol as INamedTypeSymbol;
-
-                    // if we didn't get the right symbol, just abort
-                    if (symbol == null)
+                    using (Logger.LogBlock(
+                        FunctionId.CommandHandler_FindAllReference,
+                        KeyValueLogMessage.Create(LogType.UserAction, m => m["type"] = "streaming"),
+                        cancellationToken))
                     {
-                        await context.OnCompletedAsync(cancellationToken).ConfigureAwait(false);
-                        return;
-                    }
+                        var candidateSymbolProjectPair = await FindUsagesHelpers.GetRelevantSymbolAndProjectAtPositionAsync(document, caretPosition, cancellationToken).ConfigureAwait(false);
 
-                    if (!document.Project.TryGetCompilation(out var compilation))
-                    {
-                        await context.OnCompletedAsync(cancellationToken).ConfigureAwait(false);
-                        return;
-                    }
+                        var symbol = candidateSymbolProjectPair?.symbol as INamedTypeSymbol;
 
-                    var solution = document.Project.Solution;
+                        // if we didn't get the right symbol, just abort
+                        if (symbol == null)
+                            return;
 
-                    foreach (var type in compilation.Assembly.GlobalNamespace.GetAllTypes(cancellationToken))
-                    {
-                        if (!type.MightContainExtensionMethods)
-                            continue;
+                        if (!document.Project.TryGetCompilation(out var compilation))
+                            return;
 
-                        foreach (var extMethod in type.GetMembers().OfType<IMethodSymbol>().Where(method => method.IsExtensionMethod))
+                        foreach (var type in compilation.Assembly.GlobalNamespace.GetAllTypes(cancellationToken))
                         {
-                            if (cancellationToken.IsCancellationRequested)
-                                break;
+                            if (!type.MightContainExtensionMethods)
+                                continue;
 
-                            var reducedMethod = extMethod.ReduceExtensionMethod(symbol);
-                            if (reducedMethod != null)
+                            foreach (var extMethod in type.GetMembers().OfType<IMethodSymbol>().Where(method => method.IsExtensionMethod))
                             {
-                                var loc = extMethod.Locations.First();
+                                if (cancellationToken.IsCancellationRequested)
+                                    break;
 
-                                var sourceDefinition = await SymbolFinder.FindSourceDefinitionAsync(reducedMethod, solution, cancellationToken).ConfigureAwait(false);
-
-                                // And if our definition actually is from source, then let's re-figure out what project it came from
-                                if (sourceDefinition != null)
+                                var reducedMethod = extMethod.ReduceExtensionMethod(symbol);
+                                if (reducedMethod != null)
                                 {
-                                    var originatingProject = solution.GetProject(sourceDefinition.ContainingAssembly, cancellationToken);
+                                    var loc = extMethod.Locations.First();
+
+                                    var sourceDefinition = await SymbolFinder.FindSourceDefinitionAsync(reducedMethod, solution, cancellationToken).ConfigureAwait(false);
+
+                                    // And if our definition actually is from source, then let's re-figure out what project it came from
+                                    if (sourceDefinition != null)
+                                    {
+                                        var originatingProject = solution.GetProject(sourceDefinition.ContainingAssembly, cancellationToken);
 
-                                    var definitionItem = reducedMethod.ToNonClassifiedDefinitionItem(solution, true);
+                                        var definitionItem = reducedMethod.ToNonClassifiedDefinitionItem(solution, true);
 
-                                    await context.OnDefinitionFoundAsync(definitionItem, cancellationToken).ConfigureAwait(false);
+                                        await context.OnDefinitionFoundAsync(definitionItem, cancellationToken).ConfigureAwait(false);
+                                    }
                                 }
                             }
                         }
                     }
-
+                }
+                finally
+                {
                     await context.OnCompletedAsync(cancellationToken).ConfigureAwait(false);
                 }
             }

src/EditorFeatures/Core.Cocoa/NavigationCommandHandlers/FindMemberOverloadsCommandHandler.cs

@@ -84,7 +84,7 @@ private async Task FindMemberOverloadsAsync(
                         foreach (var curSymbol in symbol.ContainingType.GetMembers()
                                                         .Where(m => m.Kind == symbol.Kind && m.Name == symbol.Name))
                         {
-                            var definitionItem = curSymbol.ToNonClassifiedDefinitionItem(document.Project.Solution, true);
+                            var definitionItem = curSymbol.ToNonClassifiedDefinitionItem(document.Project.Solution, includeHiddenLocations: true);
                             await context.OnDefinitionFoundAsync(definitionItem, cancellationToken).ConfigureAwait(false);
                         }
                     }

src/EditorFeatures/Core.Wpf/Interactive/InteractiveDocumentNavigationService.cs

@@ -7,8 +7,10 @@
 using System;
 using System.Diagnostics;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Editor.Interactive;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
+using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Navigation;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.Text;
@@ -19,6 +21,22 @@ namespace Microsoft.CodeAnalysis.Editor.Implementation.Interactive
 {
     internal sealed class InteractiveDocumentNavigationService : IDocumentNavigationService
     {
+        private readonly IThreadingContext _threadingContext;
+
+        public InteractiveDocumentNavigationService(IThreadingContext threadingContext)
+        {
+            _threadingContext = threadingContext;
+        }
+
+        public async Task<bool> CanNavigateToSpanAsync(Workspace workspace, DocumentId documentId, TextSpan textSpan, CancellationToken cancellationToken)
+        {
+            // This switch is technically not needed as the call to CanNavigateToSpan just returns 'true'.
+            // However, this abides by the contract that CanNavigateToSpan only be called on the UI thread.
+            // It also means if we ever update that method, this code will stay corrrect.
+            await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+            return CanNavigateToSpan(workspace, documentId, textSpan, cancellationToken);
+        }
+
         public bool CanNavigateToSpan(Workspace workspace, DocumentId documentId, TextSpan textSpan, CancellationToken cancellationToken)
             => true;
 
@@ -28,6 +46,12 @@ public bool CanNavigateToLineAndOffset(Workspace workspace, DocumentId documentI
         public bool CanNavigateToPosition(Workspace workspace, DocumentId documentId, int position, int virtualSpace, CancellationToken cancellationToken)
             => false;
 
+        public async Task<bool> TryNavigateToSpanAsync(Workspace workspace, DocumentId documentId, TextSpan textSpan, OptionSet options, bool allowInvalidSpan, CancellationToken cancellationToken)
+        {
+            await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+            return TryNavigateToSpan(workspace, documentId, textSpan, options, allowInvalidSpan, cancellationToken);
+        }
+
         public bool TryNavigateToSpan(Workspace workspace, DocumentId documentId, TextSpan textSpan, OptionSet options, bool allowInvalidSpan, CancellationToken cancellationToken)
         {
             if (workspace is not InteractiveWindowWorkspace interactiveWorkspace)

src/EditorFeatures/Core.Wpf/Interactive/InteractiveDocumentNavigationServiceFactory.cs

@@ -2,10 +2,9 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-#nullable disable
-
 using System;
 using System.Composition;
+using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Navigation;
@@ -19,8 +18,8 @@ internal sealed class InteractiveDocumentNavigationServiceFactory : IWorkspaceSe
 
         [ImportingConstructor]
         [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
-        public InteractiveDocumentNavigationServiceFactory()
-            => _singleton = new InteractiveDocumentNavigationService();
+        public InteractiveDocumentNavigationServiceFactory(IThreadingContext threadingContext)
+            => _singleton = new InteractiveDocumentNavigationService(threadingContext);
 
         public IWorkspaceService CreateService(HostWorkspaceServices workspaceServices)
             => _singleton;

src/EditorFeatures/Core.Wpf/Peek/PeekableItemFactory.cs

@@ -67,9 +67,8 @@ public async Task<IEnumerable<IPeekableItem>> GetPeekableItemsAsync(
             var symbolNavigationService = solution.Workspace.Services.GetService<ISymbolNavigationService>();
             var definitionItem = symbol.ToNonClassifiedDefinitionItem(solution, includeHiddenLocations: true);
 
-            if (symbolNavigationService.WouldNavigateToSymbol(
-                    definitionItem, solution, cancellationToken,
-                    out var filePath, out var lineNumber, out var charOffset))
+            var result = await symbolNavigationService.WouldNavigateToSymbolAsync(definitionItem, cancellationToken).ConfigureAwait(false);
+            if (result is var (filePath, lineNumber, charOffset))
             {
                 var position = new LinePosition(lineNumber, charOffset);
                 results.Add(new ExternalFilePeekableItem(new FileLinePositionSpan(filePath, position, position), PredefinedPeekRelationships.Definitions, peekResultFactory));

src/EditorFeatures/Core/FindUsages/AbstractFindUsagesService.DefinitionTrackingContext.cs

@@ -17,7 +17,7 @@ internal abstract partial class AbstractFindUsagesService
         /// Forwards <see cref="IFindUsagesContext"/> notifications to an underlying <see cref="IFindUsagesContext"/>
         /// while also keeping track of the <see cref="DefinitionItem"/> definitions reported.
         /// 
-        /// These can then be used by <see cref="GetThirdPartyDefinitions"/> to report the
+        /// These can then be used by <see cref="GetThirdPartyDefinitionsAsync"/> to report the
         /// definitions found to third parties in case they want to add any additional definitions
         /// to the results we present.
         /// </summary>

src/EditorFeatures/Core/FindUsages/AbstractFindUsagesService.ProgressAdapter.cs

@@ -42,8 +42,8 @@ public async ValueTask OnReferenceFoundAsync(Document document, TextSpan span, C
             {
                 var documentSpan = await ClassifiedSpansAndHighlightSpanFactory.GetClassifiedDocumentSpanAsync(
                     document, span, cancellationToken).ConfigureAwait(false);
-                await _context.OnReferenceFoundAsync(new SourceReferenceItem(
-                    _definition, documentSpan, SymbolUsageInfo.None), cancellationToken).ConfigureAwait(false);
+                await _context.OnReferenceFoundAsync(
+                    new SourceReferenceItem(_definition, documentSpan, SymbolUsageInfo.None), cancellationToken).ConfigureAwait(false);
             }
         }
 

src/EditorFeatures/Core/FindUsages/AbstractFindUsagesService_FindReferences.cs

@@ -10,6 +10,7 @@
 using Microsoft.CodeAnalysis.FindSymbols;
 using Microsoft.CodeAnalysis.FindUsages;
 using Microsoft.CodeAnalysis.LanguageServices;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Remote;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Roslyn.Utilities;
@@ -23,27 +24,16 @@ async Task IFindUsagesService.FindReferencesAsync(
         {
             var definitionTrackingContext = new DefinitionTrackingContext(context);
 
-            // Need ConfigureAwait(true) here so we get back to the UI thread before calling 
-            // GetThirdPartyDefinitions.  We need to call that on the UI thread to match behavior
-            // of how the language service always worked in the past.
-            //
-            // Any async calls before GetThirdPartyDefinitions must be ConfigureAwait(true).
             await FindLiteralOrSymbolReferencesAsync(
-                document, position, definitionTrackingContext, cancellationToken).ConfigureAwait(true);
+                document, position, definitionTrackingContext, cancellationToken).ConfigureAwait(false);
 
             // After the FAR engine is done call into any third party extensions to see
             // if they want to add results.
-            var thirdPartyDefinitions = GetThirdPartyDefinitions(
-                document.Project.Solution, definitionTrackingContext.GetDefinitions(), cancellationToken);
-
-            // From this point on we can do ConfigureAwait(false) as we're not calling back 
-            // into third parties anymore.
+            var thirdPartyDefinitions = await GetThirdPartyDefinitionsAsync(
+                document.Project.Solution, definitionTrackingContext.GetDefinitions(), cancellationToken).ConfigureAwait(false);
 
             foreach (var definition in thirdPartyDefinitions)
-            {
-                // Don't need ConfigureAwait(true) here 
                 await context.OnDefinitionFoundAsync(definition, cancellationToken).ConfigureAwait(false);
-            }
         }
 
         Task IFindUsagesLSPService.FindReferencesAsync(
@@ -74,15 +64,22 @@ await FindSymbolReferencesAsync(
                 document, position, context, cancellationToken).ConfigureAwait(false);
         }
 
-        private static ImmutableArray<DefinitionItem> GetThirdPartyDefinitions(
+        private static async Task<ImmutableArray<DefinitionItem>> GetThirdPartyDefinitionsAsync(
             Solution solution,
             ImmutableArray<DefinitionItem> definitions,
             CancellationToken cancellationToken)
         {
+            using var _ = ArrayBuilder<DefinitionItem>.GetInstance(out var result);
+
             var factory = solution.Workspace.Services.GetRequiredService<IDefinitionsAndReferencesFactory>();
-            return definitions.Select(d => factory.GetThirdPartyDefinitionItem(solution, d, cancellationToken))
-                              .WhereNotNull()
-                              .ToImmutableArray();
+
+            foreach (var definition in definitions)
+            {
+                var thirdParty = await factory.GetThirdPartyDefinitionItemAsync(solution, definition, cancellationToken).ConfigureAwait(false);
+                result.AddIfNotNull(thirdParty);
+            }
+
+            return result.ToImmutable();
         }
 
         private static async Task FindSymbolReferencesAsync(

src/EditorFeatures/Core/FindUsages/IDefinitionsAndReferencesFactory.cs

@@ -16,7 +16,6 @@
 using Microsoft.CodeAnalysis.FindUsages;
 using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Host.Mef;
-using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Collections;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Roslyn.Utilities;
@@ -27,7 +26,7 @@ namespace Microsoft.CodeAnalysis.Editor.FindUsages
 
     internal interface IDefinitionsAndReferencesFactory : IWorkspaceService
     {
-        DefinitionItem? GetThirdPartyDefinitionItem(
+        Task<DefinitionItem?> GetThirdPartyDefinitionItemAsync(
             Solution solution, DefinitionItem definitionItem, CancellationToken cancellationToken);
     }
 
@@ -44,10 +43,10 @@ public DefaultDefinitionsAndReferencesFactory()
         /// Provides an extension point that allows for other workspace layers to add additional
         /// results to the results found by the FindReferences engine.
         /// </summary>
-        public virtual DefinitionItem? GetThirdPartyDefinitionItem(
+        public virtual Task<DefinitionItem?> GetThirdPartyDefinitionItemAsync(
             Solution solution, DefinitionItem definitionItem, CancellationToken cancellationToken)
         {
-            return null;
+            return SpecializedTasks.Null<DefinitionItem>();
         }
     }
 

src/EditorFeatures/Core/GoToDefinition/AbstractGoToDefinitionService.cs

@@ -11,9 +11,11 @@
 using Microsoft.CodeAnalysis.Editor.FindUsages;
 using Microsoft.CodeAnalysis.Editor.Host;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
+using Microsoft.CodeAnalysis.FindUsages;
 using Microsoft.CodeAnalysis.GoToDefinition;
 using Microsoft.CodeAnalysis.LanguageServices;
 using Microsoft.CodeAnalysis.Navigation;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Roslyn.Utilities;
 
@@ -112,16 +114,25 @@ private bool TryGoToAlternativeLocationIfAlreadyOnDefinition(
             if (interfaceImpls.Length == 0)
                 return false;
 
-            var definitions = interfaceImpls.SelectMany(
-                i => GoToDefinitionHelpers.GetDefinitions(
-                    i, solution, thirdPartyNavigationAllowed: false, cancellationToken)).ToImmutableArray();
-
             var title = string.Format(EditorFeaturesResources._0_implemented_members,
                 FindUsagesHelpers.GetDisplayName(symbol));
 
-            return _threadingContext.JoinableTaskFactory.Run(() =>
-                _streamingPresenter.TryNavigateToOrPresentItemsAsync(
-                    _threadingContext, solution.Workspace, title, definitions, cancellationToken));
+            return _threadingContext.JoinableTaskFactory.Run(async () =>
+            {
+                using var _ = ArrayBuilder<DefinitionItem>.GetInstance(out var definitions);
+                foreach (var impl in interfaceImpls)
+                {
+                    // Use ConfigureAwait(true) here.  Not for a correctness requirements, but because we're
+                    // already blocking the UI thread by being in a JTF.Run call.  So we might as well try to
+                    // continue to use the blocking UI thread to do as much work as possible instead of making
+                    // it wait for threadpool threads to be available to process the work.
+                    definitions.AddRange(await GoToDefinitionHelpers.GetDefinitionsAsync(
+                        impl, solution, thirdPartyNavigationAllowed: false, cancellationToken).ConfigureAwait(true));
+                }
+
+                return await _streamingPresenter.TryNavigateToOrPresentItemsAsync(
+                    _threadingContext, solution.Workspace, title, definitions.ToImmutable(), cancellationToken).ConfigureAwait(true);
+            });
         }
 
         private static bool IsThirdPartyNavigationAllowed(ISymbol symbolToNavigateTo, int caretPosition, Document document, CancellationToken cancellationToken)