Cache values in dependent project finder

src/Workspaces/Core/Portable/FindSymbols/FindReferences/DependentProjectsFinder.cs

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Runtime.CompilerServices;
@@ -28,7 +29,15 @@ internal static partial class DependentProjectsFinder
     /// Cache from the <see cref="MetadataId"/> for a particular <see cref="PortableExecutableReference"/> to the
     /// name of the <see cref="IAssemblySymbol"/> defined by it.
     /// </summary>
-    private static ImmutableDictionary<MetadataId, string> s_metadataIdToAssemblyName = ImmutableDictionary<MetadataId, string>.Empty;
+    private static readonly Dictionary<MetadataId, string> s_metadataIdToAssemblyName = new();
+    private static readonly SemaphoreSlim s_metadataIdToAssemblyNameGate = new(initialCount: 1);
+
+    private static readonly ConditionalWeakTable<
+         Solution,
+         Dictionary<
+             (IAssemblySymbol assembly, Project? sourceProject, SymbolVisibility visibility),
+             ImmutableArray<(Project project, bool hasInternalsAccess)>>> s_solutionToDependentProjectMap = new();
+    private static readonly SemaphoreSlim s_solutionToDependentProjectMapGate = new(initialCount: 1);
 
     public static async Task<ImmutableArray<Project>> GetDependentProjectsAsync(
         Solution solution, ImmutableArray<ISymbol> symbols, IImmutableSet<Project> projects, CancellationToken cancellationToken)
@@ -128,24 +137,59 @@ private static async Task<ImmutableArray<Project>> GetDependentProjectsWorkerAsy
         SymbolVisibility visibility,
         CancellationToken cancellationToken)
     {
-        cancellationToken.ThrowIfCancellationRequested();
+        var dictionary = s_solutionToDependentProjectMap.GetValue(solution, static _ => new());
 
-        var dependentProjects = new HashSet<(Project, bool hasInternalsAccess)>();
+        var key = (symbolOrigination.assembly, symbolOrigination.sourceProject, visibility);
+        ImmutableArray<(Project project, bool hasInternalsAccess)> dependentProjects;
 
-        // If a symbol was defined in source, then it is always visible to the project it
-        // was defined in.
-        if (symbolOrigination.sourceProject != null)
-            dependentProjects.Add((symbolOrigination.sourceProject, hasInternalsAccess: true));
+        // Check cache first.
+        using (await s_solutionToDependentProjectMapGate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
+        {
+            if (dictionary.TryGetValue(key, out dependentProjects))
+                return dependentProjects;
+        }
+
+        // Compute if not in cache.
+        dependentProjects = await ComputeDependentProjectsWorkerAsync(
+            solution, symbolOrigination, visibility, cancellationToken).ConfigureAwait(false);
+
+        // Try to add to cache, returning existing value if another thread already added it.
+        using (await s_solutionToDependentProjectMapGate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
+        {
+            if (dictionary.TryGetValue((symbolOrigination.assembly, symbolOrigination.sourceProject, visibility), out dependentProjects))
+                return dependentProjects;
+
+            return dictionary.GetOrAdd(key, dependentProjects);
+        }
+
+        static async Task<ImmutableArray<(Project project, bool hasInternalsAccess)>> ComputeDependentProjectsWorkerAsync(
+           Solution solution,
+           (IAssemblySymbol assembly, Project? sourceProject) symbolOrigination,
+           SymbolVisibility visibility,
+           CancellationToken cancellationToken)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+            using var _ = PooledHashSet<(Project, bool hasInternalsAccess)>.GetInstance(out var dependentProjects);
 
-        // If it's not private, then we need to find possible references.
-        if (visibility != SymbolVisibility.Private)
-            AddNonSubmissionDependentProjects(solution, symbolOrigination, dependentProjects, cancellationToken);
+            // If a symbol was defined in source, then it is always visible to the project it
+            // was defined in.
+            if (symbolOrigination.sourceProject != null)
+                dependentProjects.Add((symbolOrigination.sourceProject, hasInternalsAccess: true));
 
-        // submission projects are special here. The fields generated inside the Script object is private, but
-        // further submissions can bind to them.
-        await AddSubmissionDependentProjectsAsync(solution, symbolOrigination.sourceProject, dependentProjects, cancellationToken).ConfigureAwait(false);
+            // If it's not private, then we need to find possible references.
+            if (visibility != SymbolVisibility.Private)
+            {
+                await AddNonSubmissionDependentProjectsAsync(
+                    solution, symbolOrigination, dependentProjects, cancellationToken).ConfigureAwait(false);
+            }
 
-        return [.. dependentProjects];
+            // submission projects are special here. The fields generated inside the Script object is private, but
+            // further submissions can bind to them.
+            await AddSubmissionDependentProjectsAsync(solution, symbolOrigination.sourceProject, dependentProjects, cancellationToken).ConfigureAwait(false);
+
+            return [.. dependentProjects];
+        }
     }
 
     private static async Task AddSubmissionDependentProjectsAsync(
@@ -221,7 +265,7 @@ private static bool IsInternalsVisibleToAttribute(AttributeData attr)
                attrType.ContainingNamespace.ContainingNamespace.ContainingNamespace.ContainingNamespace?.IsGlobalNamespace == true;
     }
 
-    private static void AddNonSubmissionDependentProjects(
+    private static async Task AddNonSubmissionDependentProjectsAsync(
         Solution solution,
         (IAssemblySymbol assembly, Project? sourceProject) symbolOrigination,
         HashSet<(Project project, bool hasInternalsAccess)> dependentProjects,
@@ -235,7 +279,7 @@ private static void AddNonSubmissionDependentProjects(
         foreach (var project in solution.Projects)
         {
             if (!project.SupportsCompilation ||
-                !HasReferenceTo(symbolOrigination, project, cancellationToken))
+                !await HasReferenceToAsync(symbolOrigination, project, cancellationToken).ConfigureAwait(false))
             {
                 continue;
             }
@@ -270,7 +314,7 @@ private static HashSet<string> GetInternalsVisibleToSet(IAssemblySymbol assembly
         return set;
     }
 
-    private static bool HasReferenceTo(
+    private static async Task<bool> HasReferenceToAsync(
         (IAssemblySymbol assembly, Project? sourceProject) symbolOrigination,
         Project project,
         CancellationToken cancellationToken)
@@ -284,10 +328,11 @@ private static bool HasReferenceTo(
             return project.ProjectReferences.Any(p => p.ProjectId == symbolOrigination.sourceProject.Id);
 
         // Otherwise, if the symbol is from metadata, see if the project's compilation references that metadata assembly.
-        return HasReferenceToAssembly(project, symbolOrigination.assembly.Name, cancellationToken);
+        return await HasReferenceToAssemblyAsync(
+            project, symbolOrigination.assembly.Name, cancellationToken).ConfigureAwait(false);
     }
 
-    private static bool HasReferenceToAssembly(Project project, string assemblyName, CancellationToken cancellationToken)
+    private static async Task<bool> HasReferenceToAssemblyAsync(Project project, string assemblyName, CancellationToken cancellationToken)
     {
         Contract.ThrowIfFalse(project.SupportsCompilation);
 
@@ -307,14 +352,17 @@ private static bool HasReferenceToAssembly(Project project, string assemblyName,
             if (metadataId is null)
                 continue;
 
-            if (!s_metadataIdToAssemblyName.TryGetValue(metadataId, out var name))
+            using (await s_metadataIdToAssemblyNameGate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
             {
-                uncomputedReferences.Add((peReference, metadataId));
-                continue;
-            }
+                if (!s_metadataIdToAssemblyName.TryGetValue(metadataId, out var name))
+                {
+                    uncomputedReferences.Add((peReference, metadataId));
+                    continue;
+                }
 
-            if (name == assemblyName)
-                return true;
+                if (name == assemblyName)
+                    return true;
+            }
         }
 
         if (uncomputedReferences.Count == 0)
@@ -326,16 +374,24 @@ private static bool HasReferenceToAssembly(Project project, string assemblyName,
         {
             cancellationToken.ThrowIfCancellationRequested();
 
-            if (!s_metadataIdToAssemblyName.TryGetValue(metadataId, out var name))
+            using (await s_metadataIdToAssemblyNameGate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
             {
-                // Defer creating the compilation till needed.
-                CreateCompilation(project, ref compilation);
-                if (compilation.GetAssemblyOrModuleSymbol(peReference) is IAssemblySymbol { Name: string metadataAssemblyName })
-                    name = ImmutableInterlocked.GetOrAdd(ref s_metadataIdToAssemblyName, metadataId, metadataAssemblyName);
+                if (s_metadataIdToAssemblyName.TryGetValue(metadataId, out var name) && name == assemblyName)
+                    return true;
             }
 
-            if (name == assemblyName)
-                return true;
+            // Defer creating the compilation till needed.
+            CreateCompilation(project, ref compilation);
+
+            if (compilation.GetAssemblyOrModuleSymbol(peReference) is IAssemblySymbol { Name: string metadataAssemblyName })
+            {
+                using (await s_metadataIdToAssemblyNameGate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
+                {
+                    s_metadataIdToAssemblyName.TryAdd(metadataId, metadataAssemblyName);
+                    if (metadataAssemblyName == assemblyName)
+                        return true;
+                }
+            }
         }
 
         return false;