[C++20] [Coroutines] Mark coroutine done if unhandled_exception throws

According to [dcl.fct.def.coroutine]/p14:
> If the evaluation of the expression promise.unhandled_­exception()
> exits via an exception, the coroutine is considered suspended at the
> final suspend point.

But this is not implemented in clang before. This patch would implement
this feature by marking the coroutine as done at the place of
coro.end(frame, /*InUnwindPath=*/true ).

Reviewed By: rjmccall

Differential Revision: https://reviews.llvm.org/D115219

llvm/docs/Coroutines.rst

@@ -1377,17 +1377,23 @@ The `CoroSplit` pass, if the funclet bundle is present, will insert
 ``cleanupret from %tok unwind to caller`` before
 the `coro.end`_ intrinsic and will remove the rest of the block.
 
+In the unwind path (when the argument is `true`), `coro.end` will mark the coroutine
+as done, making it undefined behavior to resume the coroutine again and causing 
+`llvm.coro.done` to return `true`.  This is not necessary in the normal path because
+the coroutine will already be marked as done by the final suspend.
+
 The following table summarizes the handling of `coro.end`_ intrinsic.
 
-+--------------------------+-------------------+-------------------------------+
-|                          | In Start Function | In Resume/Destroy Functions   |
-+--------------------------+-------------------+-------------------------------+
-|unwind=false              | nothing           |``ret void``                   |
-+------------+-------------+-------------------+-------------------------------+
-|            | WinEH       | nothing           |``cleanupret unwind to caller``|
-|unwind=true +-------------+-------------------+-------------------------------+
-|            | Landingpad  | nothing           | nothing                       |
-+------------+-------------+-------------------+-------------------------------+
++--------------------------+------------------------+-------------------------------+
+|                          | In Start Function      | In Resume/Destroy Functions   |
++--------------------------+------------------------+-------------------------------+
+|unwind=false              | nothing                |``ret void``                   |
++------------+-------------+------------------------+-------------------------------+
+|            | WinEH       | mark coroutine as done |``cleanupret unwind to caller``|
+|            |             |                        |mark coroutine done            |
+|unwind=true +-------------+------------------------+-------------------------------+
+|            | Landingpad  | mark coroutine as done | mark coroutine done           |
++------------+-------------+------------------------+-------------------------------+
 
 
 'llvm.coro.end.async' Intrinsic

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -280,6 +280,27 @@ static void replaceFallthroughCoroEnd(AnyCoroEndInst *End,
   BB->getTerminator()->eraseFromParent();
 }
 
+// Mark a coroutine as done, which implies that the coroutine is finished and
+// never get resumed.
+//
+// In resume-switched ABI, the done state is represented by storing zero in
+// ResumeFnAddr.
+//
+// NOTE: We couldn't omit the argument `FramePtr`. It is necessary because the
+// pointer to the frame in splitted function is not stored in `Shape`.
+static void markCoroutineAsDone(IRBuilder<> &Builder, const coro::Shape &Shape,
+                                Value *FramePtr) {
+  assert(
+      Shape.ABI == coro::ABI::Switch &&
+      "markCoroutineAsDone is only supported for Switch-Resumed ABI for now.");
+  auto *GepIndex = Builder.CreateStructGEP(
+      Shape.FrameTy, FramePtr, coro::Shape::SwitchFieldIndex::Resume,
+      "ResumeFn.addr");
+  auto *NullPtr = ConstantPointerNull::get(cast<PointerType>(
+      Shape.FrameTy->getTypeAtIndex(coro::Shape::SwitchFieldIndex::Resume)));
+  Builder.CreateStore(NullPtr, GepIndex);
+}
+
 /// Replace an unwind call to llvm.coro.end.
 static void replaceUnwindCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,
                                  Value *FramePtr, bool InResume,
@@ -288,10 +309,18 @@ static void replaceUnwindCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,
 
   switch (Shape.ABI) {
   // In switch-lowering, this does nothing in the main function.
-  case coro::ABI::Switch:
+  case coro::ABI::Switch: {
+    // In C++'s specification, the coroutine should be marked as done
+    // if promise.unhandled_exception() throws.  The frontend will
+    // call coro.end(true) along this path.
+    //
+    // FIXME: We should refactor this once there is other language
+    // which uses Switch-Resumed style other than C++.
+    markCoroutineAsDone(Builder, Shape, FramePtr);
     if (!InResume)
       return;
     break;
+  }
   // In async lowering this does nothing.
   case coro::ABI::Async:
     break;
@@ -364,13 +393,9 @@ static void createResumeEntryBlock(Function &F, coro::Shape &Shape) {
     auto *Save = S->getCoroSave();
     Builder.SetInsertPoint(Save);
     if (S->isFinal()) {
-      // Final suspend point is represented by storing zero in ResumeFnAddr.
-      auto *GepIndex = Builder.CreateStructGEP(FrameTy, FramePtr,
-                                 coro::Shape::SwitchFieldIndex::Resume,
-                                  "ResumeFn.addr");
-      auto *NullPtr = ConstantPointerNull::get(cast<PointerType>(
-          FrameTy->getTypeAtIndex(coro::Shape::SwitchFieldIndex::Resume)));
-      Builder.CreateStore(NullPtr, GepIndex);
+      // The coroutine should be marked done if it reaches the final suspend
+      // point.
+      markCoroutineAsDone(Builder, Shape, FramePtr);
     } else {
       auto *GepIndex = Builder.CreateStructGEP(
           FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");

llvm/test/Transforms/Coroutines/coro-split-eh-00.ll

@@ -67,6 +67,9 @@ eh.resume:
 ; CHECK-NEXT:      %lpval = landingpad { i8*, i32 }
 ; CHECK-NEXT:         cleanup
 ; CHECK-NEXT:      call void @print(i32 2)
+; Checks that the coroutine would be marked as done if it exits in unwinding path.
+; CHECK-NEXT:      %[[RESUME_ADDR:.+]] = getelementptr inbounds %[[FRAME_TY:.+]], %[[FRAME_TY]]* %FramePtr, i32 0, i32 0
+; CHECK-NEXT:      store void (%[[FRAME_TY]]*)* null, void (%[[FRAME_TY]]*)** %[[RESUME_ADDR]], align 8
 ; CHECK-NEXT:      resume { i8*, i32 } %lpval
 
 declare i8* @llvm.coro.free(token, i8*)

llvm/test/Transforms/Coroutines/coro-split-eh-01.ll

@@ -61,6 +61,9 @@ cleanup.cont:
 ; CHECK:      lpad:
 ; CHECK-NEXT:   %tok = cleanuppad within none []
 ; CHECK-NEXT:   call void @print(i32 2)
+; Checks that the coroutine would be marked as done if it exits in unwinding path.
+; CHECK-NEXT:   %[[RESUME_ADDR:.+]] = getelementptr inbounds %[[FRAME_TY:.+]], %[[FRAME_TY]]* %FramePtr, i32 0, i32 0
+; CHECK-NEXT:   store void (%[[FRAME_TY]]*)* null, void (%[[FRAME_TY]]*)** %[[RESUME_ADDR]], align 8
 ; CHECK-NEXT:   cleanupret from %tok unwind to caller
 
 declare i8* @llvm.coro.free(token, i8*)