Support devirtualization for virtual methods that require runtime lookups

src/coreclr/inc/corinfo.h

@@ -1576,6 +1576,7 @@ struct CORINFO_DEVIRTUALIZATION_INFO
     // [In] arguments of resolveVirtualMethod
     //
     CORINFO_METHOD_HANDLE       virtualMethod;
+    CORINFO_METHOD_HANDLE       callerMethod;
     CORINFO_CLASS_HANDLE        objClass;
     CORINFO_CONTEXT_HANDLE      context;
     CORINFO_RESOLVED_TOKEN     *pResolvedTokenVirtualMethod;

src/coreclr/inc/jiteeversionguid.h

@@ -37,11 +37,11 @@
 
 #include <minipal/guid.h>
 
-constexpr GUID JITEEVersionIdentifier = { /* fcb1b400-696c-4425-a8a7-bb082430a217 */
-    0xfcb1b400,
-    0x696c,
-    0x4425,
-    {0xa8, 0xa7, 0xbb, 0x08, 0x24, 0x30, 0xa2, 0x17}
+constexpr GUID JITEEVersionIdentifier = { /* 01657eff-500b-4529-a977-df9eb508cd6a */
+    0x01657eff,
+    0x500b,
+    0x4529,
+    {0xa9, 0x77, 0xdf, 0x9e, 0xb5, 0x08, 0xcd, 0x6a}
   };
 
 #endif // JIT_EE_VERSIONING_GUID_H

src/coreclr/jit/compiler.h

@@ -5042,6 +5042,7 @@ class Compiler
         CORINFO_RESOLVED_TOKEN* pResolvedToken;          // Resolved token for the target method, used by R2R.
         CORINFO_RESOLVED_TOKEN* pUnboxedResolvedToken;   // Resolved token and method handle for the unboxed entry.
         CORINFO_LOOKUP*         pInstParamLookup;        // All the information needed for the instantiation parameter lookup.
+        GenTree*                runtimeLookupContext;    // Runtime context tree to use for pInstParamLookup, or nullptr.
         CORINFO_SIG_INFO*       pMethSig;                // The devirted method signature.
         bool                    objIsNonNull;            // True if the receiver is known non-null.
         bool                    hadImplicitNullCheck;    // True if the original call's null check was implicit.
@@ -7213,9 +7214,11 @@ class Compiler
                                                 CORINFO_METHOD_HANDLE dispatcherHnd);
     GenTree* getLookupTree(CORINFO_LOOKUP* pLookup,
                            GenTreeFlags    handleFlags,
-                           void*           compileTimeHandle);
+                           void*           compileTimeHandle,
+                           GenTree*        runtimeLookupContext = nullptr);
     GenTree* getRuntimeLookupTree(CORINFO_LOOKUP* pLookup,
-                                  void*           compileTimeHandle);
+                                  void*           compileTimeHandle,
+                                  GenTree*        runtimeLookupContext = nullptr);
     GenTree* getVirtMethodPointerTree(GenTree*                thisPtr,
                                       CORINFO_RESOLVED_TOKEN* pResolvedToken,
                                       CORINFO_CALL_INFO*      pCallInfo);
@@ -8108,7 +8111,8 @@ class Compiler
                                          bool                    isInterface,
                                          CORINFO_METHOD_HANDLE   baseMethod,
                                          CORINFO_CLASS_HANDLE    baseClass,
-                                         CORINFO_CONTEXT_HANDLE* pContextHandle);
+                                         CORINFO_CONTEXT_HANDLE* pContextHandle,
+                                         CORINFO_RESOLVED_TOKEN* pResolvedToken);
 
     bool isCompatibleMethodGDV(GenTreeCall* call, CORINFO_METHOD_HANDLE gdvTarget);
 

src/coreclr/jit/fginline.cpp

@@ -608,16 +608,17 @@ class SubstitutePlaceholdersAndDevirtualizeWalker : public GenTreeVisitor<Substi
                 }
 #endif // DEBUG
 
-                CORINFO_METHOD_HANDLE  method                 = call->gtLateDevirtualizationInfo->methodHnd;
-                CORINFO_CONTEXT_HANDLE context                = call->gtLateDevirtualizationInfo->exactContextHnd;
-                InlineContext*         inlinersContext        = call->gtInlineContext;
-                unsigned               methodFlags            = 0;
-                const bool             isLateDevirtualization = true;
-                const bool             explicitTailCall       = call->IsTailPrefixedCall();
+                CORINFO_METHOD_HANDLE   method                 = call->gtLateDevirtualizationInfo->methodHnd;
+                CORINFO_CONTEXT_HANDLE  context                = call->gtLateDevirtualizationInfo->exactContextHnd;
+                InlineContext*          inlinersContext        = call->gtInlineContext;
+                CORINFO_RESOLVED_TOKEN* pResolvedToken         = &call->gtLateDevirtualizationInfo->resolvedToken;
+                unsigned                methodFlags            = 0;
+                const bool              isLateDevirtualization = true;
+                const bool              explicitTailCall       = call->IsTailPrefixedCall();
 
                 CORINFO_CONTEXT_HANDLE contextInput = context;
                 context                             = nullptr;
-                m_compiler->impDevirtualizeCall(call, nullptr, &method, &methodFlags, &contextInput, &context,
+                m_compiler->impDevirtualizeCall(call, pResolvedToken, &method, &methodFlags, &contextInput, &context,
                                                 isLateDevirtualization, explicitTailCall);
 
                 if (!call->IsDevirtualizationCandidate(m_compiler))

src/coreclr/jit/importer.cpp

@@ -1628,10 +1628,17 @@ GenTree* Compiler::impRuntimeLookupToTree(CORINFO_LOOKUP* pLookup, void* compile
         assert(pRuntimeLookup->indirections != 0);
         GenTreeCall* helperCall = gtNewRuntimeLookupHelperCallNode(pRuntimeLookup, ctxTree, compileTimeHandle);
 
-        // Spilling it to a temp improves CQ (mainly in Tier0)
-        unsigned callLclNum = lvaGrabTemp(true DEBUGARG("spilling helperCall"));
-        impStoreToTemp(callLclNum, helperCall, CHECK_SPILL_NONE);
-        return gtNewLclvNode(callLclNum, helperCall->TypeGet());
+        if (opts.OptimizationEnabled())
+        {
+            return helperCall;
+        }
+        else
+        {
+            // Spilling it to a temp improves CQ (mainly in Tier0)
+            unsigned callLclNum = lvaGrabTemp(true DEBUGARG("spilling helperCall"));
+            impStoreToTemp(callLclNum, helperCall, CHECK_SPILL_NONE);
+            return gtNewLclvNode(callLclNum, helperCall->TypeGet());
+        }
     }
 
     // Size-check is not expected without testForNull

src/coreclr/jit/importercalls.cpp

@@ -1086,7 +1086,7 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
         else if (call->AsCall()->IsDelegateInvoke())
         {
             considerGuardedDevirtualization(call->AsCall(), rawILOffset, false, call->AsCall()->gtCallMethHnd,
-                                            NO_CLASS_HANDLE, nullptr);
+                                            NO_CLASS_HANDLE, nullptr, pResolvedToken);
         }
     }
 
@@ -1367,6 +1367,7 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
             info->methodHnd                            = callInfo->hMethod;
             info->exactContextHnd                      = exactContextHnd;
             info->ilLocation                           = impCurStmtDI.GetLocation();
+            info->resolvedToken                        = *pResolvedToken;
             call->AsCall()->gtLateDevirtualizationInfo = info;
         }
     }
@@ -8013,7 +8014,8 @@ void Compiler::considerGuardedDevirtualization(GenTreeCall*            call,
                                                bool                    isInterface,
                                                CORINFO_METHOD_HANDLE   baseMethod,
                                                CORINFO_CLASS_HANDLE    baseClass,
-                                               CORINFO_CONTEXT_HANDLE* pContextHandle)
+                                               CORINFO_CONTEXT_HANDLE* pContextHandle,
+                                               CORINFO_RESOLVED_TOKEN* pResolvedToken)
 {
     JITDUMP("Considering guarded devirtualization at IL offset %u (0x%x)\n", ilOffset, ilOffset);
 
@@ -8092,9 +8094,10 @@ void Compiler::considerGuardedDevirtualization(GenTreeCall*            call,
                 //
                 CORINFO_DEVIRTUALIZATION_INFO dvInfo;
                 dvInfo.virtualMethod               = baseMethod;
+                dvInfo.callerMethod                = call->gtInlineContext->GetCallee();
                 dvInfo.objClass                    = exactCls;
                 dvInfo.context                     = originalContext;
-                dvInfo.pResolvedTokenVirtualMethod = nullptr;
+                dvInfo.pResolvedTokenVirtualMethod = pResolvedToken;
 
                 JITDUMP("GDV exact: resolveVirtualMethod (method %p class %p context %p)\n", dvInfo.virtualMethod,
                         dvInfo.objClass, dvInfo.context);
@@ -8177,6 +8180,7 @@ void Compiler::considerGuardedDevirtualization(GenTreeCall*            call,
             // Figure out which method will be called.
             //
             dvInfo.virtualMethod               = baseMethod;
+            dvInfo.callerMethod                = call->gtInlineContext->GetCallee();
             dvInfo.objClass                    = likelyClass;
             dvInfo.context                     = originalContext;
             dvInfo.pResolvedTokenVirtualMethod = nullptr;
@@ -9229,7 +9233,8 @@ void Compiler::impDevirtualizeCall(GenTreeCall*            call,
             return;
         }
 
-        considerGuardedDevirtualization(call, ilOffset, isInterface, baseMethod, baseClass, pContextHandle);
+        considerGuardedDevirtualization(call, ilOffset, isInterface, baseMethod, baseClass, pContextHandle,
+                                        pResolvedToken);
 
         return;
     }
@@ -9273,7 +9278,8 @@ void Compiler::impDevirtualizeCall(GenTreeCall*            call,
             return;
         }
 
-        considerGuardedDevirtualization(call, ilOffset, isInterface, baseMethod, baseClass, pContextHandle);
+        considerGuardedDevirtualization(call, ilOffset, isInterface, baseMethod, baseClass, pContextHandle,
+                                        pResolvedToken);
         return;
     }
 
@@ -9285,11 +9291,9 @@ void Compiler::impDevirtualizeCall(GenTreeCall*            call,
         JITDUMP("--- base class is interface\n");
     }
 
-    // Fetch the method that would be called based on the declared type of 'this',
-    // and prepare to fetch the method attributes.
-    //
     CORINFO_DEVIRTUALIZATION_INFO dvInfo;
     dvInfo.virtualMethod               = baseMethod;
+    dvInfo.callerMethod                = call->gtInlineContext->GetCallee();
     dvInfo.objClass                    = objClass;
     dvInfo.context                     = *pContextHandle;
     dvInfo.detail                      = CORINFO_DEVIRTUALIZATION_UNKNOWN;
@@ -9300,6 +9304,56 @@ void Compiler::impDevirtualizeCall(GenTreeCall*            call,
 
     info.compCompHnd->resolveVirtualMethod(&dvInfo);
 
+    GenTree* runtimeLookupContext = nullptr;
+
+    if (isLateDevirtualization && dvInfo.instParamLookup.lookupKind.needsRuntimeLookup)
+    {
+        // Late devirtualization may revisit a call that was imported in an inlinee.
+        // If token context is not the current root context, runtime lookups can be rooted in the wrong generic context.
+        //
+        CORINFO_CONTEXT_HANDLE tokenContext       = pResolvedToken->tokenContext;
+        const SIZE_T           tokenContextHandle = (SIZE_T)tokenContext & ~CORINFO_CONTEXTFLAGS_MASK;
+
+        const bool isMethodContext = ((SIZE_T)tokenContext & CORINFO_CONTEXTFLAGS_MASK) == CORINFO_CONTEXTFLAGS_METHOD;
+        const bool isCurrentContext =
+            (tokenContext == METHOD_BEING_COMPILED_CONTEXT()) ||
+            (isMethodContext ? ((CORINFO_METHOD_HANDLE)tokenContextHandle == info.compMethodHnd)
+                             : ((CORINFO_CLASS_HANDLE)tokenContextHandle == info.compClassHnd));
+
+        // If we don't have the right context, try recover it from the tree.
+        //
+        if (!isCurrentContext)
+        {
+            CallArg* runtimeMethodHandleArg = nullptr;
+            if ((call->gtControlExpr != nullptr) && call->gtControlExpr->OperIs(GT_CALL))
+            {
+                runtimeMethodHandleArg =
+                    call->gtControlExpr->AsCall()->gtArgs.FindWellKnownArg(WellKnownArg::RuntimeMethodHandle);
+            }
+
+            if (runtimeMethodHandleArg != nullptr && runtimeMethodHandleArg->GetNode()->OperIs(GT_RUNTIMELOOKUP))
+            {
+                if (runtimeMethodHandleArg->GetNode()->AsRuntimeLookup()->Lookup()->OperIs(GT_CALL))
+                {
+                    GenTreeCall* const helperCall =
+                        runtimeMethodHandleArg->GetNode()->AsRuntimeLookup()->Lookup()->AsCall();
+
+                    if (helperCall->IsHelperCall(CORINFO_HELP_RUNTIMEHANDLE_METHOD) ||
+                        helperCall->IsHelperCall(CORINFO_HELP_RUNTIMEHANDLE_CLASS))
+                    {
+                        runtimeLookupContext = helperCall->gtArgs.GetArgByIndex(0)->GetNode();
+                    }
+                }
+            }
+
+            if (runtimeLookupContext == nullptr)
+            {
+                JITDUMP("Late devirt needs a runtime lookup context cannot be figured out. Bail out.\n");
+                return;
+            }
+        }
+    }
+
     CORINFO_METHOD_HANDLE   derivedMethod         = dvInfo.devirtualizedMethod;
     CORINFO_CONTEXT_HANDLE  exactContext          = dvInfo.tokenLookupContext;
     CORINFO_RESOLVED_TOKEN* pDerivedResolvedToken = &dvInfo.resolvedTokenDevirtualizedMethod;
@@ -9395,7 +9449,8 @@ void Compiler::impDevirtualizeCall(GenTreeCall*            call,
             return;
         }
 
-        considerGuardedDevirtualization(call, ilOffset, isInterface, baseMethod, objClass, pContextHandle);
+        considerGuardedDevirtualization(call, ilOffset, isInterface, baseMethod, objClass, pContextHandle,
+                                        pResolvedToken);
         return;
     }
 
@@ -9410,6 +9465,7 @@ void Compiler::impDevirtualizeCall(GenTreeCall*            call,
     dcInfo.pInstParamLookup      = &dvInfo.instParamLookup;
     dcInfo.pResolvedToken        = pDerivedResolvedToken;
     dcInfo.pUnboxedResolvedToken = &dvInfo.resolvedTokenDevirtualizedUnboxedMethod;
+    dcInfo.runtimeLookupContext  = runtimeLookupContext;
     dcInfo.pMethSig              = &derivedSig;
     dcInfo.objIsNonNull          = objIsNonNull;
     dcInfo.hadImplicitNullCheck  = true;
@@ -9628,7 +9684,8 @@ void Compiler::impTransformDevirtualizedCall(GenTreeCall*            call,
                         CORINFO_METHOD_HANDLE exactMethodHandle =
                             (CORINFO_METHOD_HANDLE)((SIZE_T)dcInfo->tokenLookupContext & ~CORINFO_CONTEXTFLAGS_MASK);
 
-                        instParam = getLookupTree(dcInfo->pInstParamLookup, GTF_ICON_METHOD_HDL, exactMethodHandle);
+                        instParam = getLookupTree(dcInfo->pInstParamLookup, GTF_ICON_METHOD_HDL, exactMethodHandle,
+                                                  dcInfo->runtimeLookupContext);
                         JITDUMP("revising call to invoke unboxed entry with additional method desc arg\n");
                     }
                     else
@@ -9710,7 +9767,8 @@ void Compiler::impTransformDevirtualizedCall(GenTreeCall*            call,
                 CORINFO_METHOD_HANDLE exactMethodHandle =
                     (CORINFO_METHOD_HANDLE)((SIZE_T)dcInfo->tokenLookupContext & ~CORINFO_CONTEXTFLAGS_MASK);
 
-                instParam = getLookupTree(dcInfo->pInstParamLookup, GTF_ICON_METHOD_HDL, exactMethodHandle);
+                instParam = getLookupTree(dcInfo->pInstParamLookup, GTF_ICON_METHOD_HDL, exactMethodHandle,
+                                          dcInfo->runtimeLookupContext);
             }
             else
             {
@@ -9719,7 +9777,8 @@ void Compiler::impTransformDevirtualizedCall(GenTreeCall*            call,
                 CORINFO_CLASS_HANDLE exactClassHandle =
                     (CORINFO_CLASS_HANDLE)((SIZE_T)dcInfo->tokenLookupContext & ~CORINFO_CONTEXTFLAGS_MASK);
 
-                instParam = getLookupTree(dcInfo->pInstParamLookup, GTF_ICON_CLASS_HDL, exactClassHandle);
+                instParam = getLookupTree(dcInfo->pInstParamLookup, GTF_ICON_CLASS_HDL, exactClassHandle,
+                                          dcInfo->runtimeLookupContext);
             }
         }
     }

src/coreclr/jit/indirectcalltransformer.cpp

@@ -987,6 +987,7 @@ class IndirectCallTransformer
             dcInfo.pUnboxedResolvedToken =
                 (clsHnd != NO_CLASS_HANDLE) ? &inlineInfo->guardedMethodUnboxedResolvedToken : nullptr;
 
+            dcInfo.runtimeLookupContext = nullptr;
             dcInfo.objIsNonNull         = objIsNonNull;
             dcInfo.hadImplicitNullCheck = m_origCall->IsVirtual();
             dcInfo.isDelegateCall       = m_origCall->IsDelegateInvoke();

src/coreclr/jit/inline.h

@@ -634,6 +634,7 @@ struct LateDevirtualizationInfo
     CORINFO_METHOD_HANDLE  methodHnd;
     CORINFO_CONTEXT_HANDLE exactContextHnd;
     ILLocation             ilLocation;
+    CORINFO_RESOLVED_TOKEN resolvedToken;
 };
 
 // InlArgInfo describes inline candidate argument properties.

src/coreclr/jit/morph.cpp

@@ -5512,11 +5512,15 @@ GenTree* Compiler::fgCreateCallDispatcherAndGetResult(GenTreeCall*          orig
 //    pLookup - the lookup to get the tree for
 //    handleFlags - flags to set on the result node
 //    compileTimeHandle - compile-time handle corresponding to the lookup
+//    runtimeLookupContext - context tree to use for a runtime lookup, or nullptr to use the current context
 //
 // Return Value:
 //    A node representing the lookup tree
 //
-GenTree* Compiler::getLookupTree(CORINFO_LOOKUP* pLookup, GenTreeFlags handleFlags, void* compileTimeHandle)
+GenTree* Compiler::getLookupTree(CORINFO_LOOKUP* pLookup,
+                                 GenTreeFlags    handleFlags,
+                                 void*           compileTimeHandle,
+                                 GenTree*        runtimeLookupContext)
 {
     if (!pLookup->lookupKind.needsRuntimeLookup)
     {
@@ -5539,7 +5543,7 @@ GenTree* Compiler::getLookupTree(CORINFO_LOOKUP* pLookup, GenTreeFlags handleFla
         return gtNewIconEmbHndNode(handle, pIndirection, handleFlags, compileTimeHandle);
     }
 
-    return getRuntimeLookupTree(pLookup, compileTimeHandle);
+    return getRuntimeLookupTree(pLookup, compileTimeHandle, runtimeLookupContext);
 }
 
 //------------------------------------------------------------------------
@@ -5548,25 +5552,27 @@ GenTree* Compiler::getLookupTree(CORINFO_LOOKUP* pLookup, GenTreeFlags handleFla
 // Arguments:
 //    pLookup - the lookup to get the tree for
 //    compileTimeHandle - compile-time handle corresponding to the lookup
+//    runtimeLookupContext - context tree to use for the lookup, or nullptr to use the current context
 //
 // Return Value:
 //    A node representing the runtime lookup tree
 //
-GenTree* Compiler::getRuntimeLookupTree(CORINFO_LOOKUP* pLookup, void* compileTimeHandle)
+GenTree* Compiler::getRuntimeLookupTree(CORINFO_LOOKUP* pLookup, void* compileTimeHandle, GenTree* runtimeLookupContext)
 {
     CORINFO_RUNTIME_LOOKUP* pRuntimeLookup = &pLookup->runtimeLookup;
+    GenTree*                contextTree    = runtimeLookupContext != nullptr
+                                                 ? runtimeLookupContext
+                                                 : getRuntimeContextTree(pLookup->lookupKind.runtimeLookupKind);
 
     // If pRuntimeLookup->indirections is equal to CORINFO_USEHELPER, it specifies that a run-time helper should be
     // used; otherwise, it specifies the number of indirections via pRuntimeLookup->offsets array.
     if ((pRuntimeLookup->indirections == CORINFO_USEHELPER) || (pRuntimeLookup->indirections == CORINFO_USENULL) ||
         pRuntimeLookup->testForNull)
     {
-        return gtNewRuntimeLookupHelperCallNode(pRuntimeLookup,
-                                                getRuntimeContextTree(pLookup->lookupKind.runtimeLookupKind),
-                                                compileTimeHandle);
+        return gtNewRuntimeLookupHelperCallNode(pRuntimeLookup, contextTree, compileTimeHandle);
     }
 
-    GenTree* result = getRuntimeContextTree(pLookup->lookupKind.runtimeLookupKind);
+    GenTree* result = contextTree;
 
     ArrayStack<GenTree*> stmts(getAllocator(CMK_ArrayStack));
 

src/coreclr/tools/Common/Internal/Runtime/ReadyToRunConstants.cs

@@ -124,6 +124,7 @@ public enum DictionaryEntryKind
         DispatchStubAddrSlot = 5,
         FieldDescSlot = 6,
         DeclaringTypeHandleSlot = 7,
+        DevirtualizedMethodDescSlot = 8,
     }
 
     public enum ReadyToRunFixupKind