Merge pull request #73992 from CyrusNajmabadi/interceptsLocationWork

Support goto-def taking you from an interceptor method to the location being intercepted.

Author: CyrusNajmabadi

src/EditorFeatures/Core/EditorFeaturesResources.resx

@@ -613,6 +613,9 @@ Do you want to proceed?</value>
   <data name="_0_declarations" xml:space="preserve">
     <value>'{0}' declarations</value>
   </data>
+  <data name="_0_intercepted_locations" xml:space="preserve">
+    <value>'{0}' intercepted locations</value>
+  </data>
   <data name="An_inline_rename_session_is_active_for_identifier_0" xml:space="preserve">
     <value>An inline rename session is active for identifier '{0}'. Invoke inline rename again to access additional options. You may continue to edit the identifier being renamed at any time.</value>
     <comment>For screenreaders. {0} is the identifier being renamed.</comment>

src/EditorFeatures/Core/Navigation/AbstractDefinitionLocationService.cs

@@ -2,30 +2,27 @@
 // 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.Generic;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Classification;
 using Microsoft.CodeAnalysis.Editor.Host;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.FindUsages;
 using Microsoft.CodeAnalysis.GoToDefinition;
 using Microsoft.CodeAnalysis.LanguageService;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Extensions;
+using Microsoft.CodeAnalysis.Shared.Utilities;
 
 namespace Microsoft.CodeAnalysis.Navigation;
 
-internal abstract class AbstractDefinitionLocationService : IDefinitionLocationService
+internal abstract partial class AbstractDefinitionLocationService(
+    IThreadingContext threadingContext,
+    IStreamingFindUsagesPresenter streamingPresenter) : IDefinitionLocationService
 {
-    private readonly IThreadingContext _threadingContext;
-    private readonly IStreamingFindUsagesPresenter _streamingPresenter;
-
-    protected AbstractDefinitionLocationService(
-        IThreadingContext threadingContext,
-        IStreamingFindUsagesPresenter streamingPresenter)
-    {
-        _threadingContext = threadingContext;
-        _streamingPresenter = streamingPresenter;
-    }
+    private readonly IThreadingContext _threadingContext = threadingContext;
+    private readonly IStreamingFindUsagesPresenter _streamingPresenter = streamingPresenter;
 
     private static Task<INavigableLocation?> GetNavigableLocationAsync(
         Document document, int position, CancellationToken cancellationToken)
@@ -40,54 +37,62 @@ protected AbstractDefinitionLocationService(
 
     public async Task<DefinitionLocation?> GetDefinitionLocationAsync(Document document, int position, CancellationToken cancellationToken)
     {
+        var symbolService = document.GetRequiredLanguageService<IGoToDefinitionSymbolService>();
+
         // We want to compute this as quickly as possible so that the symbol be squiggled and navigated to.  We
         // don't want to wait on expensive operations like computing source-generators or skeletons if we can avoid
         // it.  So first try with a frozen document, then fallback to a normal document.  This mirrors how go-to-def
         // works as well.
         return await GetDefinitionLocationWorkerAsync(document.WithFrozenPartialSemantics(cancellationToken)).ConfigureAwait(false) ??
                await GetDefinitionLocationWorkerAsync(document).ConfigureAwait(false);
 
-        async Task<DefinitionLocation?> GetDefinitionLocationWorkerAsync(Document document)
+        async ValueTask<DefinitionLocation?> GetDefinitionLocationWorkerAsync(Document document)
+        {
+            return await GetControlFlowTargetLocationAsync(document).ConfigureAwait(false) ??
+                   await GetSymbolLocationAsync(document).ConfigureAwait(false);
+        }
+
+        async ValueTask<DefinitionLocation?> GetControlFlowTargetLocationAsync(Document document)
         {
-            var symbolService = document.GetRequiredLanguageService<IGoToDefinitionSymbolService>();
             var (controlFlowTarget, controlFlowSpan) = await symbolService.GetTargetIfControlFlowAsync(
                 document, position, cancellationToken).ConfigureAwait(false);
-            if (controlFlowTarget != null)
-            {
-                var location = await GetNavigableLocationAsync(
-                    document, controlFlowTarget.Value, cancellationToken).ConfigureAwait(false);
-                return location is null ? null : new DefinitionLocation(location, new DocumentSpan(document, controlFlowSpan));
-            }
-            else
-            {
-                // Try to compute the referenced symbol and attempt to go to definition for the symbol.
-                var (symbol, project, span) = await symbolService.GetSymbolProjectAndBoundSpanAsync(
-                    document, position, cancellationToken).ConfigureAwait(false);
-                if (symbol is null)
-                    return null;
-
-                // if the symbol only has a single source location, and we're already on it,
-                // try to see if there's a better symbol we could navigate to.
-                var remappedLocation = await GetAlternativeLocationIfAlreadyOnDefinitionAsync(
-                    project, position, symbol, originalDocument: document, cancellationToken).ConfigureAwait(false);
-                if (remappedLocation != null)
-                    return new DefinitionLocation(remappedLocation, new DocumentSpan(document, span));
-
-                var isThirdPartyNavigationAllowed = await IsThirdPartyNavigationAllowedAsync(
-                    symbol, position, document, cancellationToken).ConfigureAwait(false);
-
-                var location = await GoToDefinitionHelpers.GetDefinitionLocationAsync(
-                    symbol,
-                    project.Solution,
-                    _threadingContext,
-                    _streamingPresenter,
-                    thirdPartyNavigationAllowed: isThirdPartyNavigationAllowed,
-                    cancellationToken: cancellationToken).ConfigureAwait(false);
-                if (location is null)
-                    return null;
-
-                return new DefinitionLocation(location, new DocumentSpan(document, span));
-            }
+            if (controlFlowTarget == null)
+                return null;
+
+            var location = await GetNavigableLocationAsync(
+                document, controlFlowTarget.Value, cancellationToken).ConfigureAwait(false);
+            return location is null ? null : new DefinitionLocation(location, new DocumentSpan(document, controlFlowSpan));
+        }
+
+        async ValueTask<DefinitionLocation?> GetSymbolLocationAsync(Document document)
+        {
+            // Try to compute the referenced symbol and attempt to go to definition for the symbol.
+            var (symbol, project, span) = await symbolService.GetSymbolProjectAndBoundSpanAsync(
+                document, position, cancellationToken).ConfigureAwait(false);
+            if (symbol is null)
+                return null;
+
+            // if the symbol only has a single source location, and we're already on it,
+            // try to see if there's a better symbol we could navigate to.
+            var remappedLocation = await GetAlternativeLocationIfAlreadyOnDefinitionAsync(
+                project, position, symbol, originalDocument: document, cancellationToken).ConfigureAwait(false);
+            if (remappedLocation != null)
+                return new DefinitionLocation(remappedLocation, new DocumentSpan(document, span));
+
+            var isThirdPartyNavigationAllowed = await IsThirdPartyNavigationAllowedAsync(
+                symbol, position, document, cancellationToken).ConfigureAwait(false);
+
+            var location = await GoToDefinitionHelpers.GetDefinitionLocationAsync(
+                symbol,
+                project.Solution,
+                _threadingContext,
+                _streamingPresenter,
+                thirdPartyNavigationAllowed: isThirdPartyNavigationAllowed,
+                cancellationToken: cancellationToken).ConfigureAwait(false);
+            if (location is null)
+                return null;
+
+            return new DefinitionLocation(location, new DocumentSpan(document, span));
         }
     }
 
@@ -114,29 +119,118 @@ protected AbstractDefinitionLocationService(
         if (definitionDocument != originalDocument)
             return null;
 
-        // Ok, we were already on the definition. Look for better symbols we could show results
-        // for instead. For now, just see if we're on an interface member impl. If so, we can
-        // instead navigate to the actual interface member.
-        //
-        // In the future we can expand this with other mappings if appropriate.
-        var interfaceImpls = symbol.ExplicitOrImplicitInterfaceImplementations();
-        if (interfaceImpls.Length == 0)
-            return null;
-
-        var title = string.Format(EditorFeaturesResources._0_implemented_members,
-            FindUsagesHelpers.GetDisplayName(symbol));
+        // Ok, we were already on the definition. Look for better symbols we could show results for instead. This can be
+        // expanded with other mappings in the future if appropriate.
+        return await TryGetExplicitInterfaceLocationAsync().ConfigureAwait(false) ??
+               await TryGetInterceptedLocationAsync().ConfigureAwait(false);
 
-        using var _ = ArrayBuilder<DefinitionItem>.GetInstance(out var builder);
-        foreach (var impl in interfaceImpls)
+        async ValueTask<INavigableLocation?> TryGetExplicitInterfaceLocationAsync()
         {
-            builder.AddRange(await GoToDefinitionFeatureHelpers.GetDefinitionsAsync(
-                impl, solution, thirdPartyNavigationAllowed: false, cancellationToken).ConfigureAwait(false));
+            var interfaceImpls = symbol.ExplicitOrImplicitInterfaceImplementations();
+            if (interfaceImpls.Length == 0)
+                return null;
+
+            var title = string.Format(EditorFeaturesResources._0_implemented_members,
+                FindUsagesHelpers.GetDisplayName(symbol));
+
+            using var _ = ArrayBuilder<DefinitionItem>.GetInstance(out var builder);
+            foreach (var impl in interfaceImpls)
+            {
+                builder.AddRange(await GoToDefinitionFeatureHelpers.GetDefinitionsAsync(
+                    impl, solution, thirdPartyNavigationAllowed: false, cancellationToken).ConfigureAwait(false));
+            }
+
+            var definitions = builder.ToImmutable();
+
+            return await _streamingPresenter.GetStreamingLocationAsync(
+                _threadingContext, solution.Workspace, title, definitions, cancellationToken).ConfigureAwait(false);
         }
 
-        var definitions = builder.ToImmutable();
+        async ValueTask<INavigableLocation?> TryGetInterceptedLocationAsync()
+        {
+            if (symbol is not IMethodSymbol method)
+                return null;
+
+            // Find attributes of the form: [InterceptsLocationAttribute(version: 1, data: "...")];
+
+            var attributes = method.GetAttributes();
+            var interceptsLocationDatas = InterceptsLocationUtilities.GetInterceptsLocationData(attributes);
+            if (interceptsLocationDatas.Length == 0)
+                return null;
+
+            using var _ = ArrayBuilder<DocumentSpan>.GetInstance(out var documentSpans);
+
+            foreach (var (contentHash, position) in interceptsLocationDatas)
+            {
+                var document = await project.GetDocumentAsync(contentHash, cancellationToken).ConfigureAwait(false);
+
+                if (document != null)
+                {
+                    var root = await document.GetRequiredSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
 
-        return await _streamingPresenter.GetStreamingLocationAsync(
-            _threadingContext, solution.Workspace, title, definitions, cancellationToken).ConfigureAwait(false);
+                    if (position >= 0 && position < root.FullSpan.Length)
+                    {
+                        var token = root.FindToken(position);
+                        documentSpans.Add(new DocumentSpan(document, token.Span));
+                    }
+                }
+            }
+
+            documentSpans.RemoveDuplicates();
+
+            if (documentSpans.Count == 0)
+            {
+                return null;
+            }
+            else if (documentSpans.Count == 1)
+            {
+                // Just one document span this mapped to.  Navigate directly do that.
+                return await documentSpans[0].GetNavigableLocationAsync(cancellationToken).ConfigureAwait(false);
+            }
+            else
+            {
+                var title = string.Format(EditorFeaturesResources._0_intercepted_locations,
+                    FindUsagesHelpers.GetDisplayName(method));
+
+                var definitionItem = method.ToNonClassifiedDefinitionItem(solution, includeHiddenLocations: true);
+
+                var referenceItems = new List<SourceReferenceItem>(capacity: documentSpans.Count);
+                var classificationOptions = ClassificationOptions.Default with { ClassifyObsoleteSymbols = false };
+                foreach (var documentSpan in documentSpans)
+                {
+                    var classifiedSpans = await ClassifiedSpansAndHighlightSpanFactory.ClassifyAsync(
+                        documentSpan, classifiedSpans: null, classificationOptions, cancellationToken).ConfigureAwait(false);
+
+                    referenceItems.Add(new SourceReferenceItem(
+                        definitionItem, documentSpan, classifiedSpans, SymbolUsageInfo.None, additionalProperties: []));
+                }
+
+                // Multiple document spans this mapped to.  Show them all.
+                return new NavigableLocation(async (options, cancellationToken) =>
+                {
+                    // Can only navigate or present items on UI thread.
+                    await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+
+                    // We have multiple definitions, or we have definitions with multiple locations. Present this to the
+                    // user so they can decide where they want to go to.
+                    //
+                    // We ignore the cancellation token returned by StartSearch as we're in a context where
+                    // we've computed all the results and we're synchronously populating the UI with it.
+                    var (context, _) = _streamingPresenter.StartSearch(title, new StreamingFindUsagesPresenterOptions(SupportsReferences: true));
+                    try
+                    {
+                        await context.OnDefinitionFoundAsync(definitionItem, cancellationToken).ConfigureAwait(false);
+                        await context.OnReferencesFoundAsync(referenceItems.AsAsyncEnumerable(), cancellationToken).ConfigureAwait(false);
+                    }
+                    finally
+                    {
+                        await context.OnCompletedAsync(cancellationToken).ConfigureAwait(false);
+                    }
+
+                    return true;
+                });
+            }
+        }
     }
 
     private static async Task<bool> IsThirdPartyNavigationAllowedAsync(