Jakob's patch

src/coreclr/inc/corinfo.h

@@ -594,6 +594,7 @@ enum CorInfoHelpFunc
     CORINFO_HELP_VALIDATE_INDIRECT_CALL,    // CFG: Validate function pointer
     CORINFO_HELP_DISPATCH_INDIRECT_CALL,    // CFG: Validate and dispatch to pointer
 
+    CORINFO_HELP_ENSURE_NONHEAP,             // Ensure that the target was not in the heap.
     CORINFO_HELP_COUNT,
 };
 

src/coreclr/inc/jithelpers.h

@@ -157,6 +157,7 @@
     DYNAMICJITHELPER(CORINFO_HELP_ASSIGN_REF,   JIT_WriteBarrier,   METHOD__NIL)
     DYNAMICJITHELPER(CORINFO_HELP_CHECKED_ASSIGN_REF, JIT_CheckedWriteBarrier,METHOD__NIL)
     JITHELPER(CORINFO_HELP_ASSIGN_REF_ENSURE_NONHEAP, JIT_WriteBarrierEnsureNonHeapTarget,METHOD__NIL)
+    JITHELPER(CORINFO_HELP_ENSURE_NONHEAP, JIT_EnsureNonHeapTarget,METHOD__NIL)
 
     DYNAMICJITHELPER(CORINFO_HELP_ASSIGN_BYREF, JIT_ByRefWriteBarrier,METHOD__NIL)
     DYNAMICJITHELPER(CORINFO_HELP_BULK_WRITEBARRIER, NULL, METHOD__BUFFER__MEMCOPYGC)

src/coreclr/jit/flowgraph.cpp

@@ -2457,6 +2457,16 @@ PhaseStatus Compiler::fgAddInternal()
 
     noway_assert(!dbgHandle || !pDbgHandle);
 
+//#if DEBUG
+//    // TODO: to be enabled only under jit-stress
+//    if ((info.compRetBuffArg != BAD_VAR_NUM) && !opts.IsReadyToRun())
+//    {
+//        GenTree* retBuffAddr = gtNewLclvNode(info.compRetBuffArg, TYP_BYREF);
+//        fgNewStmtAtBeg(fgFirstBB, gtNewHelperCallNode(CORINFO_HELP_ENSURE_NONHEAP, TYP_VOID, retBuffAddr));
+//        madeChanges = true;
+//    }
+//#endif
+
     if (dbgHandle || pDbgHandle)
     {
         // Test the JustMyCode VM global state variable

src/coreclr/jit/importer.cpp

@@ -843,6 +843,17 @@ GenTree* Compiler::impStoreStruct(GenTree*         store,
             // TODO-Bug?: verify if flags matter here
             GenTreeFlags indirFlags = GTF_EMPTY;
             GenTree*     destAddr   = impGetNodeAddr(store, CHECK_SPILL_ALL, &indirFlags);
+
+            if (!impIsAddressInLocal(destAddr) && (!destAddr->OperIsScalarLocal() || (destAddr->AsLclVarCommon()->GetLclNum() != info.compRetBuffArg)))
+            {
+                unsigned tmp = lvaGrabTemp(false DEBUGARG("stack copy for value returned via return buffer"));
+                lvaSetStruct(tmp, srcCall->gtRetClsHnd, false);
+                GenTree* comma = gtNewOperNode(GT_COMMA, store->TypeGet(), gtNewStoreLclVarNode(tmp, srcCall), gtNewLclvNode(tmp, lvaGetDesc(tmp)->TypeGet()));
+                store->Data() = comma;
+                comma->AsOp()->gtOp1 = impStoreStruct(comma->gtGetOp1(), curLevel, pAfterStmt, di, block);
+                return impStoreStruct(store, curLevel, pAfterStmt, di, block);
+            }
+
             NewCallArg   newArg     = NewCallArg::Primitive(destAddr).WellKnown(wellKnownArgType);
 
             if (destAddr->OperIs(GT_LCL_ADDR))
@@ -953,6 +964,20 @@ GenTree* Compiler::impStoreStruct(GenTree*         store,
             // TODO-Bug?: verify if flags matter here
             GenTreeFlags indirFlags = GTF_EMPTY;
             GenTree*     destAddr   = impGetNodeAddr(store, CHECK_SPILL_ALL, &indirFlags);
+
+            //if (!impIsAddressInLocal(destAddr) && (!destAddr->OperIsScalarLocal() || (destAddr->AsLclVarCommon()->GetLclNum() != info.compRetBuffArg)))
+            //{
+            //    unsigned tmp = lvaGrabTemp(false DEBUGARG("stack copy for value returned via return buffer"));
+            //    lvaSetStruct(tmp, call->gtRetClsHnd, false);
+            //    destAddr = gtNewLclVarAddrNode(tmp, TYP_BYREF);
+            //    // Insert address of temp into existing call
+            //    call->gtArgs.InsertAfterThisOrFirst(this,
+            //        NewCallArg::Primitive(destAddr).WellKnown(WellKnownArg::RetBuffer));
+            //    // Now the store needs to copy from the new temp instead.
+            //    store->Data() = gtNewLclvNode(tmp, lvaGetDesc(tmp)->TypeGet());
+            //    return impStoreStruct(store, curLevel, pAfterStmt, di, block);
+            //}
+
             call->gtArgs.InsertAfterThisOrFirst(this,
                                                 NewCallArg::Primitive(destAddr).WellKnown(WellKnownArg::RetBuffer));
 

src/coreclr/jit/utils.cpp

@@ -1770,6 +1770,7 @@ void HelperCallProperties::init()
                 isNoGC = true;
                 FALLTHROUGH;
             case CORINFO_HELP_ASSIGN_REF_ENSURE_NONHEAP:
+            case CORINFO_HELP_ENSURE_NONHEAP:
             case CORINFO_HELP_BULK_WRITEBARRIER:
                 mutatesHeap = true;
                 break;

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

@@ -284,6 +284,7 @@ which is the right helper to use to allocate an object of a given type. */
         CORINFO_HELP_VALIDATE_INDIRECT_CALL,    // CFG: Validate function pointer
         CORINFO_HELP_DISPATCH_INDIRECT_CALL,    // CFG: Validate and dispatch to pointer
 
+        CORINFO_HELP_ENSURE_NONHEAP,             // Ensure that the target was not in the heap.
         CORINFO_HELP_COUNT,
     }
 }

src/coreclr/vm/gchelpers.cpp

@@ -1466,6 +1466,16 @@ extern "C" HCIMPL2_RAW(VOID, JIT_WriteBarrierEnsureNonHeapTarget, Object **dst,
 }
 HCIMPLEND_RAW
 
+extern "C" HCIMPL1_RAW(VOID, JIT_EnsureNonHeapTarget, void *dst)
+{
+    STATIC_CONTRACT_MODE_COOPERATIVE;
+    STATIC_CONTRACT_THROWS;
+    STATIC_CONTRACT_GC_NOTRIGGER;
+
+    _ASSERT(!GCHeapUtilities::GetGCHeap()->IsHeapPointer(dst));
+}
+HCIMPLEND_RAW
+
 // This function sets the card table with the granularity of 1 byte, to avoid ghost updates
 //    that could occur if multiple threads were trying to set different bits in the same card.
 

src/coreclr/vm/jitinterface.h

@@ -231,6 +231,7 @@ extern "C" FCDECL2(VOID, JIT_CheckedWriteBarrier, Object **dst, Object *ref);
 
 extern "C" FCDECL2(VOID, JIT_WriteBarrier, Object **dst, Object *ref);
 extern "C" FCDECL2(VOID, JIT_WriteBarrierEnsureNonHeapTarget, Object **dst, Object *ref);
+extern "C" FCDECL1(VOID, JIT_EnsureNonHeapTarget, void *dst);
 
 // ARM64 JIT_WriteBarrier uses special ABI and thus is not callable directly
 // Copied write barriers must be called at a different location

src/coreclr/vm/reflectioninvocation.cpp

@@ -492,6 +492,9 @@ extern "C" void QCALLTYPE RuntimeMethodHandle_InvokeMethod(
     // If an exception occurs a gc may happen but we are going to dump the stack anyway and we do
     // not need to protect anything.
 
+    // Allocate a local buffer for the return buffer if necessary
+    PVOID pLocalRetBuf = nullptr;
+
     {
     BEGINFORBIDGC();
 #ifdef _DEBUG
@@ -501,8 +504,19 @@ extern "C" void QCALLTYPE RuntimeMethodHandle_InvokeMethod(
     // Take care of any return arguments
     if (fHasRetBuffArg)
     {
-        PVOID pRetBuff = gc.retVal->GetData();
-        *((LPVOID*) (pTransitionBlock + argit.GetRetBuffArgOffset())) = pRetBuff;
+        _ASSERT(hasValueTypeReturn);
+        PTR_MethodTable pMT = retTH.GetMethodTable();
+        size_t localRetBufSize = retTH.GetSize();
+
+        // Allocate a local buffer. The invoked method will write the return value to this
+        // buffer which will be copied to gc.retVal later.
+        pLocalRetBuf = _alloca(localRetBufSize);
+        ZeroMemory(pLocalRetBuf, localRetBufSize);
+        *((LPVOID*) (pTransitionBlock + argit.GetRetBuffArgOffset())) = pLocalRetBuf;
+        if (pMT->ContainsGCPointers())
+        {
+            pValueClasses = new (_alloca(sizeof(ValueClassInfo))) ValueClassInfo(pLocalRetBuf, pMT, pValueClasses);
+        }
     }
 
     // copy args
@@ -572,6 +586,19 @@ extern "C" void QCALLTYPE RuntimeMethodHandle_InvokeMethod(
     // Call the method
     CallDescrWorkerWithHandler(&callDescrData);
 
+    if (fHasRetBuffArg)
+    {
+        // Copy the return value from the return buffer to the object
+        if (retTH.GetMethodTable()->ContainsGCPointers())
+        {
+            memmoveGCRefs(gc.retVal->GetData(), pLocalRetBuf, retTH.GetSize());
+        }
+        else
+        {
+            memcpyNoGCRefs(gc.retVal->GetData(), pLocalRetBuf, retTH.GetSize());
+        }
+    }
+
     // It is still illegal to do a GC here.  The return type might have/contain GC pointers.
     if (fConstructor)
     {