[clang] [CIR] Add structured control flow for coroutine suspend points (PR #213191)

[Erich Keane via cfe-commits]

================
@@ -479,54 +420,142 @@ CIRGenFunction::emitCoroutineBody(const CoroutineBodyStmt &s) {
                      .failed()) {
         return mlir::failure();
       }
-    }
 
-    mlir::Block &coroBodyBlock = coroBodyOp.getBody().back();
-    if (!coroBodyBlock.mightHaveTerminator()) {
-      mlir::OpBuilder::InsertionGuard guard(builder);
-      builder.setInsertionPointToEnd(&coroBodyBlock);
+      return mlir::success();
+    };
+
+    // Builds `final_suspend`: only emitted at all if the body can actually
+    // reach it (an explicit co_return, or falling off the end).
+    auto finalSuspendBuilder = [&]() -> mlir::LogicalResult {
+      // Note that LLVM checks CanFallthrough by looking into the availability
+      // of the insert block which is kinda brittle and unintuitive, seems to be
+      // related with how landing pads are handled.
+      //
+      // CIRGen handles this by checking pre-existing co_returns in the current
+      // scope instead.
+      //
+      // From LLVM IR Gen: const bool CanFallthrough = Builder.GetInsertBlock();
+      const bool canFallthrough = curLexScope->hasCoreturn();
+      const bool hasCoreturns = curCoro.data->coreturnCount > 0;
+      if (canFallthrough || hasCoreturns) {
+        curCoro.data->currentAwaitKind = cir::AwaitKind::Final;
+        if (emitStmt(s.getFinalSuspendStmt(), /*useCurrentScope=*/true)
+                .failed())
+          return mlir::failure();
+      }
       cir::YieldOp::create(builder, openCurlyLoc);
-    }
+      return mlir::success();
+    };
+
+    // Emit "if (coro.free(CoroId, CoroBegin)) Deallocate;"
+    auto destroyBuilder = [&]() -> mlir::LogicalResult {
+      Stmt *deallocate = s.getDeallocate();
+      if (emitStmt(deallocate, /*useCurrentScope=*/true).failed()) {
+        cgm.error(deallocate->getBeginLoc(),
+                  "failed to emit coroutine deallocation expression");
+        return mlir::failure();
+      }
+
+      cir::CoroFreeOp coroFree = curCoro.data->lastCoroFree;
 
-    // Note that LLVM checks CanFallthrough by looking into the availability
-    // of the insert block which is kinda brittle and unintuitive, seems to be
-    // related with how landing pads are handled.
-    //
-    // CIRGen handles this by checking pre-existing co_returns in the current
-    // scope instead.
-    //
-    // From LLVM IR Gen: const bool CanFallthrough = Builder.GetInsertBlock();
-    const bool canFallthrough = curLexScope->hasCoreturn();
-    const bool hasCoreturns = curCoro.data->coreturnCount > 0;
-    if (canFallthrough || hasCoreturns) {
-      curCoro.data->currentAwaitKind = cir::AwaitKind::Final;
+      if (!coroFree) {
+        cgm.error(deallocate->getBeginLoc(),
+                  "Deallocation expression does not refer to coro.free");
+        return mlir::failure();
+      }
       {
         mlir::OpBuilder::InsertionGuard guard(builder);
-        if (emitStmt(s.getFinalSuspendStmt(), /*useCurrentScope=*/true)
-                .failed())
+        builder.setInsertionPointAfter(coroFree);
+        mlir::Value isPtrNotNull =
+            builder.createPtrIsNotNull(coroFree.getResult());
+
+        llvm::SmallVector<mlir::Operation *> opsToMove;
+        mlir::Block *block = builder.getInsertionBlock();
+        mlir::Block::iterator it(isPtrNotNull.getDefiningOp());
+
+        for (++it; it != block->end(); ++it)
+          opsToMove.push_back(&*it);
+
+        auto ifOp = cir::IfOp::create(
+            builder, getLoc(deallocate->getSourceRange()), isPtrNotNull,
+            /*withElseRegion*/ false,
+            [&](mlir::OpBuilder &builder, mlir::Location loc) {
+              cir::YieldOp::create(builder, loc);
+            });
+
+        mlir::Operation *yieldOp = ifOp.getThenRegion().back().getTerminator();
+        for (auto *op : opsToMove)
+          op->moveBefore(yieldOp);
+      }
+
+      cir::YieldOp::create(builder, openCurlyLoc);
+      return mlir::success();
+    };
+
+    // Builds `exit`: coro.end(/*unwind*/ false) followed by the actual return
+    // to the caller.
+    auto exitBuilder = [&]() {
+      cir::ConstantOp nullHandler =
+          builder.getNullPtr(builder.getVoidPtrTy(), openCurlyLoc);
+      cir::ConstantOp noUnwind = builder.getBool(false, openCurlyLoc);
+      auto tkNone = cir::TokenNoneOp::create(builder, openCurlyLoc);
+      cir::CoroEndOp::create(builder, openCurlyLoc, nullHandler, noUnwind,
+                             tkNone);
+
+      if (auto *ret = cast_or_null<ReturnStmt>(s.getReturnStmt())) {
+        // Since we already emitted the return value above, so we shouldn't
+        // emit it again here.
+        Expr *previousRetValue = ret->getRetValue();
+        ret->setRetValue(nullptr);
+        if (emitStmt(ret, /*useCurrentScope=*/true).failed())
           return mlir::failure();
+        mlir::Block *block = builder.getInsertionBlock();
+        // emitReturnStmt() always creates a new insertion block after emitting
+        // the return. That block is unreachable in this case, so erase it.
+        block->erase();
+        // Set the return value back. The code generator, as the AST
+        // **Consumer**, shouldn't change the AST.
+        ret->setRetValue(previousRetValue);
       }
+
+      return mlir::success();
+    };
+
+    cir::CoroutineOp coro = cir::CoroutineOp::create(
+        builder, openCurlyLoc,
+        /*initialSuspendBuilder=*/
+        [&](mlir::OpBuilder &b, mlir::Location loc) {
+          res = initialSuspendBuilder();
+        },
+        /*bodyBuilder=*/
+        [&](mlir::OpBuilder &b, mlir::Location loc) {
+          if (res.succeeded())
----------------
erichkeane wrote:

I've not seen this pattern before, is there a reason we are doing it?  Is it really harmful to just always insert? 

https://github.com/llvm/llvm-project/pull/213191