Reduce allocations in ModuleSymbolKey.Resolve

Related to https://devdiv.visualstudio.com/DevDiv/_workitems/edit/2136177

Author: sharwell

src/Compilers/CSharp/Portable/Symbols/PublicModel/AssemblySymbol.cs

@@ -4,21 +4,18 @@
 
 #nullable disable
 
-using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics;
-using System.Globalization;
 using System.Linq;
-using System.Threading;
-using Microsoft.CodeAnalysis.Collections;
-using Microsoft.CodeAnalysis.PooledObjects;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.CSharp.Symbols.PublicModel
 {
     internal abstract class AssemblySymbol : Symbol, IAssemblySymbol
     {
+        private IEnumerable<IModuleSymbol> _lazyModules;
+
         internal abstract Symbols.AssemblySymbol UnderlyingAssemblySymbol { get; }
 
         INamespaceSymbol IAssemblySymbol.GlobalNamespace
@@ -33,10 +30,10 @@ IEnumerable<IModuleSymbol> IAssemblySymbol.Modules
         {
             get
             {
-                foreach (var module in UnderlyingAssemblySymbol.Modules)
-                {
-                    yield return module.GetPublicSymbol();
-                }
+                return InterlockedOperations.Initialize(
+                    ref _lazyModules,
+                    static self => self.UnderlyingAssemblySymbol.Modules.SelectAsArray(static module => module.GetPublicSymbol()),
+                    this);
             }
         }
 

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/SymbolKey/SymbolKey.ModuleSymbolKey.cs

@@ -2,6 +2,8 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Collections.Immutable;
+
 namespace Microsoft.CodeAnalysis;
 
 internal partial struct SymbolKey
@@ -24,12 +26,27 @@ protected sealed override SymbolKeyResolution Resolve(
                 return default;
             }
 
-            using var result = PooledArrayBuilder<IModuleSymbol>.GetInstance();
-            foreach (var assembly in containingSymbolResolution.OfType<IAssemblySymbol>())
+            using var result = PooledArrayBuilder<IModuleSymbol>.GetInstance(containingSymbolResolution.SymbolCount);
+            foreach (var symbol in containingSymbolResolution)
             {
+                if (symbol is not IAssemblySymbol assembly)
+                    continue;
+
                 // Don't check ModuleIds for equality because in practice, no-one uses them,
                 // and there is no way to set netmodule name programmatically using Roslyn
-                result.AddValuesIfNotNull(assembly.Modules);
+                var assemblyModules = assembly.Modules;
+                if (assemblyModules is ImmutableArray<IModuleSymbol> modules)
+                {
+                    // Avoid allocations if possible
+                    // https://devdiv.visualstudio.com/DevDiv/_workitems/edit/2136177
+                    result.AddValuesIfNotNull(modules);
+                }
+                else
+                {
+                    // Visual Basic implementation of IAssemblySymbol.Modules relies on covariance of IEnumerable<T>, so
+                    // the preceding concrete type check will fail.
+                    result.AddValuesIfNotNull(assemblyModules);
+                }
             }
 
             return CreateResolution(result, $"({nameof(ModuleSymbolKey)} failed)", out failureReason);