Generalize llvm.coro.suspend.retcon to allow an arbitrary number of arguments to be passed back to the continuation function.

llvm-svn: 368789

3 files changed, 93 insertions, 24 deletions

diff --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h
index 78e44988428..d2348057c24 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInternal.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInternal.h
@@ -171,7 +171,7 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
            ABI == coro::ABI::RetconOnce);
     auto FTy = CoroBegin->getFunction()->getFunctionType();
 
-    // This is checked by AnyCoroIdRetconInst::isWellFormed().
+    // The safety of all this is checked by checkWFRetconPrototype.
     if (auto STy = dyn_cast<StructType>(FTy->getReturnType())) {
       return STy->elements().slice(1);
     } else {
@@ -179,6 +179,15 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
     }
   }
 
+  ArrayRef<Type*> getRetconResumeTypes() const {
+    assert(ABI == coro::ABI::Retcon ||
+           ABI == coro::ABI::RetconOnce);
+
+    // The safety of all this is checked by checkWFRetconPrototype.
+    auto FTy = RetconLowering.ResumePrototype->getFunctionType();
+    return FTy->params().slice(1);
+  }
+
   CallingConv::ID getResumeFunctionCC() const {
     switch (ABI) {
     case coro::ABI::Switch:
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index b6862d1f67c..318a012b6a7 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -144,6 +144,7 @@ private:
   void createDeclaration();
   void replaceEntryBlock();
   Value *deriveNewFramePointer();
+  void replaceRetconSuspendUses();
   void replaceCoroSuspends();
   void replaceCoroEnds();
   void handleFinalSuspend();
@@ -402,6 +403,52 @@ static Function *createCloneDeclaration(Function &OrigF, coro::Shape &Shape,
   return NewF;
 }
 
+/// Replace uses of the active llvm.coro.suspend.retcon call with the
+/// arguments to the continuation function.
+///
+/// This assumes that the builder has a meaningful insertion point.
+void CoroCloner::replaceRetconSuspendUses() {
+  assert(Shape.ABI == coro::ABI::Retcon ||
+         Shape.ABI == coro::ABI::RetconOnce);
+
+  auto NewS = VMap[ActiveSuspend];
+  if (NewS->use_empty()) return;
+
+  // Copy out all the continuation arguments after the buffer pointer into
+  // an easily-indexed data structure for convenience.
+  SmallVector<Value*, 8> Args;
+  for (auto I = std::next(NewF->arg_begin()), E = NewF->arg_end(); I != E; ++I)
+    Args.push_back(&*I);
+
+  // If the suspend returns a single scalar value, we can just do a simple
+  // replacement.
+  if (!isa<StructType>(NewS->getType())) {
+    assert(Args.size() == 1);
+    NewS->replaceAllUsesWith(Args.front());
+    return;
+  }
+
+  // Try to peephole extracts of an aggregate return.
+  for (auto UI = NewS->use_begin(), UE = NewS->use_end(); UI != UE; ) {
+    auto EVI = dyn_cast<ExtractValueInst>((UI++)->getUser());
+    if (!EVI || EVI->getNumIndices() != 1)
+      continue;
+
+    EVI->replaceAllUsesWith(Args[EVI->getIndices().front()]);
+    EVI->eraseFromParent();
+  }
+
+  // If we have no remaining uses, we're done.
+  if (NewS->use_empty()) return;
+
+  // Otherwise, we need to create an aggregate.
+  Value *Agg = UndefValue::get(NewS->getType());
+  for (size_t I = 0, E = Args.size(); I != E; ++I)
+    Agg = Builder.CreateInsertValue(Agg, Args[I], I);
+
+  NewS->replaceAllUsesWith(Agg);
+}
+
 void CoroCloner::replaceCoroSuspends() {
   Value *SuspendResult;
 
@@ -416,15 +463,12 @@ void CoroCloner::replaceCoroSuspends() {
     SuspendResult = Builder.getInt8(isSwitchDestroyFunction() ? 1 : 0);
     break;
 
-  // In continuation lowering, replace all of the suspend uses with false to
-  // indicate that they're not unwinding resumes.  We've already mapped the
-  // active suspend to the appropriate argument, so any other suspend values
-  // that are still being used must be from previous suspends.  It's UB to try
-  // to suspend during unwind, so they must be from regular resumes.
+  // In returned-continuation lowering, the arguments from earlier
+  // continuations are theoretically arbitrary, and they should have been
+  // spilled.
   case coro::ABI::RetconOnce:
   case coro::ABI::Retcon:
-    SuspendResult = Builder.getInt1(false);
-    break;
+    return;
   }
 
   for (AnyCoroSuspendInst *CS : Shape.CoroSuspends) {
@@ -619,14 +663,11 @@ void CoroCloner::create() {
 
   case coro::ABI::Retcon:
   case coro::ABI::RetconOnce:
-    // Replace the active suspend with the should-unwind argument.
-    // Coerce it to i1 if necessary.
+    // Replace uses of the active suspend with the corresponding
+    // continuation-function arguments.
     assert(ActiveSuspend != nullptr &&
            "no active suspend when lowering a continuation-style coroutine");
-    Value *ShouldUnwind = &*std::next(NewF->arg_begin());
-    if (!ShouldUnwind->getType()->isIntegerTy(1))
-      ShouldUnwind = Builder.CreateIsNotNull(ShouldUnwind);
-    VMap[ActiveSuspend]->replaceAllUsesWith(ShouldUnwind);
+    replaceRetconSuspendUses();
     break;
   }
 
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index 2e690a78c42..6239674cc69 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -387,6 +387,7 @@ void coro::Shape::buildFrom(Function &F) {
     // Determine the result value types, and make sure they match up with
     // the values passed to the suspends.
     auto ResultTys = getRetconResultTypes();
+    auto ResumeTys = getRetconResumeTypes();
 
     for (auto AnySuspend : CoroSuspends) {
       auto Suspend = dyn_cast<CoroSuspendRetconInst>(AnySuspend);
@@ -396,6 +397,7 @@ void coro::Shape::buildFrom(Function &F) {
                            "coro.suspend.retcon");
       }
 
+      // Check that the argument types of the suspend match the results.
       auto SI = Suspend->value_begin(), SE = Suspend->value_end();
       auto RI = ResultTys.begin(), RE = ResultTys.end();
       for (; SI != SE && RI != RE; ++SI, ++RI) {
@@ -411,6 +413,26 @@ void coro::Shape::buildFrom(Function &F) {
         Prototype->getFunctionType()->dump();
         report_fatal_error("wrong number of arguments to coro.suspend.retcon");
       }
+
+      // Check that the result type of the suspend matches the resume types.
+      Type *SResultTy = Suspend->getType();
+      ArrayRef<Type*> SuspendResultTys =
+        (isa<StructType>(SResultTy)
+           ? cast<StructType>(SResultTy)->elements()
+           : SResultTy); // forms an ArrayRef using SResultTy, be careful
+      if (SuspendResultTys.size() != ResumeTys.size()) {
+        Suspend->dump();
+        Prototype->getFunctionType()->dump();
+        report_fatal_error("wrong number of results from coro.suspend.retcon");
+      }
+      for (size_t I = 0, E = ResumeTys.size(); I != E; ++I) {
+        if (SuspendResultTys[I] != ResumeTys[I]) {
+          Suspend->dump();
+          Prototype->getFunctionType()->dump();
+          report_fatal_error("result from coro.suspend.retcon does not "
+                             "match corresponding prototype function param");
+        }
+      }
     }
     break;
   }
@@ -501,7 +523,7 @@ static void fail(const Instruction *I, const char *Reason, Value *V) {
 static void checkWFRetconPrototype(const AnyCoroIdRetconInst *I, Value *V) {
   auto F = dyn_cast<Function>(V->stripPointerCasts());
   if (!F)
-    fail(I, "llvm.coro.retcon.* prototype not a Function", V);
+    fail(I, "llvm.coro.id.retcon.* prototype not a Function", V);
 
   auto FT = F->getFunctionType();
 
@@ -517,23 +539,20 @@ static void checkWFRetconPrototype(const AnyCoroIdRetconInst *I, Value *V) {
       ResultOkay = false;
     }
     if (!ResultOkay)
-      fail(I, "llvm.coro.retcon prototype must return pointer as first result",
-           F);
+      fail(I, "llvm.coro.id.retcon prototype must return pointer as first "
+              "result", F);
 
     if (FT->getReturnType() !=
           I->getFunction()->getFunctionType()->getReturnType())
-      fail(I, "llvm.coro.retcon.* prototype return type must be same as"
+      fail(I, "llvm.coro.id.retcon prototype return type must be same as"
               "current function return type", F);
   } else {
     // No meaningful validation to do here for llvm.coro.id.unique.once.
   }
 
-  if (FT->getNumParams() != 2)
-    fail(I, "llvm.coro.retcon.* prototype must take exactly two parameters", F);
-  if (!FT->getParamType(0)->isPointerTy())
-    fail(I, "llvm.coro.retcon.* prototype must take pointer as 1st param", F);
-  if (!FT->getParamType(1)->isIntegerTy()) // an i1, but not for abi purposes
-    fail(I, "llvm.coro.retcon.* prototype must take integer as 2nd param", F);
+  if (FT->getNumParams() == 0 || !FT->getParamType(0)->isPointerTy())
+    fail(I, "llvm.coro.id.retcon.* prototype must take pointer as "
+            "its first parameter", F);
 }
 
 /// Check that the given value is a well-formed allocator.