From 4818ec8cfb0c7acb111e5bbc9a71b55fcac4cdd2 Mon Sep 17 00:00:00 2001
From: Cyrus Najmabadi <cyrus.najmabadi@gmail.com>
Date: Thu, 11 Jul 2024 10:54:25 -0700
Subject: [PATCH 1/9] Simplify passing around of FAR symbol groups

---
 .../FindReferencesSearchEngine.cs             | 75 +++++++++++--------
 ...sSearchEngine_FindReferencesInDocuments.cs | 39 ++++------
 2 files changed, 60 insertions(+), 54 deletions(-)

diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
index 451eae0b79cd0..d47e85b0aeaea 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
@@ -23,6 +23,13 @@ namespace Microsoft.CodeAnalysis.FindSymbols;
 
 internal partial class FindReferencesSearchEngine
 {
+    /// <summary>
+    /// Scheduler we use when we're doing operations in the BG and we want to rate limit them to not saturate the threadpool.
+    /// </summary>
+    private static readonly TaskScheduler s_exclusiveScheduler = new ConcurrentExclusiveSchedulerPair().ExclusiveScheduler;
+
+    private static readonly ObjectPool<Dictionary<ISymbol, SymbolGroup>> s_symbolToGroupPool = new(() => new(MetadataUnifyingEquivalenceComparer.Instance));
+
     private readonly Solution _solution;
     private readonly IImmutableSet<Document>? _documents;
     private readonly ImmutableArray<IReferenceFinder> _finders;
@@ -30,8 +37,6 @@ internal partial class FindReferencesSearchEngine
     private readonly IStreamingFindReferencesProgress _progress;
     private readonly FindReferencesSearchOptions _options;
 
-    private static readonly TaskScheduler s_exclusiveScheduler = new ConcurrentExclusiveSchedulerPair().ExclusiveScheduler;
-
     /// <summary>
     /// Mapping from symbols (unified across metadata/retargeting) and the set of symbols that was produced for 
     /// them in the case of linked files across projects.  This allows references to be found to any of the unified
@@ -124,7 +129,7 @@ await RoslynParallel.ForEachAsync(
                 tuple.project, tuple.allSymbols, onReferenceFound, cancellationToken).ConfigureAwait(false)).ConfigureAwait(false);
     }
 
-    private async IAsyncEnumerable<(Project project, ImmutableArray<ISymbol> allSymbols)> GetProjectsAndSymbolsToSearchAsync(
+    private async IAsyncEnumerable<(Project project, ImmutableArray<(ISymbol symbol, SymbolGroup group)> allSymbols)> GetProjectsAndSymbolsToSearchAsync(
         SymbolSet symbolSet,
         ImmutableArray<Project> projectsToSearch,
         [EnumeratorCancellation] CancellationToken cancellationToken)
@@ -146,12 +151,11 @@ await RoslynParallel.ForEachAsync(
             // which is why we do it in this loop and not inside the concurrent project processing that happens
             // below.
             await symbolSet.InheritanceCascadeAsync(currentProject, cancellationToken).ConfigureAwait(false);
-            var allSymbols = symbolSet.GetAllSymbols();
 
             // Report any new symbols we've cascaded to to our caller.
-            await ReportGroupsAsync(allSymbols, cancellationToken).ConfigureAwait(false);
+            var allSymbolsAndGroups = await ReportGroupsAsync(symbolSet.GetAllSymbols(), cancellationToken).ConfigureAwait(false);
 
-            yield return (currentProject, allSymbols);
+            yield return (currentProject, allSymbolsAndGroups);
         }
     }
 
@@ -160,10 +164,15 @@ await RoslynParallel.ForEachAsync(
     /// them once per symbol group, but we may have to notify about new symbols each time we expand our symbol set
     /// when we walk into a new project.
     /// </summary>
-    private async Task ReportGroupsAsync(ImmutableArray<ISymbol> symbols, CancellationToken cancellationToken)
+    private async Task<ImmutableArray<(ISymbol symbol, SymbolGroup group)>> ReportGroupsAsync(
+        ImmutableArray<ISymbol> symbols, CancellationToken cancellationToken)
     {
+        var result = new FixedSizeArrayBuilder<(ISymbol symbol, SymbolGroup group)>(symbols.Length);
+
         foreach (var symbol in symbols)
-            await ReportGroupAsync(symbol, cancellationToken).ConfigureAwait(false);
+            result.Add((symbol, await ReportGroupAsync(symbol, cancellationToken).ConfigureAwait(false)));
+
+        return result.MoveToImmutable();
     }
 
     private async ValueTask<SymbolGroup> ReportGroupAsync(ISymbol symbol, CancellationToken cancellationToken)
@@ -206,10 +215,10 @@ private Task<ImmutableArray<Project>> GetProjectsToSearchAsync(
     }
 
     private async ValueTask ProcessProjectAsync(
-        Project project, ImmutableArray<ISymbol> allSymbols, Action<Reference> onReferenceFound, CancellationToken cancellationToken)
+        Project project, ImmutableArray<(ISymbol symbol, SymbolGroup group)> allSymbols, Action<Reference> onReferenceFound, CancellationToken cancellationToken)
     {
         using var _1 = PooledDictionary<ISymbol, PooledHashSet<string>>.GetInstance(out var symbolToGlobalAliases);
-        using var _2 = PooledDictionary<Document, MetadataUnifyingSymbolHashSet>.GetInstance(out var documentToSymbols);
+        using var _2 = PooledDictionary<Document, Dictionary<ISymbol, SymbolGroup>>.GetInstance(out var documentToSymbolAndGroup);
         try
         {
             // scratch hashset to place results in. Populated/inspected/cleared in inner loop.
@@ -217,7 +226,7 @@ private async ValueTask ProcessProjectAsync(
 
             await AddGlobalAliasesAsync(project, allSymbols, symbolToGlobalAliases, cancellationToken).ConfigureAwait(false);
 
-            foreach (var symbol in allSymbols)
+            foreach (var (symbol, group) in allSymbols)
             {
                 var globalAliases = TryGet(symbolToGlobalAliases, symbol);
 
@@ -230,30 +239,38 @@ await finder.DetermineDocumentsToSearchAsync(
                         _options, cancellationToken).ConfigureAwait(false);
 
                     foreach (var document in foundDocuments)
-                        GetSymbolSet(documentToSymbols, document).Add(symbol);
+                    {
+                        if (!documentToSymbolAndGroup.TryGetValue(document, out var symbolAndGroup))
+                        {
+                            symbolAndGroup = s_symbolToGroupPool.AllocateAndClear();
+                            documentToSymbolAndGroup.Add(document, symbolAndGroup);
+                        }
+
+                        symbolAndGroup[symbol] = group;
+                    }
 
                     foundDocuments.Clear();
                 }
             }
 
             await RoslynParallel.ForEachAsync(
-                documentToSymbols,
+                documentToSymbolAndGroup,
                 GetParallelOptions(cancellationToken),
                 (kvp, cancellationToken) =>
                     ProcessDocumentAsync(kvp.Key, kvp.Value, symbolToGlobalAliases, onReferenceFound, cancellationToken)).ConfigureAwait(false);
         }
         finally
         {
-            foreach (var (_, symbols) in documentToSymbols)
-                MetadataUnifyingSymbolHashSet.ClearAndFree(symbols);
+            foreach (var (_, symbolAndGroupMap) in documentToSymbolAndGroup)
+            {
+                symbolAndGroupMap.Clear();
+                s_symbolToGroupPool.Free(symbolAndGroupMap);
+            }
 
             FreeGlobalAliases(symbolToGlobalAliases);
 
             await _progressTracker.ItemCompletedAsync(cancellationToken).ConfigureAwait(false);
         }
-
-        static MetadataUnifyingSymbolHashSet GetSymbolSet<T>(PooledDictionary<T, MetadataUnifyingSymbolHashSet> dictionary, T key) where T : notnull
-            => dictionary.GetOrAdd(key, static _ => MetadataUnifyingSymbolHashSet.AllocateFromPool());
     }
 
     private static PooledHashSet<U>? TryGet<T, U>(Dictionary<T, PooledHashSet<U>> dictionary, T key) where T : notnull
@@ -261,7 +278,7 @@ static MetadataUnifyingSymbolHashSet GetSymbolSet<T>(PooledDictionary<T, Metadat
 
     private async ValueTask ProcessDocumentAsync(
         Document document,
-        MetadataUnifyingSymbolHashSet symbols,
+        Dictionary<ISymbol, SymbolGroup> symbolToSymbolGroup,
         Dictionary<ISymbol, PooledHashSet<string>> symbolToGlobalAliases,
         Action<Reference> onReferenceFound,
         CancellationToken cancellationToken)
@@ -281,39 +298,37 @@ private async ValueTask ProcessDocumentAsync(
         // Note: cascaded symbols will normally have the same name.  That's ok.  The second call to
         // FindMatchingIdentifierTokens with the same name will short circuit since it will already see the result of
         // the prior call.
-        foreach (var symbol in symbols)
+        foreach (var (symbol, _) in symbolToSymbolGroup)
         {
             if (symbol.CanBeReferencedByName)
                 cache.FindMatchingIdentifierTokens(symbol.Name, cancellationToken);
         }
 
         await RoslynParallel.ForEachAsync(
-            symbols,
+            symbolToSymbolGroup,
             GetParallelOptions(cancellationToken),
-            (symbol, cancellationToken) =>
+            (kvp, cancellationToken) =>
             {
+                var (symbol, group) = kvp;
+
                 // symbolToGlobalAliases is safe to read in parallel.  It is created fully before this point and is no
                 // longer mutated.
                 var state = new FindReferencesDocumentState(
                     cache, TryGet(symbolToGlobalAliases, symbol));
 
-                ProcessDocument(symbol, state, onReferenceFound);
+                ProcessDocument(symbol, group, state, onReferenceFound);
                 return ValueTaskFactory.CompletedTask;
             }).ConfigureAwait(false);
 
         return;
 
         void ProcessDocument(
-            ISymbol symbol, FindReferencesDocumentState state, Action<Reference> onReferenceFound)
+            ISymbol symbol, SymbolGroup group, FindReferencesDocumentState state, Action<Reference> onReferenceFound)
         {
             cancellationToken.ThrowIfCancellationRequested();
 
             using (Logger.LogBlock(FunctionId.FindReference_ProcessDocumentAsync, cancellationToken))
             {
-                // This is safe to just blindly read. We can only ever get here after the call to ReportGroupsAsync
-                // happened.  So there must be a group for this symbol in our map.
-                var group = _symbolToGroup[symbol];
-
                 // Note: nearly every finder will no-op when passed a in a symbol it's not applicable to.  So it's
                 // simple to just iterate over all of them, knowing that will quickly skip all the irrelevant ones,
                 // and only do interesting work on the single relevant one.
@@ -332,11 +347,11 @@ void ProcessDocument(
 
     private async Task AddGlobalAliasesAsync(
         Project project,
-        ImmutableArray<ISymbol> allSymbols,
+        ImmutableArray<(ISymbol symbol, SymbolGroup group)> allSymbols,
         PooledDictionary<ISymbol, PooledHashSet<string>> symbolToGlobalAliases,
         CancellationToken cancellationToken)
     {
-        foreach (var symbol in allSymbols)
+        foreach (var (symbol, _) in allSymbols)
         {
             foreach (var finder in _finders)
             {
diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
index e17a6dae28e09..bf11d643ec6c6 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
@@ -7,14 +7,12 @@
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
-using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.FindSymbols.Finders;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Shared.Utilities;
-using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.FindSymbols;
 
@@ -43,8 +41,7 @@ public async Task FindReferencesInDocumentsAsync(
         // Create and report the initial set of symbols to search for.  This includes linked and cascaded symbols. It does
         // not walk up/down the inheritance hierarchy.
         var symbolSet = await SymbolSet.DetermineInitialSearchSymbolsAsync(this, unifiedSymbols, cancellationToken).ConfigureAwait(false);
-        var allSymbols = symbolSet.ToImmutableArray();
-        await ReportGroupsAsync(allSymbols, cancellationToken).ConfigureAwait(false);
+        var allSymbolsAndGroups = await ReportGroupsAsync([.. symbolSet], cancellationToken).ConfigureAwait(false);
 
         // Process projects in dependency graph order so that any compilations built by one are available for later
         // projects. We only have to examine the projects containing the documents requested though.
@@ -55,12 +52,12 @@ public async Task FindReferencesInDocumentsAsync(
         {
             var currentProject = _solution.GetRequiredProject(projectId);
             if (projectsToSearch.Contains(currentProject))
-                await PerformSearchInProjectAsync(allSymbols, currentProject).ConfigureAwait(false);
+                await PerformSearchInProjectAsync(allSymbolsAndGroups, currentProject).ConfigureAwait(false);
         }
 
         return;
 
-        async ValueTask PerformSearchInProjectAsync(ImmutableArray<ISymbol> symbols, Project project)
+        async ValueTask PerformSearchInProjectAsync(ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols, Project project)
         {
             using var _ = PooledDictionary<ISymbol, PooledHashSet<string>>.GetInstance(out var symbolToGlobalAliases);
             try
@@ -82,7 +79,7 @@ async ValueTask PerformSearchInProjectAsync(ImmutableArray<ISymbol> symbols, Pro
         }
 
         async ValueTask PerformSearchInDocumentAsync(
-            ImmutableArray<ISymbol> symbols,
+            ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols,
             Document document,
             PooledDictionary<ISymbol, PooledHashSet<string>> symbolToGlobalAliases)
         {
@@ -92,32 +89,32 @@ async ValueTask PerformSearchInDocumentAsync(
             // of this call.
             var cache = await FindReferenceCache.GetCacheAsync(document, cancellationToken).ConfigureAwait(false);
 
-            foreach (var symbol in symbols)
+            foreach (var (symbol, group) in symbols)
             {
                 var state = new FindReferencesDocumentState(
                     cache, TryGet(symbolToGlobalAliases, symbol));
 
-                await PerformSearchInDocumentWorkerAsync(symbol, state).ConfigureAwait(false);
+                await PerformSearchInDocumentWorkerAsync(symbol, group, state).ConfigureAwait(false);
             }
         }
 
-        async ValueTask PerformSearchInDocumentWorkerAsync(ISymbol symbol, FindReferencesDocumentState state)
+        async ValueTask PerformSearchInDocumentWorkerAsync(ISymbol symbol, SymbolGroup group, FindReferencesDocumentState state)
         {
             // Always perform a normal search, looking for direct references to exactly that symbol.
-            await DirectSymbolSearchAsync(symbol, state).ConfigureAwait(false);
+            await DirectSymbolSearchAsync(symbol, group, state).ConfigureAwait(false);
 
             // Now, for symbols that could involve inheritance, look for references to the same named entity, and
             // see if it's a reference to a symbol that shares an inheritance relationship with that symbol.
             await InheritanceSymbolSearchAsync(symbol, state).ConfigureAwait(false);
         }
 
-        async ValueTask DirectSymbolSearchAsync(ISymbol symbol, FindReferencesDocumentState state)
+        async ValueTask DirectSymbolSearchAsync(ISymbol symbol, SymbolGroup group, FindReferencesDocumentState state)
         {
             await ProducerConsumer<FinderLocation>.RunAsync(
                 ProducerConsumerOptions.SingleReaderWriterOptions,
                 static (callback, args, cancellationToken) =>
                 {
-                    var (@this, symbol, state) = args;
+                    var (@this, symbol, group, state) = args;
 
                     // We don't bother calling into the finders in parallel as there's only ever one that applies for a
                     // particular symbol kind.  All the rest bail out immediately after a quick type-check.  So there's
@@ -134,28 +131,22 @@ await ProducerConsumer<FinderLocation>.RunAsync(
                 },
                 consumeItems: static async (values, args, cancellationToken) =>
                 {
-                    var (@this, symbol, state) = args;
-                    var converted = await ConvertLocationsAndReportGroupsAsync(@this, values, symbol, cancellationToken).ConfigureAwait(false);
+                    var (@this, symbol, group, state) = args;
+                    var converted = await ConvertLocationsAsync(@this, values, symbol, group, cancellationToken).ConfigureAwait(false);
                     await @this._progress.OnReferencesFoundAsync(converted, cancellationToken).ConfigureAwait(false);
                 },
-                args: (@this: this, symbol, state),
+                args: (@this: this, symbol, group, state),
                 cancellationToken).ConfigureAwait(false);
         }
 
-        static async Task<ImmutableArray<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)>> ConvertLocationsAndReportGroupsAsync(
-            FindReferencesSearchEngine @this, IAsyncEnumerable<FinderLocation> locations, ISymbol symbol, CancellationToken cancellationToken)
+        static async Task<ImmutableArray<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)>> ConvertLocationsAsync(
+            FindReferencesSearchEngine @this, IAsyncEnumerable<FinderLocation> locations, ISymbol symbol, SymbolGroup group, CancellationToken cancellationToken)
         {
-            SymbolGroup? group = null;
-
             using var _ = ArrayBuilder<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)>.GetInstance(out var result);
 
             // Transform the individual finder-location objects to "group/symbol/location" tuples.
             await foreach (var location in locations)
-            {
-                // The first time we see the location for a symbol, report its group.
-                group ??= await @this.ReportGroupAsync(symbol, cancellationToken).ConfigureAwait(false);
                 result.Add((group, symbol, location.Location));
-            }
 
             return result.ToImmutableAndClear();
         }

From a306ddae2416fb20954b2648c34de021ce517031 Mon Sep 17 00:00:00 2001
From: Cyrus Najmabadi <cyrus.najmabadi@gmail.com>
Date: Thu, 11 Jul 2024 11:03:10 -0700
Subject: [PATCH 2/9] In progress

---
 .../FindReferencesSearchEngine.cs             | 34 +++++++++++------
 ...sSearchEngine_FindReferencesInDocuments.cs | 37 +++++++++++++------
 2 files changed, 47 insertions(+), 24 deletions(-)

diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
index d47e85b0aeaea..9f1c8a96a5561 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
@@ -42,7 +42,11 @@ internal partial class FindReferencesSearchEngine
     /// them in the case of linked files across projects.  This allows references to be found to any of the unified
     /// symbols, while the user only gets a single reported group back that corresponds to that entire set.
     /// </summary>
-    private readonly ConcurrentDictionary<ISymbol, SymbolGroup> _symbolToGroup = new(MetadataUnifyingEquivalenceComparer.Instance);
+    /// <remarks>
+    /// This is a normal dictionary that is not locked.  It is only ever read and written to serially from within the
+    /// high level project-walking code in this class.
+    /// </remarks>
+    private readonly Dictionary<ISymbol, SymbolGroup> _symbolToGroup_mustBeAccessedSerially = new(MetadataUnifyingEquivalenceComparer.Instance);
 
     public FindReferencesSearchEngine(
         Solution solution,
@@ -113,7 +117,9 @@ private async Task PerformSearchAsync(
 
         // Report the initial set of symbols to the caller.
         var allSymbols = symbolSet.GetAllSymbols();
-        await ReportGroupsAsync(allSymbols, cancellationToken).ConfigureAwait(false);
+
+        // Safe to call as we're in the entry-point method, and nothing is running concurrently with this call.
+        await ReportGroupsSeriallyAsync(allSymbols, cancellationToken).ConfigureAwait(false);
 
         // Determine the set of projects we actually have to walk to find results in.  If the caller provided a
         // set of documents to search, we only bother with those.
@@ -123,13 +129,15 @@ private async Task PerformSearchAsync(
 
         // Pull off and start searching each project as soon as we can once we've done the inheritance cascade into it.
         await RoslynParallel.ForEachAsync(
-            GetProjectsAndSymbolsToSearchAsync(symbolSet, projectsToSearch, cancellationToken),
+            // ForEachAsync will serially pull on the IAsyncEnumerable returned here, kicking off the processing to then
+            // happen in parallel.
+            GetProjectsAndSymbolsToSearchSeriallyAsync(symbolSet, projectsToSearch, cancellationToken),
             GetParallelOptions(cancellationToken),
             async (tuple, cancellationToken) => await ProcessProjectAsync(
                 tuple.project, tuple.allSymbols, onReferenceFound, cancellationToken).ConfigureAwait(false)).ConfigureAwait(false);
     }
 
-    private async IAsyncEnumerable<(Project project, ImmutableArray<(ISymbol symbol, SymbolGroup group)> allSymbols)> GetProjectsAndSymbolsToSearchAsync(
+    private async IAsyncEnumerable<(Project project, ImmutableArray<(ISymbol symbol, SymbolGroup group)> allSymbols)> GetProjectsAndSymbolsToSearchSeriallyAsync(
         SymbolSet symbolSet,
         ImmutableArray<Project> projectsToSearch,
         [EnumeratorCancellation] CancellationToken cancellationToken)
@@ -152,8 +160,9 @@ await RoslynParallel.ForEachAsync(
             // below.
             await symbolSet.InheritanceCascadeAsync(currentProject, cancellationToken).ConfigureAwait(false);
 
-            // Report any new symbols we've cascaded to to our caller.
-            var allSymbolsAndGroups = await ReportGroupsAsync(symbolSet.GetAllSymbols(), cancellationToken).ConfigureAwait(false);
+            // Report any new symbols we've cascaded to to our caller.  This is safe to call here as we're abiding by
+            // the serial requirements of ReportGroupsSeriallyAsync
+            var allSymbolsAndGroups = await ReportGroupsSeriallyAsync(symbolSet.GetAllSymbols(), cancellationToken).ConfigureAwait(false);
 
             yield return (currentProject, allSymbolsAndGroups);
         }
@@ -164,18 +173,19 @@ await RoslynParallel.ForEachAsync(
     /// them once per symbol group, but we may have to notify about new symbols each time we expand our symbol set
     /// when we walk into a new project.
     /// </summary>
-    private async Task<ImmutableArray<(ISymbol symbol, SymbolGroup group)>> ReportGroupsAsync(
+    private async Task<ImmutableArray<(ISymbol symbol, SymbolGroup group)>> ReportGroupsSeriallyAsync(
         ImmutableArray<ISymbol> symbols, CancellationToken cancellationToken)
     {
         var result = new FixedSizeArrayBuilder<(ISymbol symbol, SymbolGroup group)>(symbols.Length);
 
+        // Safe to call this as we're only being called from within a serial context ourselves.
         foreach (var symbol in symbols)
-            result.Add((symbol, await ReportGroupAsync(symbol, cancellationToken).ConfigureAwait(false)));
+            result.Add((symbol, await ReportGroupSeriallyAsync(symbol, cancellationToken).ConfigureAwait(false)));
 
         return result.MoveToImmutable();
     }
 
-    private async ValueTask<SymbolGroup> ReportGroupAsync(ISymbol symbol, CancellationToken cancellationToken)
+    private async ValueTask<SymbolGroup> ReportGroupSeriallyAsync(ISymbol symbol, CancellationToken cancellationToken)
     {
         // See if this is the first time we're running across this symbol.  Note: no locks are needed
         // here between checking and then adding because this is only ever called serially from within
@@ -183,7 +193,7 @@ private async ValueTask<SymbolGroup> ReportGroupAsync(ISymbol symbol, Cancellati
         // symbols will happen later in ProcessDocumentAsync.  However, those reads will only happen
         // after the dependent symbol values were written in, so it will be safe to blindly read them
         // out.
-        if (!_symbolToGroup.TryGetValue(symbol, out var group))
+        if (!_symbolToGroup_mustBeAccessedSerially.TryGetValue(symbol, out var group))
         {
             var linkedSymbols = await SymbolFinder.FindLinkedSymbolsAsync(symbol, _solution, cancellationToken).ConfigureAwait(false);
             Contract.ThrowIfFalse(linkedSymbols.Contains(symbol), "Linked symbols did not contain the very symbol we started with.");
@@ -192,11 +202,11 @@ private async ValueTask<SymbolGroup> ReportGroupAsync(ISymbol symbol, Cancellati
             Contract.ThrowIfFalse(group.Symbols.Contains(symbol), "Symbol group did not contain the very symbol we started with.");
 
             foreach (var groupSymbol in group.Symbols)
-                _symbolToGroup.TryAdd(groupSymbol, group);
+                _symbolToGroup_mustBeAccessedSerially.TryAdd(groupSymbol, group);
 
             // Since "symbol" was in group.Symbols, and we just added links from all of group.Symbols to that group, then "symbol" 
             // better now be in _symbolToGroup.
-            Contract.ThrowIfFalse(_symbolToGroup.ContainsKey(symbol));
+            Contract.ThrowIfFalse(_symbolToGroup_mustBeAccessedSerially.ContainsKey(symbol));
 
             await _progress.OnDefinitionFoundAsync(group, cancellationToken).ConfigureAwait(false);
         }
diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
index bf11d643ec6c6..b29878f67bb49 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
@@ -41,7 +41,9 @@ public async Task FindReferencesInDocumentsAsync(
         // Create and report the initial set of symbols to search for.  This includes linked and cascaded symbols. It does
         // not walk up/down the inheritance hierarchy.
         var symbolSet = await SymbolSet.DetermineInitialSearchSymbolsAsync(this, unifiedSymbols, cancellationToken).ConfigureAwait(false);
-        var allSymbolsAndGroups = await ReportGroupsAsync([.. symbolSet], cancellationToken).ConfigureAwait(false);
+
+        // Safe to call as we're in the entry-point method, and nothing is running concurrently with this call.
+        var allSymbolsAndGroups = await ReportGroupsSeriallyAsync([.. symbolSet], cancellationToken).ConfigureAwait(false);
 
         // Process projects in dependency graph order so that any compilations built by one are available for later
         // projects. We only have to examine the projects containing the documents requested though.
@@ -51,13 +53,17 @@ public async Task FindReferencesInDocumentsAsync(
         foreach (var projectId in dependencyGraph.GetTopologicallySortedProjects(cancellationToken))
         {
             var currentProject = _solution.GetRequiredProject(projectId);
-            if (projectsToSearch.Contains(currentProject))
-                await PerformSearchInProjectAsync(allSymbolsAndGroups, currentProject).ConfigureAwait(false);
+            if (!projectsToSearch.Contains(currentProject))
+                continue;
+
+            // Safe to call as we're in the entry-point method, and it's only serially looping over the projects when
+            // calling into this.
+            await PerformSearchInProjectSeriallyAsync(allSymbolsAndGroups, currentProject).ConfigureAwait(false);
         }
 
         return;
 
-        async ValueTask PerformSearchInProjectAsync(ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols, Project project)
+        async ValueTask PerformSearchInProjectSeriallyAsync(ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols, Project project)
         {
             using var _ = PooledDictionary<ISymbol, PooledHashSet<string>>.GetInstance(out var symbolToGlobalAliases);
             try
@@ -68,8 +74,11 @@ async ValueTask PerformSearchInProjectAsync(ImmutableArray<(ISymbol symbol, Symb
 
                 foreach (var document in documents)
                 {
-                    if (document.Project == project)
-                        await PerformSearchInDocumentAsync(symbols, document, symbolToGlobalAliases).ConfigureAwait(false);
+                    if (document.Project != project)
+                        continue;
+
+                    // Safe to call as we're only in a serial context ourselves.
+                    await PerformSearchInDocumentSeriallyAsync(symbols, document, symbolToGlobalAliases).ConfigureAwait(false);
                 }
             }
             finally
@@ -78,7 +87,7 @@ async ValueTask PerformSearchInProjectAsync(ImmutableArray<(ISymbol symbol, Symb
             }
         }
 
-        async ValueTask PerformSearchInDocumentAsync(
+        async ValueTask PerformSearchInDocumentSeriallyAsync(
             ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols,
             Document document,
             PooledDictionary<ISymbol, PooledHashSet<string>> symbolToGlobalAliases)
@@ -94,18 +103,21 @@ async ValueTask PerformSearchInDocumentAsync(
                 var state = new FindReferencesDocumentState(
                     cache, TryGet(symbolToGlobalAliases, symbol));
 
-                await PerformSearchInDocumentWorkerAsync(symbol, group, state).ConfigureAwait(false);
+                // Safe to call as we're only in a serial context ourselves.
+                await PerformSearchInDocumentSeriallyWorkerAsync(symbol, group, state).ConfigureAwait(false);
             }
         }
 
-        async ValueTask PerformSearchInDocumentWorkerAsync(ISymbol symbol, SymbolGroup group, FindReferencesDocumentState state)
+        async ValueTask PerformSearchInDocumentSeriallyWorkerAsync(ISymbol symbol, SymbolGroup group, FindReferencesDocumentState state)
         {
             // Always perform a normal search, looking for direct references to exactly that symbol.
             await DirectSymbolSearchAsync(symbol, group, state).ConfigureAwait(false);
 
             // Now, for symbols that could involve inheritance, look for references to the same named entity, and
             // see if it's a reference to a symbol that shares an inheritance relationship with that symbol.
-            await InheritanceSymbolSearchAsync(symbol, state).ConfigureAwait(false);
+            //
+            // Safe to call as we're only in a serial context ourselves.
+            await InheritanceSymbolSearchSeriallyAsync(symbol, state).ConfigureAwait(false);
         }
 
         async ValueTask DirectSymbolSearchAsync(ISymbol symbol, SymbolGroup group, FindReferencesDocumentState state)
@@ -151,7 +163,7 @@ await ProducerConsumer<FinderLocation>.RunAsync(
             return result.ToImmutableAndClear();
         }
 
-        async ValueTask InheritanceSymbolSearchAsync(ISymbol symbol, FindReferencesDocumentState state)
+        async ValueTask InheritanceSymbolSearchSeriallyAsync(ISymbol symbol, FindReferencesDocumentState state)
         {
             if (InvolvesInheritance(symbol))
             {
@@ -166,7 +178,8 @@ async ValueTask InheritanceSymbolSearchAsync(ISymbol symbol, FindReferencesDocum
                     if (matched)
                     {
                         // Ensure we report this new symbol/group in case it's the first time we're seeing it.
-                        var candidateGroup = await ReportGroupAsync(candidate, cancellationToken).ConfigureAwait(false);
+                        // Safe to call this as we're only being called from within a serial context ourselves.
+                        var candidateGroup = await ReportGroupSeriallyAsync(candidate, cancellationToken).ConfigureAwait(false);
 
                         var location = AbstractReferenceFinder.CreateReferenceLocation(state, token, candidateReason, cancellationToken);
                         await _progress.OnReferencesFoundAsync([(candidateGroup, candidate, location)], cancellationToken).ConfigureAwait(false);

From d4037126c9afc898e492c2b69164b264317a0a1f Mon Sep 17 00:00:00 2001
From: Cyrus Najmabadi <cyrus.najmabadi@gmail.com>
Date: Thu, 11 Jul 2024 11:09:01 -0700
Subject: [PATCH 3/9] Make into local data

---
 .../FindReferencesSearchEngine.cs             | 75 ++++++++-----------
 ...sSearchEngine_FindReferencesInDocuments.cs | 37 ++++++---
 2 files changed, 59 insertions(+), 53 deletions(-)

diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
index 9f1c8a96a5561..7e32bb244f855 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
@@ -21,7 +21,12 @@ namespace Microsoft.CodeAnalysis.FindSymbols;
 
 using Reference = (SymbolGroup group, ISymbol symbol, ReferenceLocation location);
 
-internal partial class FindReferencesSearchEngine
+internal partial class FindReferencesSearchEngine(
+    Solution solution,
+    IImmutableSet<Document>? documents,
+    ImmutableArray<IReferenceFinder> finders,
+    IStreamingFindReferencesProgress progress,
+    FindReferencesSearchOptions options)
 {
     /// <summary>
     /// Scheduler we use when we're doing operations in the BG and we want to rate limit them to not saturate the threadpool.
@@ -30,39 +35,12 @@ internal partial class FindReferencesSearchEngine
 
     private static readonly ObjectPool<Dictionary<ISymbol, SymbolGroup>> s_symbolToGroupPool = new(() => new(MetadataUnifyingEquivalenceComparer.Instance));
 
-    private readonly Solution _solution;
-    private readonly IImmutableSet<Document>? _documents;
-    private readonly ImmutableArray<IReferenceFinder> _finders;
-    private readonly IStreamingProgressTracker _progressTracker;
-    private readonly IStreamingFindReferencesProgress _progress;
-    private readonly FindReferencesSearchOptions _options;
-
-    /// <summary>
-    /// Mapping from symbols (unified across metadata/retargeting) and the set of symbols that was produced for 
-    /// them in the case of linked files across projects.  This allows references to be found to any of the unified
-    /// symbols, while the user only gets a single reported group back that corresponds to that entire set.
-    /// </summary>
-    /// <remarks>
-    /// This is a normal dictionary that is not locked.  It is only ever read and written to serially from within the
-    /// high level project-walking code in this class.
-    /// </remarks>
-    private readonly Dictionary<ISymbol, SymbolGroup> _symbolToGroup_mustBeAccessedSerially = new(MetadataUnifyingEquivalenceComparer.Instance);
-
-    public FindReferencesSearchEngine(
-        Solution solution,
-        IImmutableSet<Document>? documents,
-        ImmutableArray<IReferenceFinder> finders,
-        IStreamingFindReferencesProgress progress,
-        FindReferencesSearchOptions options)
-    {
-        _documents = documents;
-        _solution = solution;
-        _finders = finders;
-        _progress = progress;
-        _options = options;
-
-        _progressTracker = progress.ProgressTracker;
-    }
+    private readonly Solution _solution = solution;
+    private readonly IImmutableSet<Document>? _documents = documents;
+    private readonly ImmutableArray<IReferenceFinder> _finders = finders;
+    private readonly IStreamingProgressTracker _progressTracker = progress.ProgressTracker;
+    private readonly IStreamingFindReferencesProgress _progress = progress;
+    private readonly FindReferencesSearchOptions _options = options;
 
     /// <summary>
     /// Options to control the parallelism of the search.   If we're in <see
@@ -104,6 +82,14 @@ await ProducerConsumer<Reference>.RunAsync(
     private async Task PerformSearchAsync(
         ImmutableArray<ISymbol> symbols, Action<Reference> onReferenceFound, CancellationToken cancellationToken)
     {
+        // Mapping from symbols (unified across metadata/retargeting) and the set of symbols that was produced for 
+        // them in the case of linked files across projects.  This allows references to be found to any of the unified
+        // symbols, while the user only gets a single reported group back that corresponds to that entire set.
+        //
+        // This is a normal dictionary that is not locked.  It is only ever read and written to serially from within the
+        // high level project-walking code in this method.
+        var symbolToGroup = new Dictionary<ISymbol, SymbolGroup>(MetadataUnifyingEquivalenceComparer.Instance);
+
         var unifiedSymbols = new MetadataUnifyingSymbolHashSet();
         unifiedSymbols.AddRange(symbols);
 
@@ -119,7 +105,7 @@ private async Task PerformSearchAsync(
         var allSymbols = symbolSet.GetAllSymbols();
 
         // Safe to call as we're in the entry-point method, and nothing is running concurrently with this call.
-        await ReportGroupsSeriallyAsync(allSymbols, cancellationToken).ConfigureAwait(false);
+        await ReportGroupsSeriallyAsync(allSymbols, symbolToGroup, cancellationToken).ConfigureAwait(false);
 
         // Determine the set of projects we actually have to walk to find results in.  If the caller provided a
         // set of documents to search, we only bother with those.
@@ -131,7 +117,7 @@ private async Task PerformSearchAsync(
         await RoslynParallel.ForEachAsync(
             // ForEachAsync will serially pull on the IAsyncEnumerable returned here, kicking off the processing to then
             // happen in parallel.
-            GetProjectsAndSymbolsToSearchSeriallyAsync(symbolSet, projectsToSearch, cancellationToken),
+            GetProjectsAndSymbolsToSearchSeriallyAsync(symbolSet, projectsToSearch, symbolToGroup, cancellationToken),
             GetParallelOptions(cancellationToken),
             async (tuple, cancellationToken) => await ProcessProjectAsync(
                 tuple.project, tuple.allSymbols, onReferenceFound, cancellationToken).ConfigureAwait(false)).ConfigureAwait(false);
@@ -140,6 +126,7 @@ await RoslynParallel.ForEachAsync(
     private async IAsyncEnumerable<(Project project, ImmutableArray<(ISymbol symbol, SymbolGroup group)> allSymbols)> GetProjectsAndSymbolsToSearchSeriallyAsync(
         SymbolSet symbolSet,
         ImmutableArray<Project> projectsToSearch,
+        Dictionary<ISymbol, SymbolGroup> symbolToGroup,
         [EnumeratorCancellation] CancellationToken cancellationToken)
     {
         // We need to process projects in order when updating our symbol set.  Say we have three projects (A, B
@@ -162,7 +149,8 @@ await RoslynParallel.ForEachAsync(
 
             // Report any new symbols we've cascaded to to our caller.  This is safe to call here as we're abiding by
             // the serial requirements of ReportGroupsSeriallyAsync
-            var allSymbolsAndGroups = await ReportGroupsSeriallyAsync(symbolSet.GetAllSymbols(), cancellationToken).ConfigureAwait(false);
+            var allSymbolsAndGroups = await ReportGroupsSeriallyAsync(
+                symbolSet.GetAllSymbols(), symbolToGroup, cancellationToken).ConfigureAwait(false);
 
             yield return (currentProject, allSymbolsAndGroups);
         }
@@ -174,18 +162,19 @@ await RoslynParallel.ForEachAsync(
     /// when we walk into a new project.
     /// </summary>
     private async Task<ImmutableArray<(ISymbol symbol, SymbolGroup group)>> ReportGroupsSeriallyAsync(
-        ImmutableArray<ISymbol> symbols, CancellationToken cancellationToken)
+        ImmutableArray<ISymbol> symbols, Dictionary<ISymbol, SymbolGroup> symbolToGroup, CancellationToken cancellationToken)
     {
         var result = new FixedSizeArrayBuilder<(ISymbol symbol, SymbolGroup group)>(symbols.Length);
 
         // Safe to call this as we're only being called from within a serial context ourselves.
         foreach (var symbol in symbols)
-            result.Add((symbol, await ReportGroupSeriallyAsync(symbol, cancellationToken).ConfigureAwait(false)));
+            result.Add((symbol, await ReportGroupSeriallyAsync(symbol, symbolToGroup, cancellationToken).ConfigureAwait(false)));
 
         return result.MoveToImmutable();
     }
 
-    private async ValueTask<SymbolGroup> ReportGroupSeriallyAsync(ISymbol symbol, CancellationToken cancellationToken)
+    private async ValueTask<SymbolGroup> ReportGroupSeriallyAsync(
+        ISymbol symbol, Dictionary<ISymbol, SymbolGroup> symbolToGroup, CancellationToken cancellationToken)
     {
         // See if this is the first time we're running across this symbol.  Note: no locks are needed
         // here between checking and then adding because this is only ever called serially from within
@@ -193,7 +182,7 @@ private async ValueTask<SymbolGroup> ReportGroupSeriallyAsync(ISymbol symbol, Ca
         // symbols will happen later in ProcessDocumentAsync.  However, those reads will only happen
         // after the dependent symbol values were written in, so it will be safe to blindly read them
         // out.
-        if (!_symbolToGroup_mustBeAccessedSerially.TryGetValue(symbol, out var group))
+        if (!symbolToGroup.TryGetValue(symbol, out var group))
         {
             var linkedSymbols = await SymbolFinder.FindLinkedSymbolsAsync(symbol, _solution, cancellationToken).ConfigureAwait(false);
             Contract.ThrowIfFalse(linkedSymbols.Contains(symbol), "Linked symbols did not contain the very symbol we started with.");
@@ -202,11 +191,11 @@ private async ValueTask<SymbolGroup> ReportGroupSeriallyAsync(ISymbol symbol, Ca
             Contract.ThrowIfFalse(group.Symbols.Contains(symbol), "Symbol group did not contain the very symbol we started with.");
 
             foreach (var groupSymbol in group.Symbols)
-                _symbolToGroup_mustBeAccessedSerially.TryAdd(groupSymbol, group);
+                symbolToGroup.TryAdd(groupSymbol, group);
 
             // Since "symbol" was in group.Symbols, and we just added links from all of group.Symbols to that group, then "symbol" 
             // better now be in _symbolToGroup.
-            Contract.ThrowIfFalse(_symbolToGroup_mustBeAccessedSerially.ContainsKey(symbol));
+            Contract.ThrowIfFalse(symbolToGroup.ContainsKey(symbol));
 
             await _progress.OnDefinitionFoundAsync(group, cancellationToken).ConfigureAwait(false);
         }
diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
index b29878f67bb49..55c078406f352 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
@@ -26,6 +26,14 @@ public async Task FindReferencesInDocumentsAsync(
         // the starting symbol.
         Debug.Assert(_options.UnidirectionalHierarchyCascade);
 
+        // Mapping from symbols (unified across metadata/retargeting) and the set of symbols that was produced for 
+        // them in the case of linked files across projects.  This allows references to be found to any of the unified
+        // symbols, while the user only gets a single reported group back that corresponds to that entire set.
+        //
+        // This is a normal dictionary that is not locked.  It is only ever read and written to serially from within the
+        // high level project-walking code in this method.
+        var symbolToGroup = new Dictionary<ISymbol, SymbolGroup>(MetadataUnifyingEquivalenceComparer.Instance);
+
         var unifiedSymbols = new MetadataUnifyingSymbolHashSet
         {
             originalSymbol
@@ -43,7 +51,8 @@ public async Task FindReferencesInDocumentsAsync(
         var symbolSet = await SymbolSet.DetermineInitialSearchSymbolsAsync(this, unifiedSymbols, cancellationToken).ConfigureAwait(false);
 
         // Safe to call as we're in the entry-point method, and nothing is running concurrently with this call.
-        var allSymbolsAndGroups = await ReportGroupsSeriallyAsync([.. symbolSet], cancellationToken).ConfigureAwait(false);
+        var allSymbolsAndGroups = await ReportGroupsSeriallyAsync(
+            [.. symbolSet], symbolToGroup, cancellationToken).ConfigureAwait(false);
 
         // Process projects in dependency graph order so that any compilations built by one are available for later
         // projects. We only have to examine the projects containing the documents requested though.
@@ -58,12 +67,14 @@ public async Task FindReferencesInDocumentsAsync(
 
             // Safe to call as we're in the entry-point method, and it's only serially looping over the projects when
             // calling into this.
-            await PerformSearchInProjectSeriallyAsync(allSymbolsAndGroups, currentProject).ConfigureAwait(false);
+            await PerformSearchInProjectSeriallyAsync(
+                currentProject, allSymbolsAndGroups, symbolToGroup).ConfigureAwait(false);
         }
 
         return;
 
-        async ValueTask PerformSearchInProjectSeriallyAsync(ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols, Project project)
+        async ValueTask PerformSearchInProjectSeriallyAsync(
+            Project project, ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols, Dictionary<ISymbol, SymbolGroup> symbolToGroup)
         {
             using var _ = PooledDictionary<ISymbol, PooledHashSet<string>>.GetInstance(out var symbolToGlobalAliases);
             try
@@ -78,7 +89,8 @@ async ValueTask PerformSearchInProjectSeriallyAsync(ImmutableArray<(ISymbol symb
                         continue;
 
                     // Safe to call as we're only in a serial context ourselves.
-                    await PerformSearchInDocumentSeriallyAsync(symbols, document, symbolToGlobalAliases).ConfigureAwait(false);
+                    await PerformSearchInDocumentSeriallyAsync(
+                        document, symbols, symbolToGroup, symbolToGlobalAliases).ConfigureAwait(false);
                 }
             }
             finally
@@ -88,8 +100,9 @@ async ValueTask PerformSearchInProjectSeriallyAsync(ImmutableArray<(ISymbol symb
         }
 
         async ValueTask PerformSearchInDocumentSeriallyAsync(
-            ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols,
             Document document,
+            ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols,
+            Dictionary<ISymbol, SymbolGroup> symbolToGroup,
             PooledDictionary<ISymbol, PooledHashSet<string>> symbolToGlobalAliases)
         {
             // We're doing to do all of our processing of this document at once.  This will necessitate all the
@@ -104,11 +117,13 @@ async ValueTask PerformSearchInDocumentSeriallyAsync(
                     cache, TryGet(symbolToGlobalAliases, symbol));
 
                 // Safe to call as we're only in a serial context ourselves.
-                await PerformSearchInDocumentSeriallyWorkerAsync(symbol, group, state).ConfigureAwait(false);
+                await PerformSearchInDocumentSeriallyWorkerAsync(
+                    symbol, group, symbolToGroup, state).ConfigureAwait(false);
             }
         }
 
-        async ValueTask PerformSearchInDocumentSeriallyWorkerAsync(ISymbol symbol, SymbolGroup group, FindReferencesDocumentState state)
+        async ValueTask PerformSearchInDocumentSeriallyWorkerAsync(
+            ISymbol symbol, SymbolGroup group, Dictionary<ISymbol, SymbolGroup> symbolToGroup, FindReferencesDocumentState state)
         {
             // Always perform a normal search, looking for direct references to exactly that symbol.
             await DirectSymbolSearchAsync(symbol, group, state).ConfigureAwait(false);
@@ -117,7 +132,7 @@ async ValueTask PerformSearchInDocumentSeriallyWorkerAsync(ISymbol symbol, Symbo
             // see if it's a reference to a symbol that shares an inheritance relationship with that symbol.
             //
             // Safe to call as we're only in a serial context ourselves.
-            await InheritanceSymbolSearchSeriallyAsync(symbol, state).ConfigureAwait(false);
+            await InheritanceSymbolSearchSeriallyAsync(symbol, symbolToGroup, state).ConfigureAwait(false);
         }
 
         async ValueTask DirectSymbolSearchAsync(ISymbol symbol, SymbolGroup group, FindReferencesDocumentState state)
@@ -163,7 +178,8 @@ await ProducerConsumer<FinderLocation>.RunAsync(
             return result.ToImmutableAndClear();
         }
 
-        async ValueTask InheritanceSymbolSearchSeriallyAsync(ISymbol symbol, FindReferencesDocumentState state)
+        async ValueTask InheritanceSymbolSearchSeriallyAsync(
+            ISymbol symbol, Dictionary<ISymbol, SymbolGroup> symbolToGroup, FindReferencesDocumentState state)
         {
             if (InvolvesInheritance(symbol))
             {
@@ -179,7 +195,8 @@ async ValueTask InheritanceSymbolSearchSeriallyAsync(ISymbol symbol, FindReferen
                     {
                         // Ensure we report this new symbol/group in case it's the first time we're seeing it.
                         // Safe to call this as we're only being called from within a serial context ourselves.
-                        var candidateGroup = await ReportGroupSeriallyAsync(candidate, cancellationToken).ConfigureAwait(false);
+                        var candidateGroup = await ReportGroupSeriallyAsync(
+                            candidate, symbolToGroup, cancellationToken).ConfigureAwait(false);
 
                         var location = AbstractReferenceFinder.CreateReferenceLocation(state, token, candidateReason, cancellationToken);
                         await _progress.OnReferencesFoundAsync([(candidateGroup, candidate, location)], cancellationToken).ConfigureAwait(false);

From f7259bbde6ab73e55e898e5ea14e21deec290578 Mon Sep 17 00:00:00 2001
From: Cyrus Najmabadi <cyrus.najmabadi@gmail.com>
Date: Thu, 11 Jul 2024 11:12:16 -0700
Subject: [PATCH 4/9] Capture

---
 ...sSearchEngine_FindReferencesInDocuments.cs | 20 ++++++++-----------
 1 file changed, 8 insertions(+), 12 deletions(-)

diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
index 55c078406f352..8fb5ff68a9ef6 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
@@ -67,14 +67,13 @@ public async Task FindReferencesInDocumentsAsync(
 
             // Safe to call as we're in the entry-point method, and it's only serially looping over the projects when
             // calling into this.
-            await PerformSearchInProjectSeriallyAsync(
-                currentProject, allSymbolsAndGroups, symbolToGroup).ConfigureAwait(false);
+            await PerformSearchInProjectSeriallyAsync(allSymbolsAndGroups, currentProject).ConfigureAwait(false);
         }
 
         return;
 
         async ValueTask PerformSearchInProjectSeriallyAsync(
-            Project project, ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols, Dictionary<ISymbol, SymbolGroup> symbolToGroup)
+            ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols, Project project)
         {
             using var _ = PooledDictionary<ISymbol, PooledHashSet<string>>.GetInstance(out var symbolToGlobalAliases);
             try
@@ -89,8 +88,7 @@ async ValueTask PerformSearchInProjectSeriallyAsync(
                         continue;
 
                     // Safe to call as we're only in a serial context ourselves.
-                    await PerformSearchInDocumentSeriallyAsync(
-                        document, symbols, symbolToGroup, symbolToGlobalAliases).ConfigureAwait(false);
+                    await PerformSearchInDocumentSeriallyAsync(symbols, document, symbolToGlobalAliases).ConfigureAwait(false);
                 }
             }
             finally
@@ -100,9 +98,8 @@ await PerformSearchInDocumentSeriallyAsync(
         }
 
         async ValueTask PerformSearchInDocumentSeriallyAsync(
-            Document document,
             ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols,
-            Dictionary<ISymbol, SymbolGroup> symbolToGroup,
+            Document document,
             PooledDictionary<ISymbol, PooledHashSet<string>> symbolToGlobalAliases)
         {
             // We're doing to do all of our processing of this document at once.  This will necessitate all the
@@ -117,13 +114,12 @@ async ValueTask PerformSearchInDocumentSeriallyAsync(
                     cache, TryGet(symbolToGlobalAliases, symbol));
 
                 // Safe to call as we're only in a serial context ourselves.
-                await PerformSearchInDocumentSeriallyWorkerAsync(
-                    symbol, group, symbolToGroup, state).ConfigureAwait(false);
+                await PerformSearchInDocumentSeriallyWorkerAsync(symbol, group, state).ConfigureAwait(false);
             }
         }
 
         async ValueTask PerformSearchInDocumentSeriallyWorkerAsync(
-            ISymbol symbol, SymbolGroup group, Dictionary<ISymbol, SymbolGroup> symbolToGroup, FindReferencesDocumentState state)
+            ISymbol symbol, SymbolGroup group, FindReferencesDocumentState state)
         {
             // Always perform a normal search, looking for direct references to exactly that symbol.
             await DirectSymbolSearchAsync(symbol, group, state).ConfigureAwait(false);
@@ -132,7 +128,7 @@ async ValueTask PerformSearchInDocumentSeriallyWorkerAsync(
             // see if it's a reference to a symbol that shares an inheritance relationship with that symbol.
             //
             // Safe to call as we're only in a serial context ourselves.
-            await InheritanceSymbolSearchSeriallyAsync(symbol, symbolToGroup, state).ConfigureAwait(false);
+            await InheritanceSymbolSearchSeriallyAsync(symbol, state).ConfigureAwait(false);
         }
 
         async ValueTask DirectSymbolSearchAsync(ISymbol symbol, SymbolGroup group, FindReferencesDocumentState state)
@@ -179,7 +175,7 @@ await ProducerConsumer<FinderLocation>.RunAsync(
         }
 
         async ValueTask InheritanceSymbolSearchSeriallyAsync(
-            ISymbol symbol, Dictionary<ISymbol, SymbolGroup> symbolToGroup, FindReferencesDocumentState state)
+            ISymbol symbol, FindReferencesDocumentState state)
         {
             if (InvolvesInheritance(symbol))
             {

From 682583c33ae58223308fe42ae13a3abf0897ce97 Mon Sep 17 00:00:00 2001
From: Cyrus Najmabadi <cyrus.najmabadi@gmail.com>
Date: Thu, 11 Jul 2024 11:16:06 -0700
Subject: [PATCH 5/9] Update
 src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs

---
 .../FindReferencesSearchEngine_FindReferencesInDocuments.cs    | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
index 8fb5ff68a9ef6..559ced8efa672 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
@@ -174,8 +174,7 @@ await ProducerConsumer<FinderLocation>.RunAsync(
             return result.ToImmutableAndClear();
         }
 
-        async ValueTask InheritanceSymbolSearchSeriallyAsync(
-            ISymbol symbol, FindReferencesDocumentState state)
+        async ValueTask InheritanceSymbolSearchSeriallyAsync(ISymbol symbol, FindReferencesDocumentState state)
         {
             if (InvolvesInheritance(symbol))
             {

From ac62bcdf40e5c06b2a89be2d16d8be91e3446fb6 Mon Sep 17 00:00:00 2001
From: Cyrus Najmabadi <cyrus.najmabadi@gmail.com>
Date: Thu, 11 Jul 2024 11:18:14 -0700
Subject: [PATCH 6/9] Update
 src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs

---
 .../FindReferencesSearchEngine_FindReferencesInDocuments.cs    | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
index 559ced8efa672..ccbb11e3cf36d 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
@@ -72,8 +72,7 @@ public async Task FindReferencesInDocumentsAsync(
 
         return;
 
-        async ValueTask PerformSearchInProjectSeriallyAsync(
-            ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols, Project project)
+        async ValueTask PerformSearchInProjectSeriallyAsync(ImmutableArray<(ISymbol symbol, SymbolGroup group)> symbols, Project project)
         {
             using var _ = PooledDictionary<ISymbol, PooledHashSet<string>>.GetInstance(out var symbolToGlobalAliases);
             try

From 774670af980d115119b3dd9048f13cbef2a4bc52 Mon Sep 17 00:00:00 2001
From: Cyrus Najmabadi <cyrus.najmabadi@gmail.com>
Date: Thu, 11 Jul 2024 12:31:19 -0700
Subject: [PATCH 7/9] Renames

---
 .../FindReferencesSearchEngine.cs             | 38 +++++++++----------
 1 file changed, 19 insertions(+), 19 deletions(-)

diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
index 7e32bb244f855..f2247e264638b 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
@@ -217,7 +217,7 @@ private async ValueTask ProcessProjectAsync(
         Project project, ImmutableArray<(ISymbol symbol, SymbolGroup group)> allSymbols, Action<Reference> onReferenceFound, CancellationToken cancellationToken)
     {
         using var _1 = PooledDictionary<ISymbol, PooledHashSet<string>>.GetInstance(out var symbolToGlobalAliases);
-        using var _2 = PooledDictionary<Document, Dictionary<ISymbol, SymbolGroup>>.GetInstance(out var documentToSymbolAndGroup);
+        using var _2 = PooledDictionary<Document, Dictionary<ISymbol, SymbolGroup>>.GetInstance(out var documentToSymbolsWithin);
         try
         {
             // scratch hashset to place results in. Populated/inspected/cleared in inner loop.
@@ -239,13 +239,13 @@ await finder.DetermineDocumentsToSearchAsync(
 
                     foreach (var document in foundDocuments)
                     {
-                        if (!documentToSymbolAndGroup.TryGetValue(document, out var symbolAndGroup))
+                        if (!documentToSymbolsWithin.TryGetValue(document, out var symbolsWithin))
                         {
-                            symbolAndGroup = s_symbolToGroupPool.AllocateAndClear();
-                            documentToSymbolAndGroup.Add(document, symbolAndGroup);
+                            symbolsWithin = s_symbolToGroupPool.AllocateAndClear();
+                            documentToSymbolsWithin.Add(document, symbolsWithin);
                         }
 
-                        symbolAndGroup[symbol] = group;
+                        symbolsWithin[symbol] = group;
                     }
 
                     foundDocuments.Clear();
@@ -253,17 +253,17 @@ await finder.DetermineDocumentsToSearchAsync(
             }
 
             await RoslynParallel.ForEachAsync(
-                documentToSymbolAndGroup,
+                documentToSymbolsWithin,
                 GetParallelOptions(cancellationToken),
                 (kvp, cancellationToken) =>
                     ProcessDocumentAsync(kvp.Key, kvp.Value, symbolToGlobalAliases, onReferenceFound, cancellationToken)).ConfigureAwait(false);
         }
         finally
         {
-            foreach (var (_, symbolAndGroupMap) in documentToSymbolAndGroup)
+            foreach (var (_, symbolsWithin) in documentToSymbolsWithin)
             {
-                symbolAndGroupMap.Clear();
-                s_symbolToGroupPool.Free(symbolAndGroupMap);
+                symbolsWithin.Clear();
+                s_symbolToGroupPool.Free(symbolsWithin);
             }
 
             FreeGlobalAliases(symbolToGlobalAliases);
@@ -277,7 +277,7 @@ await RoslynParallel.ForEachAsync(
 
     private async ValueTask ProcessDocumentAsync(
         Document document,
-        Dictionary<ISymbol, SymbolGroup> symbolToSymbolGroup,
+        Dictionary<ISymbol, SymbolGroup> symbolsToSearchFor,
         Dictionary<ISymbol, PooledHashSet<string>> symbolToGlobalAliases,
         Action<Reference> onReferenceFound,
         CancellationToken cancellationToken)
@@ -297,32 +297,32 @@ private async ValueTask ProcessDocumentAsync(
         // Note: cascaded symbols will normally have the same name.  That's ok.  The second call to
         // FindMatchingIdentifierTokens with the same name will short circuit since it will already see the result of
         // the prior call.
-        foreach (var (symbol, _) in symbolToSymbolGroup)
+        foreach (var (symbol, _) in symbolsToSearchFor)
         {
             if (symbol.CanBeReferencedByName)
                 cache.FindMatchingIdentifierTokens(symbol.Name, cancellationToken);
         }
 
         await RoslynParallel.ForEachAsync(
-            symbolToSymbolGroup,
+            symbolsToSearchFor,
             GetParallelOptions(cancellationToken),
             (kvp, cancellationToken) =>
             {
-                var (symbol, group) = kvp;
+                var (symbolToSearchFor, symbolGroup) = kvp;
 
                 // symbolToGlobalAliases is safe to read in parallel.  It is created fully before this point and is no
                 // longer mutated.
                 var state = new FindReferencesDocumentState(
-                    cache, TryGet(symbolToGlobalAliases, symbol));
+                    cache, TryGet(symbolToGlobalAliases, symbolToSearchFor));
 
-                ProcessDocument(symbol, group, state, onReferenceFound);
+                ProcessDocument(symbolToSearchFor, symbolGroup, state, onReferenceFound);
                 return ValueTaskFactory.CompletedTask;
             }).ConfigureAwait(false);
 
         return;
 
         void ProcessDocument(
-            ISymbol symbol, SymbolGroup group, FindReferencesDocumentState state, Action<Reference> onReferenceFound)
+            ISymbol symbolToSearchFor, SymbolGroup symbolGroup, FindReferencesDocumentState state, Action<Reference> onReferenceFound)
         {
             cancellationToken.ThrowIfCancellationRequested();
 
@@ -334,9 +334,9 @@ void ProcessDocument(
                 foreach (var finder in _finders)
                 {
                     finder.FindReferencesInDocument(
-                        symbol, state,
-                        static (loc, tuple) => tuple.onReferenceFound((tuple.group, tuple.symbol, loc.Location)),
-                        (group, symbol, onReferenceFound),
+                        symbolToSearchFor, state,
+                        static (loc, tuple) => tuple.onReferenceFound((tuple.symbolGroup, tuple.symbolToSearchFor, loc.Location)),
+                        (symbolGroup, symbolToSearchFor, onReferenceFound),
                         _options,
                         cancellationToken);
                 }

From 8cf03d3f5d94e033ee2254ae2ef5ddb078d2a2f6 Mon Sep 17 00:00:00 2001
From: Cyrus Najmabadi <cyrus.najmabadi@gmail.com>
Date: Thu, 11 Jul 2024 12:36:06 -0700
Subject: [PATCH 8/9] Renames

---
 .../FindReferences/FindReferencesSearchEngine.cs           | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
index f2247e264638b..d6070403001e2 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
@@ -239,12 +239,7 @@ await finder.DetermineDocumentsToSearchAsync(
 
                     foreach (var document in foundDocuments)
                     {
-                        if (!documentToSymbolsWithin.TryGetValue(document, out var symbolsWithin))
-                        {
-                            symbolsWithin = s_symbolToGroupPool.AllocateAndClear();
-                            documentToSymbolsWithin.Add(document, symbolsWithin);
-                        }
-
+                        var symbolsWithin = documentToSymbolsWithin.GetOrAdd(document, static _ => s_symbolToGroupPool.AllocateAndClear());
                         symbolsWithin[symbol] = group;
                     }
 

From 0b0cb3239ca922130a64031ffe1755e718d3ec00 Mon Sep 17 00:00:00 2001
From: Cyrus Najmabadi <cyrus.najmabadi@gmail.com>
Date: Thu, 11 Jul 2024 12:37:18 -0700
Subject: [PATCH 9/9] Pooling

---
 .../FindSymbols/FindReferences/FindReferencesSearchEngine.cs  | 4 ++--
 .../FindReferencesSearchEngine_FindReferencesInDocuments.cs   | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
index d6070403001e2..a7f879ba4dd68 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine.cs
@@ -88,13 +88,13 @@ private async Task PerformSearchAsync(
         //
         // This is a normal dictionary that is not locked.  It is only ever read and written to serially from within the
         // high level project-walking code in this method.
-        var symbolToGroup = new Dictionary<ISymbol, SymbolGroup>(MetadataUnifyingEquivalenceComparer.Instance);
+        using var _1 = s_symbolToGroupPool.GetPooledObject(out var symbolToGroup);
 
         var unifiedSymbols = new MetadataUnifyingSymbolHashSet();
         unifiedSymbols.AddRange(symbols);
 
         var disposable = await _progressTracker.AddSingleItemAsync(cancellationToken).ConfigureAwait(false);
-        await using var _ = disposable.ConfigureAwait(false);
+        await using var _2 = disposable.ConfigureAwait(false);
 
         // Create the initial set of symbols to search for.  As we walk the appropriate projects in the solution
         // we'll expand this set as we discover new symbols to search for in each project.
diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
index ccbb11e3cf36d..18345b5713bef 100644
--- a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
+++ b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/FindReferencesSearchEngine_FindReferencesInDocuments.cs
@@ -32,7 +32,7 @@ public async Task FindReferencesInDocumentsAsync(
         //
         // This is a normal dictionary that is not locked.  It is only ever read and written to serially from within the
         // high level project-walking code in this method.
-        var symbolToGroup = new Dictionary<ISymbol, SymbolGroup>(MetadataUnifyingEquivalenceComparer.Instance);
+        using var _ = s_symbolToGroupPool.GetPooledObject(out var symbolToGroup);
 
         var unifiedSymbols = new MetadataUnifyingSymbolHashSet
         {