Unify Peek behavior with GoToDef

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

@@ -2,45 +2,46 @@
 // 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.Collections.Generic;
 using System.Linq;
 using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.MetadataAsSource;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.VisualStudio.Language.Intellisense;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Editor.Implementation.Peek
 {
-    internal class DefinitionPeekableItem : PeekableItem
+    internal sealed class DefinitionPeekableItem : PeekableItem
     {
         private readonly Workspace _workspace;
         private readonly ProjectId _projectId;
         private readonly SymbolKey _symbolKey;
         private readonly IMetadataAsSourceFileService _metadataAsSourceFileService;
         private readonly IGlobalOptionService _globalOptions;
+        private readonly IThreadingContext _threadingContext;
 
         public DefinitionPeekableItem(
             Workspace workspace, ProjectId projectId, SymbolKey symbolKey,
             IPeekResultFactory peekResultFactory,
             IMetadataAsSourceFileService metadataAsSourceService,
-            IGlobalOptionService globalOptions)
+            IGlobalOptionService globalOptions,
+            IThreadingContext threadingContext)
             : base(peekResultFactory)
         {
             _workspace = workspace;
             _projectId = projectId;
             _symbolKey = symbolKey;
             _metadataAsSourceFileService = metadataAsSourceService;
             _globalOptions = globalOptions;
+            _threadingContext = threadingContext;
         }
 
-        public override IEnumerable<IPeekRelationship> Relationships
-        {
-            get { return SpecializedCollections.SingletonEnumerable(PredefinedPeekRelationships.Definitions); }
-        }
+        public override IEnumerable<IPeekRelationship> Relationships =>
+            SpecializedCollections.SingletonEnumerable(PredefinedPeekRelationships.Definitions);
 
         public override IPeekResultSource GetOrCreateResultSource(string relationshipName)
             => new ResultSource(this);
@@ -55,38 +56,45 @@ public ResultSource(DefinitionPeekableItem peekableItem)
             public void FindResults(string relationshipName, IPeekResultCollection resultCollection, CancellationToken cancellationToken, IFindPeekResultsCallback callback)
             {
                 if (relationshipName != PredefinedPeekRelationships.Definitions.Name)
-                {
                     return;
-                }
 
                 // Note: this is called on a background thread, but we must block the thread since the API doesn't support proper asynchrony.
+                var success = _peekableItem._threadingContext.JoinableTaskFactory.Run(async () => await FindResultsAsync(
+                    resultCollection, callback, cancellationToken).ConfigureAwait(false));
+                if (!success)
+                    callback.ReportFailure(new Exception(EditorFeaturesResources.No_information_found));
+            }
+
+            private async Task<bool> FindResultsAsync(IPeekResultCollection resultCollection, IFindPeekResultsCallback callback, CancellationToken cancellationToken)
+            {
                 var workspace = _peekableItem._workspace;
                 var solution = workspace.CurrentSolution;
                 var project = solution.GetProject(_peekableItem._projectId);
-                var compilation = project.GetCompilationAsync(cancellationToken).WaitAndGetResult_CanCallOnBackground(cancellationToken);
+                if (project is null)
+                    return false;
+
+                var compilation = await project.GetCompilationAsync(cancellationToken).ConfigureAwait(false);
+                if (compilation is null)
+                    return false;
 
                 var symbol = _peekableItem._symbolKey.Resolve(compilation, ignoreAssemblyKey: true, cancellationToken: cancellationToken).Symbol;
                 if (symbol == null)
-                {
-                    callback.ReportFailure(new Exception(EditorFeaturesResources.No_information_found));
-                    return;
-                }
+                    return false;
 
                 var sourceLocations = symbol.Locations.Where(l => l.IsInSource).ToList();
 
-                if (!sourceLocations.Any())
+                if (sourceLocations.Count == 0)
                 {
                     // It's a symbol from metadata, so we want to go produce it from metadata
                     var options = _peekableItem._globalOptions.GetMetadataAsSourceOptions(project.Services);
-                    var declarationFile = _peekableItem._metadataAsSourceFileService.GetGeneratedFileAsync(workspace, project, symbol, signaturesOnly: false, options, cancellationToken).WaitAndGetResult_CanCallOnBackground(cancellationToken);
+                    var declarationFile = await _peekableItem._metadataAsSourceFileService.GetGeneratedFileAsync(workspace, project, symbol, signaturesOnly: false, options, cancellationToken).ConfigureAwait(false);
                     var peekDisplayInfo = new PeekResultDisplayInfo(declarationFile.DocumentTitle, declarationFile.DocumentTooltip, declarationFile.DocumentTitle, declarationFile.DocumentTooltip);
                     var identifierSpan = declarationFile.IdentifierLocation.GetLineSpan().Span;
                     var entityOfInterestSpan = PeekHelpers.GetEntityOfInterestSpan(symbol, workspace, declarationFile.IdentifierLocation, cancellationToken);
                     resultCollection.Add(PeekHelpers.CreateDocumentPeekResult(declarationFile.FilePath, identifierSpan, entityOfInterestSpan, peekDisplayInfo, _peekableItem.PeekResultFactory, isReadOnly: true));
                 }
 
                 var processedSourceLocations = 0;
-
                 foreach (var declaration in sourceLocations)
                 {
                     var declarationLocation = declaration.GetMappedLineSpan();
@@ -95,6 +103,8 @@ public void FindResults(string relationshipName, IPeekResultCollection resultCol
                     resultCollection.Add(PeekHelpers.CreateDocumentPeekResult(declarationLocation.Path, declarationLocation.Span, entityOfInterestSpan, _peekableItem.PeekResultFactory));
                     callback.ReportProgress(100 * ++processedSourceLocations / sourceLocations.Count);
                 }
+
+                return true;
             }
         }
     }

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

@@ -6,17 +6,21 @@
 
 using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.ComponentModel.Composition;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Editor.Peek;
+using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.FindSymbols;
 using Microsoft.CodeAnalysis.FindUsages;
+using Microsoft.CodeAnalysis.GoToDefinition;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.MetadataAsSource;
 using Microsoft.CodeAnalysis.Navigation;
 using Microsoft.CodeAnalysis.Options;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.VisualStudio.Language.Intellisense;
 
 namespace Microsoft.CodeAnalysis.Editor.Implementation.Peek
@@ -26,52 +30,51 @@ internal class PeekableItemFactory : IPeekableItemFactory
     {
         private readonly IMetadataAsSourceFileService _metadataAsSourceFileService;
         private readonly IGlobalOptionService _globalOptions;
+        private readonly IThreadingContext _threadingContext;
 
         [ImportingConstructor]
         [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
-        public PeekableItemFactory(IMetadataAsSourceFileService metadataAsSourceFileService, IGlobalOptionService globalOptions)
+        public PeekableItemFactory(
+            IMetadataAsSourceFileService metadataAsSourceFileService,
+            IGlobalOptionService globalOptions,
+            IThreadingContext threadingContext)
         {
             _metadataAsSourceFileService = metadataAsSourceFileService;
             _globalOptions = globalOptions;
+            _threadingContext = threadingContext;
         }
 
         public async Task<IEnumerable<IPeekableItem>> GetPeekableItemsAsync(
-            ISymbol symbol, Project project,
+            ISymbol symbol,
+            Project project,
             IPeekResultFactory peekResultFactory,
             CancellationToken cancellationToken)
         {
             if (symbol == null)
-            {
                 throw new ArgumentNullException(nameof(symbol));
-            }
 
             if (project == null)
-            {
                 throw new ArgumentNullException(nameof(project));
-            }
 
             if (peekResultFactory == null)
-            {
                 throw new ArgumentNullException(nameof(peekResultFactory));
-            }
-
-            var results = new List<IPeekableItem>();
 
             var solution = project.Solution;
-            var sourceDefinition = await SymbolFinder.FindSourceDefinitionAsync(symbol, solution, cancellationToken).ConfigureAwait(false);
+            symbol = await SymbolFinder.FindSourceDefinitionAsync(symbol, solution, cancellationToken).ConfigureAwait(false) ?? symbol;
+            symbol = await GoToDefinitionHelpers.TryGetPreferredSymbolAsync(solution, symbol, cancellationToken).ConfigureAwait(false);
+            if (symbol is null)
+                return ImmutableArray<IPeekableItem>.Empty;
 
-            // 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);
-
-                project = originatingProject ?? project;
-            }
+            // if we mapped the symbol, then get the new project it is contained in.
+            var originatingProject = solution.GetProject(symbol.ContainingAssembly, cancellationToken);
+            project = originatingProject ?? project;
 
-            var symbolNavigationService = solution.Services.GetService<ISymbolNavigationService>();
             var definitionItem = symbol.ToNonClassifiedDefinitionItem(solution, includeHiddenLocations: true);
 
+            var symbolNavigationService = solution.Services.GetService<ISymbolNavigationService>();
             var result = await symbolNavigationService.GetExternalNavigationSymbolLocationAsync(definitionItem, cancellationToken).ConfigureAwait(false);
+
+            using var _ = ArrayBuilder<IPeekableItem>.GetInstance(out var results);
             if (result is var (filePath, linePosition))
             {
                 results.Add(new ExternalFilePeekableItem(new FileLinePositionSpan(filePath, linePosition, linePosition), PredefinedPeekRelationships.Definitions, peekResultFactory));
@@ -85,12 +88,13 @@ public async Task<IEnumerable<IPeekableItem>> GetPeekableItemsAsync(
                 {
                     if (firstLocation.IsInSource || _metadataAsSourceFileService.IsNavigableMetadataSymbol(symbol))
                     {
-                        results.Add(new DefinitionPeekableItem(solution.Workspace, project.Id, symbolKey, peekResultFactory, _metadataAsSourceFileService, _globalOptions));
+                        results.Add(new DefinitionPeekableItem(
+                            solution.Workspace, project.Id, symbolKey, peekResultFactory, _metadataAsSourceFileService, _globalOptions, _threadingContext));
                     }
                 }
             }
 
-            return results;
+            return results.ToImmutable();
         }
     }
 }

src/EditorFeatures/Core/GoToDefinition/GoToDefinitionHelpers.cs

@@ -4,6 +4,7 @@
 
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -21,45 +22,14 @@ namespace Microsoft.CodeAnalysis.GoToDefinition
     internal static class GoToDefinitionHelpers
     {
         public static async Task<ImmutableArray<DefinitionItem>> GetDefinitionsAsync(
-            ISymbol symbol,
+            ISymbol? symbol,
             Solution solution,
             bool thirdPartyNavigationAllowed,
             CancellationToken cancellationToken)
         {
-            var alias = symbol as IAliasSymbol;
-            if (alias != null)
-            {
-                if (alias.Target is INamespaceSymbol ns && ns.IsGlobalNamespace)
-                {
-                    return ImmutableArray.Create<DefinitionItem>();
-                }
-            }
-
-            // VB global import aliases have a synthesized SyntaxTree.
-            // We can't go to the definition of the alias, so use the target type.
-
-            if (alias != null)
-            {
-                var sourceLocations = NavigableItemFactory.GetPreferredSourceLocations(
-                    solution, symbol, cancellationToken);
-
-                if (sourceLocations.All(l => solution.GetDocument(l.SourceTree) == null))
-                {
-                    symbol = alias.Target;
-                }
-            }
-
-            var definition = await SymbolFinder.FindSourceDefinitionAsync(symbol, solution, cancellationToken).ConfigureAwait(false);
-            cancellationToken.ThrowIfCancellationRequested();
-
-            symbol = definition ?? symbol;
-
-            // If it is a partial method declaration with no body, choose to go to the implementation
-            // that has a method body.
-            if (symbol is IMethodSymbol method)
-            {
-                symbol = method.PartialImplementationPart ?? symbol;
-            }
+            symbol = await TryGetPreferredSymbolAsync(solution, symbol, cancellationToken).ConfigureAwait(false);
+            if (symbol is null)
+                return ImmutableArray.Create<DefinitionItem>();
 
             using var _ = ArrayBuilder<DefinitionItem>.GetInstance(out var definitions);
 
@@ -98,6 +68,44 @@ public static async Task<ImmutableArray<DefinitionItem>> GetDefinitionsAsync(
             return definitions.ToImmutable();
         }
 
+        public static async Task<ISymbol?> TryGetPreferredSymbolAsync(
+            Solution solution, ISymbol? symbol, CancellationToken cancellationToken)
+        {
+            // VB global import aliases have a synthesized SyntaxTree.
+            // We can't go to the definition of the alias, so use the target type.
+
+            var alias = symbol as IAliasSymbol;
+            if (alias != null)
+            {
+                if (alias.Target is INamespaceSymbol ns && ns.IsGlobalNamespace)
+                    return null;
+            }
+
+            // VB global import aliases have a synthesized SyntaxTree.
+            // We can't go to the definition of the alias, so use the target type.
+
+            if (alias != null)
+            {
+                var sourceLocations = NavigableItemFactory.GetPreferredSourceLocations(
+                    solution, symbol, cancellationToken);
+
+                if (sourceLocations.All(l => solution.GetDocument(l.SourceTree) == null))
+                    symbol = alias.Target;
+            }
+
+            var definition = await SymbolFinder.FindSourceDefinitionAsync(symbol, solution, cancellationToken).ConfigureAwait(false);
+            cancellationToken.ThrowIfCancellationRequested();
+
+            symbol = definition ?? symbol;
+
+            // If it is a partial method declaration with no body, choose to go to the implementation
+            // that has a method body.
+            if (symbol is IMethodSymbol method)
+                symbol = method.PartialImplementationPart ?? symbol;
+
+            return symbol;
+        }
+
         public static async Task<bool> TryNavigateToLocationAsync(
             ISymbol symbol,
             Solution solution,

src/EditorFeatures/Test2/Peek/PeekTests.vb

@@ -181,7 +181,6 @@ End Module
                 Assert.Equal(1, result.Items.Count)
                 result.AssertNavigatesToIdentifier(0, "Identifier")
             End Using
-
         End Sub
 
         <WpfFact, WorkItem(820363, "http://vstfdevdiv:8080/DevDiv2/DevDiv/_workitems/edit/820363")>
@@ -223,6 +222,36 @@ public class Component
             End Using
         End Sub
 
+        <WpfFact, WorkItem(64615, "https://github.com/dotnet/roslyn/issues/64615")>
+        Public Sub TestPartialMethods()
+            Using workspace = CreateTestWorkspace(<Workspace>
+                                                      <Project Language="C#" CommonReferences="true">
+                                                          <Document><![CDATA[
+public partial class D
+{
+    public void M()
+    {
+        $$PartialMethod();
+    }
+
+    partial void PartialMethod();
+}
+                                                          ]]></Document>
+                                                          <Document><![CDATA[
+public partial class D
+{
+    partial void {|Identifier:PartialMethod|}() { }
+}
+                                                          ]]></Document>
+                                                      </Project>
+                                                  </Workspace>)
+                Dim result = GetPeekResultCollection(workspace)
+
+                Assert.Equal(1, result.Items.Count)
+                result.AssertNavigatesToIdentifier(0, "Identifier")
+            End Using
+        End Sub
+
         Private Shared Function CreateTestWorkspace(element As XElement) As TestWorkspace
             Return TestWorkspace.Create(element, composition:=EditorTestCompositions.EditorFeaturesWpf)
         End Function