Find results in related 'partial type parts' when doing a scoped nav-to search.

src/Analyzers/Core/Analyzers/RemoveUnnecessarySuppressions/AbstractRemoveUnnecessaryPragmaSuppressionsDiagnosticAnalyzer.cs

@@ -745,10 +745,11 @@ private async Task<bool> ProcessSuppressMessageAttributesAsync(
             return false;
         }
 
-        using var pooledDeclarationNodes = SyntaxFacts.GetTopLevelAndMethodLevelMembers(root);
-        var declarationNodes = pooledDeclarationNodes.Object;
+        // Specifies false for discardLargeInstances as these objects commonly exceed the default ArrayBuilder capacity threshold.
+        using var _1 = ArrayBuilder<SyntaxNode>.GetInstance(discardLargeInstances: false, out var declarationNodes);
+        this.SyntaxFacts.AddTopLevelAndMethodLevelMembers(root, declarationNodes);
 
-        using var _ = PooledHashSet<ISymbol>.GetInstance(out var processedPartialSymbols);
+        using var _2 = PooledHashSet<ISymbol>.GetInstance(out var processedPartialSymbols);
         if (declarationNodes.Count > 0)
         {
             foreach (var node in declarationNodes)

src/EditorFeatures/CSharpTest/NavigateTo/NavigateToSearcherTests.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Immutable;
 using System.Linq;
 using System.Threading;
@@ -12,6 +13,7 @@
 using Microsoft.CodeAnalysis.NavigateTo;
 using Microsoft.CodeAnalysis.Navigation;
 using Microsoft.CodeAnalysis.PatternMatching;
+using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Test.Utilities;
 using Microsoft.CodeAnalysis.Text;
 using Moq;
@@ -354,6 +356,151 @@ public class D
         Assert.True(searchGeneratedDocumentsAsyncCalled);
     }
 
+    [Fact, WorkItem("https://github.com/dotnet/roslyn/issues/77051")]
+    public async Task DocumentScopeRelatedDocuments_Inheritance()
+    {
+        using var workspace = EditorTestWorkspace.Create("""
+            <Workspace>
+                <Project Language="C#" AssemblyName="Assembly1" CommonReferences="true">
+                    <Document FilePath="z:\\file1.cs">
+                    public class C : Base
+                    {
+                        // Starting search here.
+                        void Goo1() { }
+                    }
+                    </Document>
+                    <Document FilePath="z:\\file2.cs">
+                    public class Base
+                    {
+                        // Should find this.
+                        void Goo2() { }
+                    }
+                    public class Other
+                    {
+                        // Should not find this.
+                        void Goo3() { }
+                    }
+                    </Document>
+                </Project>
+            </Workspace>
+            """, composition: FirstActiveAndVisibleComposition);
+
+        var hostMock = new Mock<INavigateToSearcherHost>(MockBehavior.Strict);
+        hostMock.Setup(h => h.IsFullyLoadedAsync(It.IsAny<CancellationToken>())).Returns(() => new ValueTask<bool>(true));
+
+        var project = workspace.CurrentSolution.Projects.Single();
+        var searchService = project.GetRequiredLanguageService<INavigateToSearchService>();
+
+        hostMock.Setup(h => h.GetNavigateToSearchService(It.IsAny<Project>())).Returns(() => searchService);
+
+        var callback = new TestNavigateToSearchCallback();
+
+        var searcher = NavigateToSearcher.Create(
+            workspace.CurrentSolution,
+            callback,
+            "Goo",
+            kinds: searchService.KindsProvided,
+            hostMock.Object);
+
+        await searcher.SearchAsync(NavigateToSearchScope.Document, CancellationToken.None);
+
+        Assert.Equal(2, callback.Results.Count);
+
+        var firstDocument = project.Documents.Single(d => d.FilePath!.Contains("file1"));
+        var secondDocument = project.Documents.Single(d => d.FilePath!.Contains("file2"));
+
+        var firstDocumentResult = Assert.Single(callback.Results, r => r.NavigableItem.Document.Id == firstDocument.Id);
+        var secondDocumentResult = Assert.Single(callback.Results, r => r.NavigableItem.Document.Id == secondDocument.Id);
+
+        Assert.Equal("Goo1", firstDocumentResult.Name);
+        Assert.Equal("Goo2", secondDocumentResult.Name);
+    }
+
+    [Fact, WorkItem("https://github.com/dotnet/roslyn/issues/77051")]
+    public async Task DocumentScopeRelatedDocuments_Partial()
+    {
+        using var workspace = EditorTestWorkspace.Create("""
+            <Workspace>
+                <Project Language="C#" AssemblyName="Assembly1" CommonReferences="true">
+                    <Document FilePath="z:\\file1.cs">
+                    public partial class C
+                    {
+                        // Starting search here.
+                        void Goo1() { }
+                    }
+                    </Document>
+                    <Document FilePath="z:\\file2.cs">
+                    public class Base
+                    {
+                        // Should not find this.
+                        void Goo2() { }
+                    }
+                    public partial class C
+                    {
+                        // Should find this.
+                        void Goo3() { }
+                    }
+                    </Document>
+                </Project>
+            </Workspace>
+            """, composition: FirstActiveAndVisibleComposition);
+
+        var hostMock = new Mock<INavigateToSearcherHost>(MockBehavior.Strict);
+        hostMock.Setup(h => h.IsFullyLoadedAsync(It.IsAny<CancellationToken>())).Returns(() => new ValueTask<bool>(true));
+
+        var project = workspace.CurrentSolution.Projects.Single();
+        var searchService = project.GetRequiredLanguageService<INavigateToSearchService>();
+
+        hostMock.Setup(h => h.GetNavigateToSearchService(It.IsAny<Project>())).Returns(() => searchService);
+
+        var callback = new TestNavigateToSearchCallback();
+
+        var searcher = NavigateToSearcher.Create(
+            workspace.CurrentSolution,
+            callback,
+            "Goo",
+            kinds: searchService.KindsProvided,
+            hostMock.Object);
+
+        await searcher.SearchAsync(NavigateToSearchScope.Document, CancellationToken.None);
+
+        Assert.Equal(2, callback.Results.Count);
+
+        var firstDocument = project.Documents.Single(d => d.FilePath!.Contains("file1"));
+        var secondDocument = project.Documents.Single(d => d.FilePath!.Contains("file2"));
+
+        var firstDocumentResult = Assert.Single(callback.Results, r => r.NavigableItem.Document.Id == firstDocument.Id);
+        var secondDocumentResult = Assert.Single(callback.Results, r => r.NavigableItem.Document.Id == secondDocument.Id);
+
+        Assert.Equal("Goo1", firstDocumentResult.Name);
+        Assert.Equal("Goo3", secondDocumentResult.Name);
+    }
+
+    private sealed class TestNavigateToSearchCallback : INavigateToSearchCallback
+    {
+        public readonly ConcurrentBag<INavigateToSearchResult> Results = [];
+
+        public void Done(bool isFullyLoaded)
+        {
+        }
+
+        public void ReportIncomplete()
+        {
+        }
+
+        public Task AddResultsAsync(ImmutableArray<INavigateToSearchResult> results, CancellationToken cancellationToken)
+        {
+            foreach (var result in results)
+                this.Results.Add(result);
+
+            return Task.CompletedTask;
+        }
+
+        public void ReportProgress(int current, int maximum)
+        {
+        }
+    }
+
     private sealed class MockAdvancedNavigateToSearchService : IAdvancedNavigateToSearchService
     {
         public IImmutableSet<string> KindsProvided => AbstractNavigateToSearchService.AllKinds;

src/EditorFeatures/Core/Editor/GoToAdjacentMemberCommandHandler.cs

@@ -11,6 +11,7 @@
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.LanguageService;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.Commanding;
@@ -103,12 +104,11 @@ private bool ExecuteCommandImpl(EditorCommandArgs args, bool gotoNextMember, Com
     /// </summary>
     internal static int? GetTargetPosition(ISyntaxFactsService service, SyntaxNode root, int caretPosition, bool next)
     {
-        using var pooledMembers = service.GetMethodLevelMembers(root);
-        var members = pooledMembers.Object;
+        // Specifies false for discardLargeInstances as these objects commonly exceed the default ArrayBuilder capacity threshold.
+        using var _ = ArrayBuilder<SyntaxNode>.GetInstance(discardLargeInstances: false, out var members);
+        service.AddMethodLevelMembers(root, members);
         if (members.Count == 0)
-        {
             return null;
-        }
 
         var starts = members.Select(m => MemberStart(m)).ToArray();
         var index = Array.BinarySearch(starts, caretPosition);

src/Features/CSharp/Portable/NavigateTo/CSharpNavigateToSearchService.cs

@@ -12,6 +12,4 @@ namespace Microsoft.CodeAnalysis.CSharp.NavigateTo;
 [ExportLanguageService(typeof(INavigateToSearchService), LanguageNames.CSharp), Shared]
 [method: ImportingConstructor]
 [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
-internal class CSharpNavigateToSearchService() : AbstractNavigateToSearchService
-{
-}
+internal sealed class CSharpNavigateToSearchService() : AbstractNavigateToSearchService;

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

@@ -8,9 +8,14 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.LanguageService;
 using Microsoft.CodeAnalysis.PatternMatching;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Remote;
+using Microsoft.CodeAnalysis.Shared;
+using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Shared.Utilities;
+using Microsoft.CodeAnalysis.Text;
 using Microsoft.CodeAnalysis.Threading;
 using Roslyn.Utilities;
 
@@ -36,16 +41,16 @@ public async Task SearchDocumentAsync(
             await client.TryInvokeAsync<IRemoteNavigateToSearchService>(
                 document.Project,
                 (service, solutionInfo, callbackId, cancellationToken) =>
-                service.SearchDocumentAsync(solutionInfo, document.Id, searchPattern, [.. kinds], callbackId, cancellationToken),
+                service.SearchDocumentAndRelatedDocumentsAsync(solutionInfo, document.Id, searchPattern, [.. kinds], callbackId, cancellationToken),
                 callback, cancellationToken).ConfigureAwait(false);
 
             return;
         }
 
-        await SearchDocumentInCurrentProcessAsync(document, searchPattern, kinds, onItemsFound, cancellationToken).ConfigureAwait(false);
+        await SearchDocumentAndRelatedDocumentsInCurrentProcessAsync(document, searchPattern, kinds, onItemsFound, cancellationToken).ConfigureAwait(false);
     }
 
-    public static async Task SearchDocumentInCurrentProcessAsync(
+    public static async Task SearchDocumentAndRelatedDocumentsInCurrentProcessAsync(
         Document document,
         string searchPattern,
         IImmutableSet<string> kinds,
@@ -55,12 +60,85 @@ public static async Task SearchDocumentInCurrentProcessAsync(
         var (patternName, patternContainerOpt) = PatternMatcher.GetNameAndContainer(searchPattern);
         var declaredSymbolInfoKindsSet = new DeclaredSymbolInfoKindSet(kinds);
 
-        var results = new ConcurrentSet<RoslynNavigateToItem>();
-        await SearchSingleDocumentAsync(
-            document, patternName, patternContainerOpt, declaredSymbolInfoKindsSet, t => results.Add(t), cancellationToken).ConfigureAwait(false);
+        // In parallel, search both the document requested, and any relevant 'related documents' we find for it. For the
+        // original document, search the entirety of it (by passing 'null' in for the 'spans' argument).  For related
+        // documents, only search the spans of the partial-types/inheriting-types that we find for the types in this
+        // starting document.
+        await Task.WhenAll(
+            SearchDocumentsInCurrentProcessAsync([(document, spans: null)]),
+            SearchRelatedDocumentsInCurrentProcessAsync()).ConfigureAwait(false);
+
+        Task SearchDocumentsInCurrentProcessAsync(ImmutableArray<(Document document, NormalizedTextSpanCollection? spans)> documentAndSpans)
+            => ProducerConsumer<RoslynNavigateToItem>.RunParallelAsync(
+                documentAndSpans,
+                produceItems: static async (documentAndSpan, onItemFound, args, cancellationToken) =>
+                {
+                    var (patternName, patternContainerOpt, declaredSymbolInfoKindsSet, onItemsFound) = args;
+                    await SearchSingleDocumentAsync(
+                        documentAndSpan.document, patternName, patternContainerOpt, declaredSymbolInfoKindsSet,
+                        item =>
+                        {
+                            // Ensure that the results found while searching the single document intersect the desired
+                            // subrange of the document we're searching in.  For the primary document this will always
+                            // succeed (since we're searching the full document).  But for related documents this may fail
+                            // if the results is not in the span of any of the types in those files we're searching.
+                            if (documentAndSpan.spans is null || documentAndSpan.spans.IntersectsWith(item.DeclaredSymbolInfo.Span))
+                                onItemFound(item);
+                        },
+                        cancellationToken).ConfigureAwait(false);
+                },
+                consumeItems: static (values, args, cancellationToken) => args.onItemsFound(values, default, cancellationToken),
+                args: (patternName, patternContainerOpt, declaredSymbolInfoKindsSet, onItemsFound),
+                cancellationToken);
+
+        async Task SearchRelatedDocumentsInCurrentProcessAsync()
+        {
+            var relatedDocuments = await GetRelatedDocumentsAsync().ConfigureAwait(false);
+            await SearchDocumentsInCurrentProcessAsync(relatedDocuments).ConfigureAwait(false);
+        }
 
-        if (results.Count > 0)
-            await onItemsFound([.. results], default, cancellationToken).ConfigureAwait(false);
+        async Task<ImmutableArray<(Document document, NormalizedTextSpanCollection? spans)>> GetRelatedDocumentsAsync()
+        {
+            // For C#/VB we define 'related documents' as those containing types in the inheritance chain of types in
+            // the originating file (as well as all partial parts of the original and inheritance types).  This way a
+            // user can search for symbols scoped to the 'current document' and still get results for the members found
+            // in partial parts.
+
+            var solution = document.Project.Solution;
+            var root = await document.GetRequiredSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
+            var semanticModel = await document.GetRequiredNullableDisabledSemanticModelAsync(cancellationToken).ConfigureAwait(false);
+            var syntaxFacts = document.GetRequiredLanguageService<ISyntaxFactsService>();
+
+            using var _ = ArrayBuilder<SyntaxNode>.GetInstance(out var topLevelNodes);
+            syntaxFacts.AddTopLevelMembers(root, topLevelNodes);
+
+            // Keep track of all of the interesting spans in each document we find. Note: we will convert this to a
+            // NormalizedTextSpanCollection before returning it.  That way the span of an outer partial type will
+            // encompass the span of an inner one and we won't get duplicates for the same symbol.
+            var documentToTextSpans = new MultiDictionary<Document, TextSpan>();
+
+            foreach (var topLevelMember in topLevelNodes)
+            {
+                if (semanticModel.GetDeclaredSymbol(topLevelMember, cancellationToken) is not INamedTypeSymbol namedTypeSymbol)
+                    continue;
+
+                foreach (var type in namedTypeSymbol.GetBaseTypesAndThis())
+                {
+                    foreach (var reference in type.DeclaringSyntaxReferences)
+                    {
+                        var relatedDocument = solution.GetDocument(reference.SyntaxTree);
+                        if (relatedDocument is null)
+                            continue;
+
+                        documentToTextSpans.Add(relatedDocument, reference.Span);
+                    }
+                }
+            }
+
+            // Ensure we don't search the original document we were already searching.
+            documentToTextSpans.Remove(document);
+            return documentToTextSpans.SelectAsArray(kvp => (kvp.Key, new NormalizedTextSpanCollection(kvp.Value)))!;
+        }
     }
 
     public async Task SearchProjectsAsync(

src/Features/Core/Portable/NavigateTo/IRemoteNavigateToSearchService.cs

@@ -18,7 +18,7 @@ namespace Microsoft.CodeAnalysis.NavigateTo;
 
 internal interface IRemoteNavigateToSearchService
 {
-    ValueTask SearchDocumentAsync(Checksum solutionChecksum, DocumentId documentId, string searchPattern, ImmutableArray<string> kinds, RemoteServiceCallbackId callbackId, CancellationToken cancellationToken);
+    ValueTask SearchDocumentAndRelatedDocumentsAsync(Checksum solutionChecksum, DocumentId documentId, string searchPattern, ImmutableArray<string> kinds, RemoteServiceCallbackId callbackId, CancellationToken cancellationToken);
     ValueTask SearchProjectsAsync(Checksum solutionChecksum, ImmutableArray<ProjectId> projectIds, ImmutableArray<DocumentId> priorityDocumentIds, string searchPattern, ImmutableArray<string> kinds, RemoteServiceCallbackId callbackId, CancellationToken cancellationToken);
 
     ValueTask SearchGeneratedDocumentsAsync(Checksum solutionChecksum, ImmutableArray<ProjectId> projectIds, string searchPattern, ImmutableArray<string> kinds, RemoteServiceCallbackId callbackId, CancellationToken cancellationToken);

src/Features/VisualBasic/Portable/NavigateTo/VisualBasicNavigateToSearchService.vb

@@ -8,7 +8,7 @@ Imports Microsoft.CodeAnalysis.NavigateTo
 
 Namespace Microsoft.CodeAnalysis.VisualBasic.NavigateTo
     <ExportLanguageService(GetType(INavigateToSearchService), LanguageNames.VisualBasic), [Shared]>
-    Friend Class VisualBasicNavigateToSearchService
+    Friend NotInheritable Class VisualBasicNavigateToSearchService
         Inherits AbstractNavigateToSearchService
 
         <ImportingConstructor>

src/LanguageServer/Protocol/Features/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer.IncrementalMemberEditAnalyzer.cs

@@ -3,7 +3,6 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
-using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics;
 using System.Linq;
@@ -201,8 +200,10 @@ async Task ExecuteAnalyzersAsync(
                 }
 
                 var syntaxFacts = document.GetRequiredLanguageService<ISyntaxFactsService>();
-                using var pooledMembers = syntaxFacts.GetMethodLevelMembers(root);
-                var members = pooledMembers.Object;
+
+                // Specifies false for discardLargeInstances as these objects commonly exceed the default ArrayBuilder capacity threshold.
+                using var _ = ArrayBuilder<SyntaxNode>.GetInstance(discardLargeInstances: false, out var members);
+                syntaxFacts.AddMethodLevelMembers(root, members);
 
                 var memberSpans = members.SelectAsArray(member => member.FullSpan);
                 var changedMemberId = members.IndexOf(changedMember);