Handle async continuation in TransitionFrame and ArgIterator

docs/design/coreclr/botr/clr-abi.md

@@ -116,12 +116,12 @@ To return `Continuation` we use a volatile/calee-trash register that cannot be u
 | risc-v  | a2  |
 
 ### Passing `Continuation` argument
-The `Continuation` parameter is passed at the same position as generic instantiation parameter or immediately after, if both present.
+The `Continuation` parameter is passed at the same position as generic instantiation parameter or immediately after, if both present. For x86 the argument order is reversed.
 
 ```
 call(["this" pointer] [return buffer pointer] [generics context] [continuation] [userargs])   // not x86
 
-call(["this" pointer] [return buffer pointer] [userargs] [generics context] [continuation])   // x86
+call(["this" pointer] [return buffer pointer] [userargs] [continuation] [generics context])   // x86
 ```
 
 ## AMD64-only: by-value value types

src/coreclr/jit/importercalls.cpp

@@ -899,18 +899,18 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
         }
         else
         {
-            if (instParam != nullptr)
-            {
-                call->AsCall()->gtArgs.PushBack(this,
-                                                NewCallArg::Primitive(instParam).WellKnown(WellKnownArg::InstParam));
-            }
-
             if (call->AsCall()->IsAsync())
             {
                 call->AsCall()->gtArgs.PushBack(this, NewCallArg::Primitive(gtNewNull(), TYP_REF)
                                                           .WellKnown(WellKnownArg::AsyncContinuation));
             }
 
+            if (instParam != nullptr)
+            {
+                call->AsCall()->gtArgs.PushBack(this,
+                                                NewCallArg::Primitive(instParam).WellKnown(WellKnownArg::InstParam));
+            }
+
             if (varArgsCookie != nullptr)
             {
                 call->AsCall()->gtArgs.PushBack(this, NewCallArg::Primitive(varArgsCookie)

src/coreclr/jit/lclvars.cpp

@@ -387,12 +387,12 @@ void Compiler::lvaInitArgs(bool hasRetBuffArg)
     //-------------------------------------------------------------------------
     lvaInitUserArgs(&varNum, numUserArgsToSkip, numUserArgs);
 #if !USER_ARGS_COME_LAST
+    lvaInitAsyncContinuation(&varNum);
+
     //@GENERICS: final instantiation-info argument for shared generic methods
     // and shared generic struct instance methods
     lvaInitGenericsCtxt(&varNum);
 
-    lvaInitAsyncContinuation(&varNum);
-
     /* If the method is varargs, process the varargs cookie */
     lvaInitVarArgsHandle(&varNum);
 #endif

src/coreclr/vm/callingconvention.h

@@ -641,6 +641,7 @@ class ArgIteratorTemplate : public ARGITERATOR_BASE
     int GetRetBuffArgOffset();
     int GetVASigCookieOffset();
     int GetParamTypeArgOffset();
+    int GetAsyncContinuationArgOffset();
 
     //------------------------------------------------------------
     // Each time this is called, this returns a byte offset of the next
@@ -1000,21 +1001,26 @@ class ArgIteratorTemplate : public ARGITERATOR_BASE
 #endif
 
     enum {
-        ITERATION_STARTED               = 0x0001,   // Started iterating over arguments
-        SIZE_OF_ARG_STACK_COMPUTED      = 0x0002,
-        RETURN_FLAGS_COMPUTED           = 0x0004,
-        RETURN_HAS_RET_BUFFER           = 0x0008,   // Cached value of HasRetBuffArg
+        ITERATION_STARTED                 = 0x0001,   // Started iterating over arguments
+        SIZE_OF_ARG_STACK_COMPUTED        = 0x0002,
+        RETURN_FLAGS_COMPUTED             = 0x0004,
+        RETURN_HAS_RET_BUFFER             = 0x0008,   // Cached value of HasRetBuffArg
 
 #ifdef TARGET_X86
-        PARAM_TYPE_REGISTER_MASK        = 0x0030,
-        PARAM_TYPE_REGISTER_STACK       = 0x0010,
-        PARAM_TYPE_REGISTER_ECX         = 0x0020,
-        PARAM_TYPE_REGISTER_EDX         = 0x0030,
+        PARAM_TYPE_REGISTER_MASK          = 0x0030,
+        PARAM_TYPE_REGISTER_STACK         = 0x0010,
+        PARAM_TYPE_REGISTER_ECX           = 0x0020,
+        PARAM_TYPE_REGISTER_EDX           = 0x0030,
+
+        ASYNC_CONTINUATION_REGISTER_MASK  = 0x00C0,
+        ASYNC_CONTINUATION_REGISTER_STACK = 0x0040,
+        ASYNC_CONTINUATION_REGISTER_ECX   = 0x0080,
+        ASYNC_CONTINUATION_REGISTER_EDX   = 0x00C0,
 #endif
 
-        METHOD_INVOKE_NEEDS_ACTIVATION  = 0x0040,   // Flag used by ArgIteratorForMethodInvoke
+        METHOD_INVOKE_NEEDS_ACTIVATION    = 0x0100,   // Flag used by ArgIteratorForMethodInvoke
 
-        RETURN_FP_SIZE_SHIFT            = 8,        // The rest of the flags is cached value of GetFPReturnSize
+        RETURN_FP_SIZE_SHIFT              = 10,       // The rest of the flags is cached value of GetFPReturnSize
     };
 
     void ComputeReturnFlags();
@@ -1158,6 +1164,68 @@ int ArgIteratorTemplate<ARGITERATOR_BASE>::GetParamTypeArgOffset()
 #endif
 }
 
+template<class ARGITERATOR_BASE>
+int ArgIteratorTemplate<ARGITERATOR_BASE>::GetAsyncContinuationArgOffset()
+{
+    CONTRACTL
+    {
+        INSTANCE_CHECK;
+        if (FORBIDGC_LOADER_USE_ENABLED()) NOTHROW; else THROWS;
+        if (FORBIDGC_LOADER_USE_ENABLED()) GC_NOTRIGGER; else GC_TRIGGERS;
+        if (FORBIDGC_LOADER_USE_ENABLED()) FORBID_FAULT; else { INJECT_FAULT(COMPlusThrowOM()); }
+        MODE_ANY;
+    }
+    CONTRACTL_END
+
+    _ASSERTE(this->HasAsyncContinuation());
+
+#ifdef TARGET_X86
+    // x86 is special as always
+    if (!(m_dwFlags & SIZE_OF_ARG_STACK_COMPUTED))
+        ForceSigWalk();
+
+    switch (m_dwFlags & ASYNC_CONTINUATION_REGISTER_MASK)
+    {
+    case ASYNC_CONTINUATION_REGISTER_ECX:
+        return TransitionBlock::GetOffsetOfArgumentRegisters() + offsetof(ArgumentRegisters, ECX);
+    case ASYNC_CONTINUATION_REGISTER_EDX:
+        return TransitionBlock::GetOffsetOfArgumentRegisters() + offsetof(ArgumentRegisters, EDX);
+    default:
+        break;
+    }
+
+    // If the async continuation is a stack arg, then it comes last unless
+    // there also is a param type arg on the stack, in which case it comes
+    // before it.
+    if (this->HasParamType() && (m_dwFlags & PARAM_TYPE_REGISTER_MASK) == PARAM_TYPE_REGISTER_STACK)
+    {
+        return sizeof(TransitionBlock) + sizeof(void*);
+    }
+
+    return sizeof(TransitionBlock);
+#else
+    // The hidden arg is after this, retbuf and param type arguments by default.
+    int ret = TransitionBlock::GetOffsetOfArgumentRegisters();
+
+    if (this->HasThis())
+    {
+        ret += TARGET_POINTER_SIZE;
+    }
+
+    if (this->HasRetBuffArg() && IsRetBuffPassedAsFirstArg())
+    {
+        ret += TARGET_POINTER_SIZE;
+    }
+
+    if (this->HasParamType())
+    {
+        ret += TARGET_POINTER_SIZE;
+    }
+
+    return ret;
+#endif
+}
+
 // To avoid corner case bugs, limit maximum size of the arguments with sufficient margin
 #define MAX_ARG_SIZE 0xFFFFFF
 
@@ -1190,6 +1258,11 @@ int ArgIteratorTemplate<ARGITERATOR_BASE>::GetNextOffset()
         {
             numRegistersUsed++;
         }
+
+        if (this->HasAsyncContinuation())
+        {
+            numRegistersUsed++;
+        }
 #endif
 
 #ifdef TARGET_X86
@@ -2002,6 +2075,23 @@ void ArgIteratorTemplate<ARGITERATOR_BASE>::ForceSigWalk()
         }
     }
 
+    if (this->HasAsyncContinuation())
+    {
+        DWORD asyncContFlags = 0;
+        if (numRegistersUsed < NUM_ARGUMENT_REGISTERS)
+        {
+            numRegistersUsed++;
+            asyncContFlags = (numRegistersUsed == 1) ?
+                ASYNC_CONTINUATION_REGISTER_ECX : ASYNC_CONTINUATION_REGISTER_EDX;
+        }
+        else
+        {
+            nSizeOfArgStack += sizeof(void *);
+            asyncContFlags = ASYNC_CONTINUATION_REGISTER_STACK;
+        }
+        m_dwFlags |= asyncContFlags;
+    }
+
     if (this->HasParamType())
     {
         DWORD paramTypeFlags = 0;
@@ -2139,25 +2229,19 @@ class ArgIteratorBase
     BOOL HasParamType()
     {
         LIMITED_METHOD_CONTRACT;
-        return m_pSig->GetCallingConventionInfo() & CORINFO_CALLCONV_PARAMTYPE;
+        return m_pSig->HasGenericContextArg();
     }
 
-    BOOL IsVarArg()
-    {
-        LIMITED_METHOD_CONTRACT;
-        return m_pSig->IsVarArg() || m_pSig->IsTreatAsVarArg();
-    }
-
-    BOOL IsAsyncCall()
+    BOOL HasAsyncContinuation()
     {
         LIMITED_METHOD_CONTRACT;
-        return m_pSig->IsAsyncCall();
+        return m_pSig->HasAsyncContinuation();
     }
 
-    BOOL HasAsyncContinuation()
+    BOOL IsVarArg()
     {
         LIMITED_METHOD_CONTRACT;
-        return m_pSig->HasAsyncContinuation();
+        return m_pSig->IsVarArg() || m_pSig->IsTreatAsVarArg();
     }
 
     DWORD NumFixedArgs()

src/coreclr/vm/frames.cpp

@@ -1530,6 +1530,17 @@ void TransitionFrame::PromoteCallerStackHelper(promote_func* fn, ScanContext* sc
             (fn)(PTR_PTR_Object(pThis), sc, CHECK_APP_DOMAIN);
     }
 
+    // Promote async continuation for async methods
+    if (argit.HasAsyncContinuation())
+    {
+        PTR_PTR_VOID pAsyncCont = dac_cast<PTR_PTR_VOID>(pTransitionBlock + argit.GetAsyncContinuationArgOffset());
+        LOG((LF_GC, INFO3,
+             "    async continuation argument at " FMT_ADDR "promoted from" FMT_ADDR "\n",
+             DBG_ADDR(pAsyncCont), DBG_ADDR(*pAsyncCont) ));
+
+        (fn)(PTR_PTR_Object(pAsyncCont), sc, CHECK_APP_DOMAIN);
+    }
+
     int argOffset;
     while ((argOffset = argit.GetNextOffset()) != TransitionBlock::InvalidOffset)
     {

src/coreclr/vm/gccover.cpp

@@ -444,13 +444,27 @@ void ReplaceInstrAfterCall(PBYTE instrToReplace, MethodDesc* callMD)
         return;
     }
 
-    if (IsPointerReturnKind(returnKind))
+    if (callMD->IsAsyncMethod())
     {
-        *instrToReplace = INTERRUPT_INSTR_PROTECT_FIRST_RET;
+        if (IsPointerReturnKind(returnKind))
+        {
+            *instrToReplace = INTERRUPT_INSTR_PROTECT_CONT_AND_RET;
+        }
+        else
+        {
+            *instrToReplace = INTERRUPT_INSTR_PROTECT_CONT;
+        }
     }
     else
     {
-        *instrToReplace = INTERRUPT_INSTR;
+        if (IsPointerReturnKind(returnKind))
+        {
+            *instrToReplace = INTERRUPT_INSTR_PROTECT_RET;
+        }
+        else
+        {
+            *instrToReplace = INTERRUPT_INSTR;
+        }
     }
 }
 
@@ -639,14 +653,22 @@ void DoGcStress (PCONTEXT regs, NativeCodeVersion nativeCodeVersion)
     // `if (!IsGcCoverageInterruptInstruction(instrPtr))` after we read `*instrPtr`.
 
     bool atCall;
-    bool afterCallProtect = false;
+    bool afterCallRetProtect = false;
+    bool afterCallContProtect = false;
 
     BYTE instrVal = *instrPtr;
     atCall = (instrVal == INTERRUPT_INSTR_CALL);
 
-    if (instrVal == INTERRUPT_INSTR_PROTECT_FIRST_RET)
+    if (instrVal == INTERRUPT_INSTR_PROTECT_RET ||
+        instrVal == INTERRUPT_INSTR_PROTECT_CONT_AND_RET)
     {
-        afterCallProtect = true;
+        afterCallRetProtect = true;
+    }
+
+    if (instrVal == INTERRUPT_INSTR_PROTECT_CONT ||
+        instrVal == INTERRUPT_INSTR_PROTECT_CONT_AND_RET)
+    {
+        afterCallContProtect = true;
     }
 
     if (!IsGcCoverageInterruptInstruction(instrPtr))
@@ -842,15 +864,20 @@ void DoGcStress (PCONTEXT regs, NativeCodeVersion nativeCodeVersion)
 
     // The legacy X86 GC encoder does not encode the state of return registers at
     // call sites, so we must add an extra frame to protect returns.
-    DWORD_PTR retValReg = 0;
+    DWORD_PTR protRegs[2] = {};
 
-    if (afterCallProtect)
+    if (afterCallRetProtect)
     {
-        retValReg = regs->Eax;
+        protRegs[0] = regs->Eax;
+    }
+
+    if (afterCallContProtect)
+    {
+        protRegs[1] = regs->Ecx;
     }
 
     _ASSERTE(sizeof(OBJECTREF) == sizeof(DWORD_PTR));
-    GCFrame gcFrame(pThread, (OBJECTREF*)&retValReg, 1, TRUE);
+    GCFrame gcFrame(pThread, (OBJECTREF*)protRegs, 2, TRUE);
 
     MethodDesc *pMD = nativeCodeVersion.GetMethodDesc();
     LOG((LF_GCROOTS, LL_EVERYTHING, "GCCOVER: Doing GC at method %s::%s offset 0x%x\n",
@@ -876,9 +903,14 @@ void DoGcStress (PCONTEXT regs, NativeCodeVersion nativeCodeVersion)
 
     CONSISTENCY_CHECK(!pThread->HasPendingGCStressInstructionUpdate());
 
-    if (afterCallProtect)
+    if (afterCallRetProtect)
+    {
+        regs->Eax = protRegs[0];
+    }
+
+    if (afterCallContProtect)
     {
-        regs->Eax = retValReg;
+        regs->Ecx = protRegs[1];
     }
 
     if (!Thread::UseRedirectForGcStress())

src/coreclr/vm/gccover.h

@@ -58,7 +58,9 @@ typedef DPTR(GCCoverageInfo) PTR_GCCoverageInfo; // see code:GCCoverageInfo::sav
 
 #if defined(TARGET_X86)
 #define INTERRUPT_INSTR_CALL                   0xFA    // X86 CLI instruction
-#define INTERRUPT_INSTR_PROTECT_FIRST_RET      0xFB    // X86 STI instruction, protect the first return register
+#define INTERRUPT_INSTR_PROTECT_RET            0xFB    // X86 STI instruction, protect the first return register
+#define INTERRUPT_INSTR_PROTECT_CONT           0xEC    // X86 IN instruction, protect the continuation register
+#define INTERRUPT_INSTR_PROTECT_CONT_AND_RET   0xED    // X86 IN instruction, protect both continuation and return registers
 #endif
 
 #elif defined(TARGET_ARM)
@@ -117,7 +119,9 @@ inline bool IsGcCoverageInterruptInstructionVal(UINT32 instrVal)
     {
     case INTERRUPT_INSTR:
     case INTERRUPT_INSTR_CALL:
-    case INTERRUPT_INSTR_PROTECT_FIRST_RET:
+    case INTERRUPT_INSTR_PROTECT_RET:
+    case INTERRUPT_INSTR_PROTECT_CONT:
+    case INTERRUPT_INSTR_PROTECT_CONT_AND_RET:
         return true;
     default:
         return false;