Speed up virtual method resolution in the compiler

src/coreclr/tools/Common/Compiler/CachingVirtualMethodAlgorithm.cs

@@ -0,0 +1,86 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Diagnostics;
+
+using Internal.TypeSystem;
+
+namespace ILCompiler
+{
+    public class CachingVirtualMethodAlgorithm : MetadataVirtualMethodAlgorithm
+    {
+        private readonly ConcurrentDictionary<TypeDesc, (MethodDesc Slot, MethodDesc Implementation)[]> _vtableCache
+            = new ConcurrentDictionary<TypeDesc, (MethodDesc Slot, MethodDesc Implementation)[]>();
+
+        private readonly Func<TypeDesc, (MethodDesc Slot, MethodDesc Implementation)[]> _vtableCreator;
+
+        public CachingVirtualMethodAlgorithm()
+            => _vtableCreator = ComputeVtable;
+
+        private (MethodDesc Slot, MethodDesc Implementation)[] ComputeVtable(TypeDesc type)
+        {
+            // Do not recompute the same things for all the generics
+            Debug.Assert(type.IsTypeDefinition);
+
+            var result = new ArrayBuilder<(MethodDesc Slot, MethodDesc Implementation)>(3 /* At least Equals/GetHashCode/ToString */);
+
+            foreach (MethodDesc slotDecl in base.ComputeAllVirtualSlots(type))
+            {
+                result.Add((slotDecl, base.FindVirtualFunctionTargetMethodOnObjectType(slotDecl, type)));
+            }
+
+            return result.ToArray();
+        }
+
+        public override IEnumerable<MethodDesc> ComputeAllVirtualSlots(TypeDesc type)
+        {
+            // This just enumerates virtual methods, not worth caching.
+            if (type.IsInterface)
+                return base.ComputeAllVirtualSlots(type);
+
+            return GetCachedSlots(this, type);
+
+            static IEnumerable<MethodDesc> GetCachedSlots(CachingVirtualMethodAlgorithm thisObj, TypeDesc type)
+            {
+                foreach ((MethodDesc slot, _) in thisObj._vtableCache.GetOrAdd(type.GetTypeDefinition(), thisObj._vtableCreator))
+                {
+                    yield return type.FindMethodOnTypeWithMatchingTypicalMethod(slot);
+                }
+            }
+        }
+
+        public override MethodDesc FindVirtualFunctionTargetMethodOnObjectType(MethodDesc targetMethod, TypeDesc objectType)
+        {
+            MetadataType uninstantiatedType = (MetadataType)objectType.GetTypeDefinition();
+            MethodDesc targetMethodDefinition = targetMethod.GetMethodDefinition();
+
+            MethodDesc slotDefiningMethod = targetMethodDefinition;
+            if (uninstantiatedType != objectType)
+            {
+                slotDefiningMethod = uninstantiatedType.FindMethodOnTypeWithMatchingTypicalMethod(slotDefiningMethod);
+            }
+            slotDefiningMethod = FindSlotDefiningMethodForVirtualMethod(slotDefiningMethod);
+
+            foreach (var kvp in _vtableCache.GetOrAdd(objectType.GetTypeDefinition(), _vtableCreator))
+            {
+                if (kvp.Slot == slotDefiningMethod)
+                {
+                    MethodDesc result = kvp.Implementation;
+
+                    if (uninstantiatedType != objectType)
+                        result = objectType.FindMethodOnTypeWithMatchingTypicalMethod(result);
+
+                    if (targetMethod != targetMethodDefinition)
+                        result = result.MakeInstantiatedMethod(targetMethod.Instantiation);
+
+                    return result;
+                }
+            }
+
+            return null;
+        }
+    }
+}

src/coreclr/tools/Common/Compiler/CompilerTypeSystemContext.Async.cs

@@ -21,7 +21,7 @@ public static IEnumerable<MethodDesc> GetAllMethodsAndAsyncVariants(this TypeDes
 
     public partial class CompilerTypeSystemContext
     {
-        private sealed class AsyncAwareVirtualMethodResolutionAlgorithm : MetadataVirtualMethodAlgorithm
+        private sealed class AsyncAwareVirtualMethodResolutionAlgorithm : CachingVirtualMethodAlgorithm
         {
             private readonly CompilerTypeSystemContext _context;
 

src/coreclr/tools/Common/Compiler/DependencyAnalysis/GVMDependenciesNode.cs

@@ -32,7 +32,7 @@ public GVMDependenciesNode(MethodDesc method)
             Debug.Assert(method.GetCanonMethodTarget(CanonicalFormKind.Specific) == method);
             Debug.Assert(method.HasInstantiation);
             Debug.Assert(method.IsVirtual);
-            Debug.Assert(MetadataVirtualMethodAlgorithm.FindSlotDefiningMethodForVirtualMethod(method) == method);
+            Debug.Assert(MetadataVirtualMethodAlgorithm.FindSlotDefiningMethodForVirtualMethod(method.GetMethodDefinition()) == method.GetMethodDefinition());
 
             _method = method;
         }

src/coreclr/tools/Common/Compiler/GenericCycleDetection/GraphBuilder.cs

@@ -270,7 +270,7 @@ private void LookForVirtualOverrides(EcmaMethod method)
                 // of the implementation to the generic parameters of the declaration - any call to the
                 // declaration will be modeled as if the declaration was calling into the implementation.
 
-                var decl = (EcmaMethod)MetadataVirtualMethodAlgorithm.FindSlotDefiningMethodForVirtualMethod(method).GetTypicalMethodDefinition();
+                var decl = (EcmaMethod)MetadataVirtualMethodAlgorithm.FindSlotDefiningMethodForVirtualMethod(method.GetMethodDefinition()).GetTypicalMethodDefinition();
                 if (decl != method)
                 {
                     RecordBinding(this, decl.Instantiation, method.Instantiation);

src/coreclr/tools/Common/TypeSystem/Common/MetadataVirtualMethodAlgorithm.cs

@@ -239,6 +239,7 @@ private static MethodDesc FindVirtualFunctionTargetMethodOnObjectType(MethodDesc
 
             // Step 2, convert targetMethod to method in type hierarchy of uninstantiated form
             targetMethod = targetMethod.GetMethodDefinition();
+            MethodDesc initialTargetMethodDefinition = targetMethod;
             if (uninstantiatedType != objectType)
             {
                 targetMethod = uninstantiatedType.FindMethodOnTypeWithMatchingTypicalMethod(targetMethod);
@@ -259,7 +260,7 @@ private static MethodDesc FindVirtualFunctionTargetMethodOnObjectType(MethodDesc
             {
                 resolutionTarget = objectType.FindMethodOnTypeWithMatchingTypicalMethod(resolutionTarget);
             }
-            if (initialTargetMethod.HasInstantiation)
+            if (initialTargetMethod != initialTargetMethodDefinition)
             {
                 resolutionTarget = resolutionTarget.MakeInstantiatedMethod(initialTargetMethod.Instantiation);
             }
@@ -441,6 +442,8 @@ public static MethodDesc FindSlotDefiningMethodForVirtualMethod(MethodDesc metho
             if (method == null)
                 return method;
 
+            Debug.Assert(method.GetMethodDefinition() == method);
+
             DefType currentType = method.OwningType.BaseType;
 
             // Loop until a newslot method is found

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/Compilation.cs

@@ -123,7 +123,13 @@ public DelegateCreationInfo GetDelegateCtor(TypeDesc delegateType, MethodDesc ta
                 followVirtualDispatch = false;
 
             if (followVirtualDispatch)
-                target = MetadataVirtualMethodAlgorithm.FindSlotDefiningMethodForVirtualMethod(target);
+            {
+                MethodDesc originalTarget = target;
+                MethodDesc targetDefinition = target.GetMethodDefinition();
+                target = MetadataVirtualMethodAlgorithm.FindSlotDefiningMethodForVirtualMethod(targetDefinition);
+                if (originalTarget != targetDefinition)
+                    target = target.MakeInstantiatedMethod(originalTarget.Instantiation);
+            }
 
             return DelegateCreationInfo.Create(delegateType, target, constrainedType, NodeFactory, followVirtualDispatch);
         }

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/DependencyAnalysis/ReflectionVirtualInvokeMapNode.cs

@@ -86,11 +86,11 @@ public static void GetVirtualInvokeMapDependencies(ref DependencyList dependenci
                     factory.NecessaryTypeSymbol(method.OwningType.ConvertToCanonForm(CanonicalFormKind.Specific)),
                     "Reflection virtual invoke owning type");
 
-                MethodDesc slotDefiningMethod = MetadataVirtualMethodAlgorithm.FindSlotDefiningMethodForVirtualMethod(method);
+                MethodDesc methodDefinition = method.GetMethodDefinition();
+                MethodDesc slotDefiningMethod = MetadataVirtualMethodAlgorithm.FindSlotDefiningMethodForVirtualMethod(methodDefinition);
                 if (method.HasInstantiation)
                 {
-                    // FindSlotDefiningMethod might uninstantiate. We might want to fix the method not to do that.
-                    if (slotDefiningMethod.IsMethodDefinition)
+                    if (method != methodDefinition)
                         slotDefiningMethod = factory.TypeSystemContext.GetInstantiatedMethod(slotDefiningMethod, method.Instantiation);
                     dependencies.Add(factory.GVMDependencies(slotDefiningMethod.GetCanonMethodTarget(CanonicalFormKind.Specific)), "GVM callable reflectable method");
                 }

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/UsageBasedMetadataManager.cs

@@ -631,7 +631,8 @@ public override void GetDependenciesDueToDelegateCreation(ref CombinedDependency
 
         public override void GetDependenciesForOverridingMethod(ref CombinedDependencyList dependencies, NodeFactory factory, MethodDesc decl, MethodDesc impl)
         {
-            Debug.Assert(decl.IsVirtual && MetadataVirtualMethodAlgorithm.FindSlotDefiningMethodForVirtualMethod(decl) == decl);
+            Debug.Assert(decl.IsVirtual
+                && MetadataVirtualMethodAlgorithm.FindSlotDefiningMethodForVirtualMethod(decl.GetMethodDefinition()) == decl.GetMethodDefinition());
 
             // If a virtual method slot is a target of a delegate, all implementations become reflection visible
             // to support Delegate.GetMethodInfo().

src/coreclr/tools/aot/ILCompiler.Compiler/ILCompiler.Compiler.csproj

@@ -36,6 +36,7 @@
     <Compile Include="..\..\Common\Compiler\AsyncContinuationType.cs" Link="Compiler\AsyncContinuationType.cs" />
     <Compile Include="..\..\Common\Compiler\AsyncMethodVariant.cs" Link="Compiler\AsyncMethodVariant.cs" />
     <Compile Include="..\..\Common\Compiler\AsyncMethodVariant.Mangling.cs" Link="Compiler\AsyncMethodVariant.Mangling.cs" />
+    <Compile Include="..\..\Common\Compiler\CachingVirtualMethodAlgorithm.cs" Link="Compiler\CachingVirtualMethodAlgorithm.cs" />
     <Compile Include="..\..\Common\Compiler\CompilerTypeSystemContext.Async.cs" Link="Compiler\CompilerTypeSystemContext.Async.cs" />
     <Compile Include="..\..\Common\Compiler\CompilerTypeSystemContext.Wasm.cs" Link="Compiler\CompilerTypeSystemContext.Wasm.cs" />
     <Compile Include="..\..\Common\Compiler\Win32Resources\ResourceData.cs" Link="Compiler\Win32Resources\ResourceData.cs" />

src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/DependencyAnalysis/ReadyToRunCodegenNodeFactory.cs

@@ -154,8 +154,8 @@ public GVMDependenciesNode GVMDependencies(MethodDesc method)
         {
             Debug.Assert(method.IsVirtual);
             MethodDesc canonMethod = method.GetCanonMethodTarget(CanonicalFormKind.Specific);
-            canonMethod = MetadataVirtualMethodAlgorithm.FindSlotDefiningMethodForVirtualMethod(canonMethod);
-            return _gvmDependenciesNode.GetOrAdd(canonMethod);
+            MethodDesc canonSlotMethodDefinition = MetadataVirtualMethodAlgorithm.FindSlotDefiningMethodForVirtualMethod(canonMethod.GetMethodDefinition());
+            return _gvmDependenciesNode.GetOrAdd(canonSlotMethodDefinition.MakeInstantiatedMethod(canonMethod.Instantiation));
         }
 
         private NodeCache<MethodDesc, VirtualMethodUseNode> _virtualMethodUseNodes;

src/coreclr/tools/aot/ILCompiler.ReadyToRun/ILCompiler.ReadyToRun.csproj

@@ -64,6 +64,7 @@
     <Compile Include="..\..\Common\Compiler\AsyncContinuationType.cs" Link="Compiler\AsyncContinuationType.cs" />
     <Compile Include="..\..\Common\Compiler\AsyncMethodVariant.cs" Link="Compiler\AsyncMethodVariant.cs" />
     <Compile Include="..\..\Common\Compiler\AsyncMethodVariant.Mangling.cs" Link="Compiler\AsyncMethodVariant.Mangling.cs" />
+    <Compile Include="..\..\Common\Compiler\CachingVirtualMethodAlgorithm.cs" Link="Compiler\CachingVirtualMethodAlgorithm.cs" />
     <Compile Include="..\..\Common\Compiler\CodeGenerationFailedException.cs" Link="Compiler\CodeGenerationFailedException.cs" />
     <Compile Include="..\..\Common\Compiler\CompilationBuilder.cs" Link="Compiler\CompilationBuilder.cs" />
     <Compile Include="..\..\Common\Compiler\CompilationModuleGroup.cs" Link="Compiler\CompilationModuleGroup.cs" />