Implement AsyncHelpers.TailAwait in the interpreter and use it for instantiating and unboxing stubs on CoreCLR

src/coreclr/interpreter/compiler.cpp

@@ -3194,7 +3194,13 @@ bool InterpCompiler::EmitNamedIntrinsicCall(NamedIntrinsic ni, bool nonVirtualCa
             m_pStackPointer--;
             AddIns(INTOP_RET_VOID);
             return true;
-
+        case NI_System_Runtime_CompilerServices_AsyncHelpers_TailAwait:
+            if ((m_methodInfo->options & CORINFO_ASYNC_SAVE_CONTEXTS) != 0)
+            {
+                BADCODE("TailAwait is not supported in async methods that capture contexts");
+            }
+            m_nextAwaitIsTail = true;
+            return true;
         case NI_System_Runtime_CompilerServices_AsyncHelpers_AsyncCallContinuation:
             if (m_methodInfo->args.isAsyncCall())
             {
@@ -4620,7 +4626,8 @@ void InterpCompiler::EmitCall(CORINFO_RESOLVED_TOKEN* pConstrainedToken, bool re
                     ni == NI_System_Runtime_CompilerServices_RuntimeHelpers_SetNextCallGenericContext ||
                     ni == NI_System_Runtime_CompilerServices_RuntimeHelpers_SetNextCallAsyncContinuation ||
                     ni == NI_System_Runtime_CompilerServices_AsyncHelpers_AsyncCallContinuation ||
-                    ni == NI_System_Runtime_CompilerServices_AsyncHelpers_AsyncSuspend);
+                    ni == NI_System_Runtime_CompilerServices_AsyncHelpers_AsyncSuspend ||
+                    ni == NI_System_Runtime_CompilerServices_AsyncHelpers_TailAwait);
             if ((InterpConfig.InterpMode() == 3) || isMustExpand)
             {
                 if (EmitNamedIntrinsicCall(ni, callInfo.kind == CORINFO_CALL, resolvedCallToken.hClass, callInfo.hMethod, callInfo.sig))
@@ -5416,8 +5423,15 @@ void SetSlotToTrue(TArray<bool, MemPoolAllocator> &gcRefMap, int32_t slotOffset)
     gcRefMap.Set(slotIndex, true);
 }
 
-void InterpCompiler::EmitSuspend(const CORINFO_CALL_INFO &callInfo, ContinuationContextHandling ContinuationContextHandling)
+void InterpCompiler::EmitSuspend(const CORINFO_CALL_INFO &callInfo, ContinuationContextHandling continuationContextHandling)
 {
+    if (m_nextAwaitIsTail)
+    {
+        AddIns(INTOP_RET_EXISTING_CONTINUATION);
+        m_nextAwaitIsTail = false;
+        return;
+    }
+
     CORINFO_LOOKUP_KIND kindForAllocationContinuation;
     m_compHnd->getLocationOfThisType(m_methodHnd, &kindForAllocationContinuation);
 
@@ -5440,7 +5454,7 @@ void InterpCompiler::EmitSuspend(const CORINFO_CALL_INFO &callInfo, Continuation
     }
     bool needsEHHandling = m_pCBB->enclosingTryBlockCount > (m_isAsyncMethodWithContextSaveRestore ? 1 : 0);
 
-    bool captureContinuationContext = ContinuationContextHandling == ContinuationContextHandling::ContinueOnCapturedContext;
+    bool captureContinuationContext = continuationContextHandling == ContinuationContextHandling::ContinueOnCapturedContext;
 
     // 1. For each IL var that is live across the await/continuation point. Add it to the list of live vars
     // 2. For each stack var that is live other than the return value from the call
@@ -5663,7 +5677,7 @@ void InterpCompiler::EmitSuspend(const CORINFO_CALL_INFO &callInfo, Continuation
         flags |= CORINFO_CONTINUATION_HAS_CONTINUATION_CONTEXT;
     }
 
-    if (ContinuationContextHandling == ContinuationContextHandling::ContinueOnThreadPool)
+    if (continuationContextHandling == ContinuationContextHandling::ContinueOnThreadPool)
     {
         flags |= CORINFO_CONTINUATION_CONTINUE_ON_THREAD_POOL;
     }

src/coreclr/interpreter/compiler.h

@@ -673,6 +673,11 @@ class InterpCompiler
     int32_t m_nextCallGenericContextVar;
     int32_t m_nextCallAsyncContinuationVar;
 
+    // If true, the next await should be done as a tail await that just
+    // directly returns the continuation of the call instead of creating a new
+    // suspension point.
+    bool m_nextAwaitIsTail = false;
+
     // Table of mappings of leave instructions to the first finally call island the leave
     // needs to execute.
     TArray<LeavesTableEntry, MemPoolAllocator> m_leavesTable;

src/coreclr/interpreter/inc/intops.def

@@ -450,6 +450,7 @@ OPDEF(INTOP_HANDLE_CONTINUATION_RESUME, "handle.continuation.resume", 2, 0, 0, I
 OPDEF(INTOP_CHECK_FOR_CONTINUATION, "check.for.continuation", 3, 0, 1, InterpOpNoArgs)
 OPDEF(INTOP_CAPTURE_CONTEXT_ON_SUSPEND, "capture.context.on.suspend", 4, 1, 1, InterpOpHandleContinuationPt2)
 OPDEF(INTOP_RESTORE_CONTEXTS_ON_SUSPEND, "restore.contexts.on.suspend", 6, 1, 3, InterpOpHandleContinuationPt2)
+OPDEF(INTOP_RET_EXISTING_CONTINUATION, "ret.existing.continuation", 1, 0, 0, InterpOpNoArgs)
 
 // Intrinsics
 OPDEF(INTOP_COMPARE_EXCHANGE_U1, "compare.exchange.u1", 5, 1, 3, InterpOpNoArgs)

src/coreclr/interpreter/intrinsics.cpp

@@ -123,6 +123,8 @@ NamedIntrinsic GetNamedIntrinsic(COMP_HANDLE compHnd, CORINFO_METHOD_HANDLE comp
                     return NI_System_Runtime_CompilerServices_AsyncHelpers_AsyncCallContinuation;
                 else if (!strcmp(methodName, "Await"))
                     return NI_System_Runtime_CompilerServices_AsyncHelpers_Await;
+                else if (!strcmp(methodName, "TailAwait"))
+                    return NI_System_Runtime_CompilerServices_AsyncHelpers_TailAwait;
             }
         }
         else if (!strcmp(namespaceName, "System.Runtime.InteropServices"))

src/coreclr/vm/corelib.h

@@ -715,6 +715,7 @@ DEFINE_METHOD(ASYNC_HELPERS,      CAPTURE_CONTEXTS,          CaptureContexts, No
 DEFINE_METHOD(ASYNC_HELPERS,      RESTORE_CONTEXTS,          RestoreContexts, NoSig)
 DEFINE_METHOD(ASYNC_HELPERS,      RESTORE_CONTEXTS_ON_SUSPENSION, RestoreContextsOnSuspension, NoSig)
 DEFINE_METHOD(ASYNC_HELPERS,      ASYNC_CALL_CONTINUATION,   AsyncCallContinuation, NoSig)
+DEFINE_METHOD(ASYNC_HELPERS,      TAIL_AWAIT,                TailAwait, NoSig)
 
 #ifdef FEATURE_INTERPRETER
 DEFINE_METHOD(ASYNC_HELPERS,      RESUME_INTERPRETER_CONTINUATION, ResumeInterpreterContinuation, NoSig)

src/coreclr/vm/interpexec.cpp

@@ -4320,6 +4320,19 @@ do                                                                      \
                     goto EXIT_FRAME;
                 }
 
+                case INTOP_RET_EXISTING_CONTINUATION:
+                {
+                    if (pInterpreterFrame->GetContinuation() == NULL)
+                    {
+                        // No continuation returned
+                        ip++;
+                        break;
+                    }
+
+                    // Otherwise exit without modifying current continuation
+                    goto EXIT_FRAME;
+                }
+
                 case INTOP_HANDLE_CONTINUATION_RESUME:
                 {
                     InterpAsyncSuspendData *pAsyncSuspendData = (InterpAsyncSuspendData*)pMethod->pDataItems[ip[1]];

src/coreclr/vm/prestub.cpp

@@ -1498,6 +1498,11 @@ Stub * CreateUnboxingILStubForValueTypeMethods(MethodDesc* pTargetMD)
     // Push the target address
     pCode->EmitLDC((TADDR)pTargetMD->GetMultiCallableAddrOfCode(CORINFO_ACCESS_ANY));
 
+    if (pTargetMD->IsAsyncMethod())
+    {
+        pCode->EmitCALL(METHOD__ASYNC_HELPERS__TAIL_AWAIT, 0, 0);
+    }
+
     // Do the calli
     pCode->EmitCALLI(TOKEN_ILSTUB_TARGET_SIG, msig.NumFixedArgs() + 1, msig.IsReturnTypeVoid() ? 0 : 1);
     pCode->EmitRET();
@@ -1587,6 +1592,11 @@ Stub * CreateInstantiatingILStub(MethodDesc* pTargetMD, void* pHiddenArg)
     // Push the target address
     pCode->EmitLDC((TADDR)pTargetMD->GetMultiCallableAddrOfCode(CORINFO_ACCESS_ANY));
 
+    if (pTargetMD->IsAsyncMethod())
+    {
+        pCode->EmitCALL(METHOD__ASYNC_HELPERS__TAIL_AWAIT, 0, 0);
+    }
+
     // Do the calli
     pCode->EmitCALLI(TOKEN_ILSTUB_TARGET_SIG, msig.NumFixedArgs() + (msig.HasThis() ? 1 : 0), msig.IsReturnTypeVoid() ? 0 : 1);
     pCode->EmitRET();