NativeAOT: Guarded devirtualization

src/coreclr/inc/corinfo.h

@@ -2790,6 +2790,15 @@ class ICorStaticInfo
             CORINFO_CLASS_HANDLE       *vcTypeRet       /* OUT */
             ) = 0;
 
+    // Obtains a list of exact classes for a given base type. Returns 0 if the number of 
+    // the exact classes is greater than maxExactClasses or if more types might be loaded
+    // in future.
+    virtual int getExactClasses(
+                CORINFO_CLASS_HANDLE  baseType,            /* IN */
+                int                   maxExactClasses,     /* IN */
+                CORINFO_CLASS_HANDLE* exactClsRet          /* OUT */
+                ) = 0;
+
     // If the Arg is a CORINFO_TYPE_CLASS fetch the class handle associated with it
     virtual CORINFO_CLASS_HANDLE getArgClass (
             CORINFO_SIG_INFO*           sig,            /* IN */

src/coreclr/inc/icorjitinfoimpl_generated.h

@@ -430,6 +430,11 @@ CorInfoTypeWithMod getArgType(
           CORINFO_ARG_LIST_HANDLE args,
           CORINFO_CLASS_HANDLE* vcTypeRet) override;
 
+int getExactClasses(
+          CORINFO_CLASS_HANDLE baseType,
+          int maxExactClasses,
+          CORINFO_CLASS_HANDLE* exactClsRet) override;
+
 CORINFO_CLASS_HANDLE getArgClass(
           CORINFO_SIG_INFO* sig,
           CORINFO_ARG_LIST_HANDLE args) override;

src/coreclr/inc/jiteeversionguid.h

@@ -43,11 +43,11 @@ typedef const GUID *LPCGUID;
 #define GUID_DEFINED
 #endif // !GUID_DEFINED
 
-constexpr GUID JITEEVersionIdentifier = { /* 4efa8fe2-8489-4b61-aac9-b4df74af15b7 */
-    0x4efa8fe2,
-    0x8489,
-    0x4b61,
-    {0xaa, 0xc9, 0xb4, 0xdf, 0x74, 0xaf, 0x15, 0xb7}
+constexpr GUID JITEEVersionIdentifier = { /* 1b9551b8-21f4-4233-9c90-f3eabd6a322b */
+    0x1b9551b8,
+    0x21f4,
+    0x4233,
+    {0x9c, 0x90, 0xf3, 0xea, 0xbd, 0x6a, 0x32, 0x2b}
   };
 
 //////////////////////////////////////////////////////////////////////////////////////////////////////////

src/coreclr/jit/ICorJitInfo_API_names.h

@@ -107,6 +107,7 @@ DEF_CLR_API(allocateArray)
 DEF_CLR_API(freeArray)
 DEF_CLR_API(getArgNext)
 DEF_CLR_API(getArgType)
+DEF_CLR_API(getExactClasses)
 DEF_CLR_API(getArgClass)
 DEF_CLR_API(getHFAType)
 DEF_CLR_API(GetErrorHRESULT)

src/coreclr/jit/ICorJitInfo_API_wrapper.hpp

@@ -1015,6 +1015,17 @@ CorInfoTypeWithMod WrapICorJitInfo::getArgType(
     return temp;
 }
 
+int WrapICorJitInfo::getExactClasses(
+          CORINFO_CLASS_HANDLE baseType,
+          int maxExactClasses,
+          CORINFO_CLASS_HANDLE* exactClsRet)
+{
+    API_ENTER(getExactClasses);
+    int temp = wrapHnd->getExactClasses(baseType, maxExactClasses, exactClsRet);
+    API_LEAVE(getExactClasses);
+    return temp;
+}
+
 CORINFO_CLASS_HANDLE WrapICorJitInfo::getArgClass(
           CORINFO_SIG_INFO* sig,
           CORINFO_ARG_LIST_HANDLE args)

src/coreclr/jit/gentree.cpp

@@ -17771,6 +17771,16 @@ CORINFO_CLASS_HANDLE Compiler::gtGetClassHandle(GenTree* tree, bool* pIsExact, b
         }
     }
 
+    if ((objClass != NO_CLASS_HANDLE) && !*pIsExact && IsTargetAbi(CORINFO_NATIVEAOT_ABI))
+    {
+        CORINFO_CLASS_HANDLE exactClass;
+        if (info.compCompHnd->getExactClasses(objClass, 1, &exactClass) == 1)
+        {
+            *pIsExact = true;
+            objClass  = exactClass;
+        }
+    }
+
     return objClass;
 }
 

src/coreclr/jit/lclvars.cpp

@@ -3174,6 +3174,16 @@ void Compiler::lvaSetClass(unsigned varNum, CORINFO_CLASS_HANDLE clsHnd, bool is
         return;
     }
 
+    if (clsHnd != NO_CLASS_HANDLE && !isExact && IsTargetAbi(CORINFO_NATIVEAOT_ABI))
+    {
+        CORINFO_CLASS_HANDLE exactClass;
+        if (info.compCompHnd->getExactClasses(clsHnd, 1, &exactClass) == 1)
+        {
+            isExact = true;
+            clsHnd  = exactClass;
+        }
+    }
+
     // Else we should have a type handle.
     assert(clsHnd != nullptr);
 

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

@@ -211,6 +211,8 @@ protected virtual MethodDesc ResolveVirtualMethod(MethodDesc declMethod, DefType
         /// so it can answer this question.
         /// </remarks>
         public virtual bool CanConstructType(TypeDesc type) => true;
+
+        public virtual TypeDesc[] GetImplementingClasses(TypeDesc type) => null;
 #endif
     }
 }

src/coreclr/tools/Common/JitInterface/ThunkGenerator/ThunkInput.txt

@@ -258,6 +258,7 @@ FUNCTIONS
     void freeArray(void*array);
     CORINFO_ARG_LIST_HANDLE getArgNext(CORINFO_ARG_LIST_HANDLE args);
     CorInfoTypeWithMod getArgType(CORINFO_SIG_INFO* sig, CORINFO_ARG_LIST_HANDLE args, CORINFO_CLASS_HANDLE* vcTypeRet);
+    int getExactClasses(CORINFO_CLASS_HANDLE baseType, int maxExactClasses, CORINFO_CLASS_HANDLE* exactClsRet);
     CORINFO_CLASS_HANDLE getArgClass(CORINFO_SIG_INFO* sig, CORINFO_ARG_LIST_HANDLE args);
     CorInfoHFAElemType getHFAType(CORINFO_CLASS_HANDLE hClass);
     JITINTERFACE_HRESULT GetErrorHRESULT(struct _EXCEPTION_POINTERS *pExceptionPointers);

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

@@ -219,6 +219,11 @@ public bool IsEffectivelySealed(TypeDesc type)
             return _devirtualizationManager.IsEffectivelySealed(type);
         }
 
+        public TypeDesc[] GetImplementingClasses(TypeDesc type)
+        {
+            return _devirtualizationManager.GetImplementingClasses(type);
+        }
+
         public bool IsEffectivelySealed(MethodDesc method)
         {
             return _devirtualizationManager.IsEffectivelySealed(method);

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

@@ -367,6 +367,8 @@ private class ScannedDevirtualizationManager : DevirtualizationManager
             private HashSet<TypeDesc> _constructedTypes = new HashSet<TypeDesc>();
             private HashSet<TypeDesc> _canonConstructedTypes = new HashSet<TypeDesc>();
             private HashSet<TypeDesc> _unsealedTypes = new HashSet<TypeDesc>();
+            private Dictionary<TypeDesc, HashSet<TypeDesc>> _interfaceImplementators = new();
+            private HashSet<TypeDesc> _disqualifiedInterfaces = new();
 
             public ScannedDevirtualizationManager(ImmutableArray<DependencyNodeCore<NodeFactory>> markedNodes)
             {
@@ -381,7 +383,32 @@ public ScannedDevirtualizationManager(ImmutableArray<DependencyNodeCore<NodeFact
 
                     if (type != null)
                     {
-                        if (!type.IsInterface)
+                        if (type.IsCanonicalSubtype(CanonicalFormKind.Any))
+                        {
+                            foreach (DefType baseInterface in type.RuntimeInterfaces)
+                            {
+                                // If the interface is implemented on a template type, there might be
+                                // no real upper bound on the number of actual classes implementing it
+                                // due to MakeGenericType.
+                                if (CanAssumeWholeProgramViewOnInterfaceUse(baseInterface))
+                                    _disqualifiedInterfaces.Add(baseInterface);
+                            }
+                        }
+
+                        if (type.IsInterface)
+                        {
+                            if (((MetadataType)type).IsDynamicInterfaceCastableImplementation())
+                            {
+                                foreach (DefType baseInterface in type.RuntimeInterfaces)
+                                {
+                                    // If the interface is implemented through IDynamicInterfaceCastable, there might be
+                                    // no real upper bound on the number of actual classes implementing it.
+                                    if (CanAssumeWholeProgramViewOnInterfaceUse(baseInterface))
+                                        _disqualifiedInterfaces.Add(baseInterface);
+                                }
+                            }
+                        }
+                        else
                         {
                             //
                             // We collect this information:
@@ -390,15 +417,28 @@ public ScannedDevirtualizationManager(ImmutableArray<DependencyNodeCore<NodeFact
                             // 2. What types are the base types of other types
                             //    This is needed for optimizations. We use this information to effectively
                             //    seal types that are not base types for any other type.
+                            // 3. What types implement interfaces for which use we can assume whole
+                            //    program view.
                             //
 
                             if (!type.IsCanonicalSubtype(CanonicalFormKind.Any))
                                 _constructedTypes.Add(type);
 
+                            if (type is not MetadataType { IsAbstract: true })
+                            {
+                                // Record all interfaces this class implements to _interfaceImplementators
+                                foreach (DefType baseInterface in type.RuntimeInterfaces)
+                                {
+                                    if (CanAssumeWholeProgramViewOnInterfaceUse(baseInterface))
+                                    {
+                                        RecordImplementation(baseInterface, type);
+                                    }
+                                }
+                            }
+
                             TypeDesc canonType = type.ConvertToCanonForm(CanonicalFormKind.Specific);
                             _canonConstructedTypes.Add(canonType.GetClosestDefType());
 
-                            bool hasNonAbstractTypeInHierarchy = canonType is not MetadataType mdType || !mdType.IsAbstract;
                             TypeDesc baseType = canonType.BaseType;
                             bool added = true;
                             while (baseType != null && added)
@@ -412,6 +452,61 @@ public ScannedDevirtualizationManager(ImmutableArray<DependencyNodeCore<NodeFact
                 }
             }
 
+            private static bool CanAssumeWholeProgramViewOnInterfaceUse(DefType interfaceType)
+            {
+                if (!interfaceType.HasInstantiation)
+                {
+                    return true;
+                }
+
+                foreach (GenericParameterDesc genericParam in interfaceType.GetTypeDefinition().Instantiation)
+                {
+                    if (genericParam.Variance != GenericVariance.None)
+                    {
+                        // If the interface has any variance, this gets complicated.
+                        // Skip for now.
+                        return false;
+                    }
+                }
+
+                if (((CompilerTypeSystemContext)interfaceType.Context).IsGenericArrayInterfaceType(interfaceType))
+                {
+                    // Interfaces implemented by arrays also behave covariantly on arrays even though
+                    // they're not actually variant. Skip for now.
+                    return false;
+                }
+
+                if (interfaceType.IsCanonicalSubtype(CanonicalFormKind.Any)
+                    || interfaceType.ConvertToCanonForm(CanonicalFormKind.Specific) != interfaceType
+                    || interfaceType.Context.SupportsUniversalCanon)
+                {
+                    // If the interface has a canonical form, we might not have a full view of all implementers.
+                    // E.g. if we have:
+                    // class Fooer<T> : IFooable<T> { }
+                    // class Doer<T> : IFooable<T> { }
+                    // And we instantiated Fooer<string>, but not Doer<string>. But we do have code for Doer<__Canon>.
+                    // We might think we can devirtualize IFooable<string> to Fooer<string>, but someone could
+                    // typeof(Doer<>).MakeGenericType(typeof(string)) and break our whole program view.
+                    // This is only a problem if canonical form of the interface exists.
+                    return false;
+                }
+
+                return true;
+            }
+
+            private void RecordImplementation(TypeDesc type, TypeDesc implType)
+            {
+                Debug.Assert(!implType.IsInterface);
+
+                HashSet<TypeDesc> implList;
+                if (!_interfaceImplementators.TryGetValue(type, out implList))
+                {
+                    implList = new();
+                    _interfaceImplementators[type] = implList;
+                }
+                implList.Add(implType);
+            }
+
             public override bool IsEffectivelySealed(TypeDesc type)
             {
                 // If we know we scanned a type that derives from this one, this for sure can't be reported as sealed.
@@ -456,6 +551,24 @@ protected override MethodDesc ResolveVirtualMethod(MethodDesc declMethod, DefTyp
             }
 
             public override bool CanConstructType(TypeDesc type) => _constructedTypes.Contains(type);
+
+            public override TypeDesc[] GetImplementingClasses(TypeDesc type)
+            {
+                if (_disqualifiedInterfaces.Contains(type))
+                    return null;
+
+                if (type.IsInterface && _interfaceImplementators.TryGetValue(type, out HashSet<TypeDesc> implementations))
+                {
+                    var types = new TypeDesc[implementations.Count];
+                    int index = 0;
+                    foreach (TypeDesc implementation in implementations)
+                    {
+                        types[index++] = implementation;
+                    }
+                    return types;
+                }
+                return null;
+            }
         }
 
         private class ScannedInliningPolicy : IInliningPolicy

src/coreclr/tools/aot/ILCompiler.ReadyToRun/JitInterface/CorInfoImpl.ReadyToRun.cs

@@ -2969,5 +2969,11 @@ private void updateEntryPointForTailCall(ref CORINFO_CONST_LOOKUP entryPoint)
 
             entryPoint = CreateConstLookupToSymbol(newEntryPoint);
         }
+
+        private int getExactClasses(CORINFO_CLASS_STRUCT_* baseType, int maxExactClasses, CORINFO_CLASS_STRUCT_** exactClsRet)
+        {
+            // Not implemented for R2R yet
+            return 0;
+        }
     }
 }

src/coreclr/tools/aot/ILCompiler.RyuJit/JitInterface/CorInfoImpl.RyuJit.cs

@@ -2154,5 +2154,43 @@ private void getFieldInfo(ref CORINFO_RESOLVED_TOKEN pResolvedToken, CORINFO_MET
             // TODO: We need to implement access checks for fields and methods.  See JitInterface.cpp in mrtjit
             //       and STS::AccessCheck::CanAccess.
         }
+
+        private int getExactClasses(CORINFO_CLASS_STRUCT_* baseType, int maxExactClasses, CORINFO_CLASS_STRUCT_** exactClsRet)
+        {
+            MetadataType type = HandleToObject(baseType) as MetadataType;
+            if (type == null)
+            {
+                return 0;
+            }
+
+            // type is already sealed, return it
+            if (_compilation.IsEffectivelySealed(type))
+            {
+                *exactClsRet = baseType;
+                return 1;
+            }
+
+            if (!type.IsInterface)
+            {
+                // TODO: handle classes
+                return 0;
+            }
+
+            TypeDesc[] implClasses = _compilation.GetImplementingClasses(type);
+            if (implClasses == null || implClasses.Length > maxExactClasses)
+            {
+                return 0;
+            }
+
+            int index = 0;
+            foreach (TypeDesc implClass in implClasses)
+            {
+                Debug.Assert(!implClass.IsCanonicalSubtype(CanonicalFormKind.Any));
+                exactClsRet[index++] = ObjectToHandle(implClass);
+            }
+
+            Debug.Assert(index <= maxExactClasses);
+            return index;
+        }
     }
 }

src/coreclr/tools/aot/jitinterface/jitinterface.h

@@ -114,6 +114,7 @@ struct JitInterfaceCallbacks
     void (* freeArray)(void * thisHandle, CorInfoExceptionClass** ppException, void* array);
     CORINFO_ARG_LIST_HANDLE (* getArgNext)(void * thisHandle, CorInfoExceptionClass** ppException, CORINFO_ARG_LIST_HANDLE args);
     CorInfoTypeWithMod (* getArgType)(void * thisHandle, CorInfoExceptionClass** ppException, CORINFO_SIG_INFO* sig, CORINFO_ARG_LIST_HANDLE args, CORINFO_CLASS_HANDLE* vcTypeRet);
+    int (* getExactClasses)(void * thisHandle, CorInfoExceptionClass** ppException, CORINFO_CLASS_HANDLE baseType, int maxExactClasses, CORINFO_CLASS_HANDLE* exactClsRet);
     CORINFO_CLASS_HANDLE (* getArgClass)(void * thisHandle, CorInfoExceptionClass** ppException, CORINFO_SIG_INFO* sig, CORINFO_ARG_LIST_HANDLE args);
     CorInfoHFAElemType (* getHFAType)(void * thisHandle, CorInfoExceptionClass** ppException, CORINFO_CLASS_HANDLE hClass);
     JITINTERFACE_HRESULT (* GetErrorHRESULT)(void * thisHandle, CorInfoExceptionClass** ppException, struct _EXCEPTION_POINTERS* pExceptionPointers);
@@ -1203,6 +1204,17 @@ class JitInterfaceWrapper : public ICorJitInfo
     return temp;
 }
 
+    virtual int getExactClasses(
+          CORINFO_CLASS_HANDLE baseType,
+          int maxExactClasses,
+          CORINFO_CLASS_HANDLE* exactClsRet)
+{
+    CorInfoExceptionClass* pException = nullptr;
+    int temp = _callbacks->getExactClasses(_thisHandle, &pException, baseType, maxExactClasses, exactClsRet);
+    if (pException != nullptr) throw pException;
+    return temp;
+}
+
     virtual CORINFO_CLASS_HANDLE getArgClass(
           CORINFO_SIG_INFO* sig,
           CORINFO_ARG_LIST_HANDLE args)

src/coreclr/tools/superpmi/superpmi-shared/lwmlist.h

@@ -51,6 +51,7 @@ LWM(GetAddrOfCaptureThreadGlobal, DWORD, DLDL)
 LWM(GetArgClass, Agnostic_GetArgClass_Key, Agnostic_GetArgClass_Value)
 LWM(GetArgNext, DWORDLONG, DWORDLONG)
 LWM(GetArgType, Agnostic_GetArgType_Key, Agnostic_GetArgType_Value)
+LWM(GetExactClasses, DLD, DLD)
 LWM(GetArrayInitializationData, DLD, DWORDLONG)
 LWM(GetArrayRank, DWORDLONG, DWORD)
 LWM(GetArrayIntrinsicID, DWORDLONG, DWORD)