Return type from exact method context for call returns

src/coreclr/jit/fginline.cpp

@@ -562,9 +562,9 @@ class SubstitutePlaceholdersAndDevirtualizeWalker : public GenTreeVisitor<Substi
                 const bool             isLateDevirtualization = true;
                 const bool             explicitTailCall       = call->IsTailPrefixedCall();
 
-                if ((call->gtCallMoreFlags & GTF_CALL_M_HAS_LATE_DEVIRT_INFO) != 0)
+                if ((call->gtCallMoreFlags & GTF_CALL_M_HAS_EXACT_CONTEXT_INFO) != 0)
                 {
-                    context = call->gtLateDevirtualizationInfo->exactContextHnd;
+                    context = call->gtExactContextInfo->exactContextHnd;
                     // Note: we might call this multiple times for the same trees.
                     // If the devirtualization below succeeds, the call becomes
                     // non-virtual and we won't get here again. If it does not

src/coreclr/jit/gentree.cpp

@@ -19224,13 +19224,25 @@ CORINFO_CLASS_HANDLE Compiler::gtGetClassHandle(GenTree* tree, bool* pIsExact, b
             }
             else if (call->gtCallType == CT_USER_FUNC)
             {
-                // For user calls, we can fetch the approximate return
-                // type info from the method handle. Unfortunately
-                // we've lost the exact context, so this is the best
-                // we can do for now.
                 CORINFO_METHOD_HANDLE method     = call->gtCallMethHnd;
                 CORINFO_CLASS_HANDLE  exactClass = nullptr;
-                CORINFO_SIG_INFO      sig;
+                if ((call->gtCallMoreFlags & GTF_CALL_M_HAS_EXACT_CONTEXT_INFO) != 0)
+                {
+                    CORINFO_CONTEXT_HANDLE context = call->gtExactContextInfo->exactContextHnd;
+
+                    if (context != nullptr)
+                    {
+                        if (((SIZE_T)context & CORINFO_CONTEXTFLAGS_MASK) == CORINFO_CONTEXTFLAGS_CLASS)
+                        {
+                            exactClass = CORINFO_CLASS_HANDLE((SIZE_T)context & ~CORINFO_CONTEXTFLAGS_MASK);
+                        }
+                        else
+                        {
+                            method = CORINFO_METHOD_HANDLE((SIZE_T)context & ~CORINFO_CONTEXTFLAGS_MASK);
+                        }
+                    }
+                }
+                CORINFO_SIG_INFO sig;
                 eeGetMethodSig(method, &sig, exactClass);
                 if (sig.retType == CORINFO_TYPE_VOID)
                 {

src/coreclr/jit/gentree.h

@@ -174,7 +174,7 @@ struct BasicBlock;
 enum BasicBlockFlags : uint64_t;
 struct InlineCandidateInfo;
 struct HandleHistogramProfileCandidateInfo;
-struct LateDevirtualizationInfo;
+struct ExactContextInfo;
 
 typedef unsigned short AssertionIndex;
 
@@ -4169,7 +4169,7 @@ enum GenTreeCallFlags : unsigned int
     GTF_CALL_M_ALLOC_SIDE_EFFECTS      = 0x00100000, // this is a call to an allocator with side effects
     GTF_CALL_M_SUPPRESS_GC_TRANSITION  = 0x00200000, // suppress the GC transition (i.e. during a pinvoke) but a separate GC safe point is required.
     GTF_CALL_M_EXPANDED_EARLY          = 0x00800000, // the Virtual Call target address is expanded and placed in gtControlExpr in Morph rather than in Lower
-    GTF_CALL_M_HAS_LATE_DEVIRT_INFO    = 0x01000000, // this call has late devirtualzation info
+    GTF_CALL_M_HAS_EXACT_CONTEXT_INFO  = 0x01000000, // this call has exact context info
     GTF_CALL_M_LDVIRTFTN_INTERFACE     = 0x02000000, // ldvirtftn on an interface type
     GTF_CALL_M_CAST_CAN_BE_EXPANDED    = 0x04000000, // this cast (helper call) can be expanded if it's profitable. To be removed.
     GTF_CALL_M_CAST_OBJ_NONNULL        = 0x08000000, // if we expand this specific cast we don't need to check the input object for null
@@ -5691,7 +5691,7 @@ struct GenTreeCall final : public GenTree
         jitstd::vector<InlineCandidateInfo*>* gtInlineCandidateInfoList;
 
         HandleHistogramProfileCandidateInfo* gtHandleHistogramProfileCandidateInfo;
-        LateDevirtualizationInfo*            gtLateDevirtualizationInfo;
+        ExactContextInfo*                    gtExactContextInfo;
         CORINFO_GENERIC_HANDLE compileTimeHelperArgumentHandle; // Used to track type handle argument of dynamic helpers
         void*                  gtDirectCallAddress; // Used to pass direct call address between lower and codegen
     };

src/coreclr/jit/importercalls.cpp

@@ -1398,20 +1398,20 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
             }
             else
             {
-                // If the call is virtual, and has a generics context, and is not going to have a class probe,
-                // record the context for possible use during late devirt.
+                // If the call has a generic context, and is not going to have a class probe,
+                // record the context for possible future use in devirtualization.
                 //
                 // If we ever want to devirt at Tier0, and/or see issues where OSR methods under PGO lose
                 // important devirtualizations, we'll want to allow both a class probe and a captured context.
                 //
-                if (origCall->IsVirtual() && (origCall->gtCallType != CT_INDIRECT) && (exactContextHnd != nullptr) &&
+                if ((origCall->gtCallType != CT_INDIRECT) && (exactContextHnd != nullptr) &&
                     (origCall->gtHandleHistogramProfileCandidateInfo == nullptr))
                 {
                     JITDUMP("\nSaving context %p for call [%06u]\n", dspPtr(exactContextHnd), dspTreeID(origCall));
-                    origCall->gtCallMoreFlags |= GTF_CALL_M_HAS_LATE_DEVIRT_INFO;
-                    LateDevirtualizationInfo* const info = new (this, CMK_Inlining) LateDevirtualizationInfo;
-                    info->exactContextHnd                = exactContextHnd;
-                    origCall->gtLateDevirtualizationInfo = info;
+                    origCall->gtCallMoreFlags |= GTF_CALL_M_HAS_EXACT_CONTEXT_INFO;
+                    ExactContextInfo* const info = new (this, CMK_Inlining) ExactContextInfo;
+                    info->exactContextHnd        = exactContextHnd;
+                    origCall->gtExactContextInfo = info;
                 }
 
                 if (isFatPointerCandidate)
@@ -8007,7 +8007,7 @@ void Compiler::impDevirtualizeCall(GenTreeCall*            call,
     // it's a union field used for other things by virtual
     // stubs)
     call->ClearInlineInfo();
-    call->gtCallMoreFlags &= ~GTF_CALL_M_HAS_LATE_DEVIRT_INFO;
+    call->gtCallMoreFlags &= ~GTF_CALL_M_HAS_EXACT_CONTEXT_INFO;
 
 #if defined(DEBUG)
     if (verbose)

src/coreclr/jit/inline.h

@@ -621,11 +621,11 @@ struct InlineCandidateInfo : public HandleHistogramProfileCandidateInfo
     InlineContext*         inlinersContext;
 };
 
-// LateDevirtualizationInfo
+// ExactContextInfo
 //
-// Used to fill in missing contexts during late devirtualization.
+// Used to fill in missing contexts during devirtualization.
 //
-struct LateDevirtualizationInfo
+struct ExactContextInfo
 {
     CORINFO_CONTEXT_HANDLE exactContextHnd;
 };

src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs

@@ -1196,8 +1196,10 @@ private void getMethodSig(CORINFO_METHOD_STRUCT_* ftn, CORINFO_SIG_INFO* sig, CO
                     }
                     else
                     {
-                        Debug.Assert(type.HasSameTypeDefinition(method.OwningType));
+                        Debug.Assert(type.HasSameTypeDefinition(method.OwningType),
+                            $"{type.GetDisplayName()} is a different type definition from the parent of {method.GetDisplayName()}");
                         Instantiation methodInst = method.Instantiation;
+                        Debug.Assert(type is InstantiatedType, $"{type.GetDisplayName()} is not a {nameof(InstantiatedType)}");
                         method = _compilation.TypeSystemContext.GetMethodForInstantiatedType(method.GetTypicalMethodDefinition(), (InstantiatedType)type);
                         if (methodInst.Length > 0)
                         {

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

@@ -1462,7 +1462,14 @@ private void getCallInfo(ref CORINFO_RESOLVED_TOKEN pResolvedToken, CORINFO_RESO
                     // the method with a method the intrinsic expands into. If it's not the special intrinsic,
                     // method stays unchanged.
                     var methodIL = (MethodIL)HandleToObject((void*)pResolvedToken.tokenScope);
-                    targetMethod = _compilation.ExpandIntrinsicForCallsite(targetMethod, methodIL.OwningMethod);
+                    MethodDesc callsiteMethod = _compilation.ExpandIntrinsicForCallsite(targetMethod, methodIL.OwningMethod);
+                    if (targetMethod != callsiteMethod)
+                    {
+                        Debug.Assert(!pResult->exactContextNeedsRuntimeLookup);
+                        Debug.Assert(!callsiteMethod.HasInstantiation && !targetMethod.HasInstantiation);
+                        pResult->contextHandle = contextFromType(callsiteMethod.OwningType);
+                        targetMethod = callsiteMethod;
+                    }
 
                     // For multidim array Address method, we pretend the method requires a hidden instantiation argument
                     // (even though it doesn't need one). We'll actually swap the method out for a differnt one with