Merge pull request swiftlang#10132 from nate-chandler/coro-popless

[Coro] RetconOnceDynamic: Popless return on null allocator.

Author: nate-chandler

llvm/include/llvm/AsmParser/LLToken.h

@@ -161,6 +161,7 @@ enum Kind {
   kw_anyregcc,
   kw_swiftcc,
   kw_swifttailcc,
+  kw_swiftcorocc,
   kw_preserve_mostcc,
   kw_preserve_allcc,
   kw_preserve_nonecc,

llvm/include/llvm/IR/CallingConv.h

@@ -270,6 +270,10 @@ namespace CallingConv {
     /// Preserve X1-X15, X19-X29, SP, Z0-Z31, P0-P15.
     AArch64_SME_ABI_Support_Routines_PreserveMost_From_X1 = 111,
 
+    /// This follows the Swift calling convention in how arguments are passed
+    /// but doesn't clean up the stack on a return.
+    SwiftCoro = 112, // FIXME: allocate
+
     /// The highest possible ID. Must be some 2^k - 1.
     MaxID = 1023
   };

llvm/include/llvm/IR/Intrinsics.td

@@ -1754,6 +1754,10 @@ def int_coro_await_suspend_handle : Intrinsic<[],
                                               [llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
                                               [Throws]>;
 
+// FIXME: enforce musttail
+// XXX: attrs;  not throws, wb DefaultAttrsIntrinsic
+def int_coro_return : Intrinsic<[llvm_any_ty], [LLVMMatchType<0>], []>;
+
 // Coroutine Lowering Intrinsics. Used internally by coroutine passes.
 
 def int_coro_subfn_addr : DefaultAttrsIntrinsic<

llvm/lib/AsmParser/LLLexer.cpp

@@ -619,6 +619,7 @@ lltok::Kind LLLexer::LexIdentifier() {
   KEYWORD(x86_regcallcc);
   KEYWORD(swiftcc);
   KEYWORD(swifttailcc);
+  KEYWORD(swiftcorocc);
   KEYWORD(anyregcc);
   KEYWORD(preserve_mostcc);
   KEYWORD(preserve_allcc);

llvm/lib/AsmParser/LLParser.cpp

@@ -2187,6 +2187,7 @@ void LLParser::parseOptionalDLLStorageClass(unsigned &Res) {
 ///   ::= 'ghccc'
 ///   ::= 'swiftcc'
 ///   ::= 'swifttailcc'
+///   ::= 'swiftcorocc'
 ///   ::= 'x86_intrcc'
 ///   ::= 'hhvmcc'
 ///   ::= 'hhvm_ccc'
@@ -2252,6 +2253,7 @@ bool LLParser::parseOptionalCallingConv(unsigned &CC) {
   case lltok::kw_ghccc:          CC = CallingConv::GHC; break;
   case lltok::kw_swiftcc:        CC = CallingConv::Swift; break;
   case lltok::kw_swifttailcc:    CC = CallingConv::SwiftTail; break;
+  case lltok::kw_swiftcorocc:    CC = CallingConv::SwiftCoro; break;
   case lltok::kw_x86_intrcc:     CC = CallingConv::X86_INTR; break;
   case lltok::kw_hhvmcc:
     CC = CallingConv::DUMMY_HHVM;

llvm/lib/IR/AsmWriter.cpp

@@ -343,6 +343,7 @@ static void PrintCallingConv(unsigned cc, raw_ostream &Out) {
   case CallingConv::SPIR_KERNEL:   Out << "spir_kernel"; break;
   case CallingConv::Swift:         Out << "swiftcc"; break;
   case CallingConv::SwiftTail:     Out << "swifttailcc"; break;
+  case CallingConv::SwiftCoro:     Out << "swiftcorocc"; break;
   case CallingConv::X86_INTR:      Out << "x86_intrcc"; break;
   case CallingConv::DUMMY_HHVM:
     Out << "hhvmcc";

llvm/lib/Transforms/Coroutines/CoroCleanup.cpp

@@ -72,6 +72,10 @@ bool Lowerer::lower(Function &F) {
       case Intrinsic::coro_id_async:
         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
         break;
+      case Intrinsic::coro_return:
+        // FIXME: Remove this case with backend support.
+        II->replaceAllUsesWith(II->getArgOperand(0));
+        break;
       case Intrinsic::coro_subfn_addr:
         lowerSubFn(Builder, cast<CoroSubFnInst>(II));
         break;

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -1967,7 +1967,34 @@ static void splitRetconCoroutine(Function &F, coro::Shape &Shape,
           RetV = Builder.CreateInsertValue(RetV, ReturnPHIs[I], I);
       }
 
-      Builder.CreateRet(RetV);
+      if (Shape.ABI == coro::ABI::RetconOnceDynamic &&
+          F.getCallingConv() == CallingConv::SwiftCoro) {
+        //  %retval = ...
+        //  %null_allocator = icmp %1, null
+        //  br i1 %null_allocator, label %popless, label %normal
+        // popless:
+        //  ret %retval
+        // normal:
+        //  %popless_retval = musttail call i64 @llvm.coro.return(%retval)
+        //  ret %popless_retval
+        auto *NullAllocator = Builder.CreateCmp(
+            CmpInst::Predicate::ICMP_EQ, Shape.RetconLowering.Allocator,
+            ConstantPointerNull::get(
+                cast<PointerType>(Shape.RetconLowering.Allocator->getType())));
+        auto *PoplessReturnBB = BasicBlock::Create(
+            F.getContext(), "coro.return.popless", &F, NewSuspendBB);
+        auto *NormalReturnBB = BasicBlock::Create(
+            F.getContext(), "coro.return.normal", &F, NewSuspendBB);
+        Builder.CreateCondBr(NullAllocator, PoplessReturnBB, NormalReturnBB);
+        IRBuilder<> PoplessBuilder(PoplessReturnBB);
+        auto *WrapRetV = PoplessBuilder.CreateIntrinsic(
+            RetV->getType(), Intrinsic::coro_return, {RetV});
+        PoplessBuilder.CreateRet(WrapRetV);
+        IRBuilder<> NormalBuilder(NormalReturnBB);
+        NormalBuilder.CreateRet(RetV);
+      } else {
+        Builder.CreateRet(RetV);
+      }
     }
 
     // Branch to the return block.