[clang] [CIR] Implement FlattenCFG for coroutine AwaitOp, CoroBodyOp, and CoReturnOp (PR #203802)

[via cfe-commits]

================
@@ -2110,12 +2118,339 @@ class CIRTryOpFlattening : public mlir::OpRewritePattern<cir::TryOp> {
   }
 };
 
+static void
+collectCoReturnsOps(mlir::Region &region,
+                    llvm::SmallVectorImpl<cir::CoReturnOp> &coReturnOps) {
+  region.walk(
+      [&](cir::CoReturnOp coReturnOp) { coReturnOps.push_back(coReturnOp); });
+}
+
+// Flatten `cir.coroutine`. It runs once every nested op (awaits, scopes,
+// cleanups, ...) has been flattened. The initialSuspend, body and finalSuspend
+// regions become the body of a cleanup scope whose cleanup is the destroy
+// region; the exit region follows the scope:
+//
+//   cir.cleanup.scope {
+//     initialSuspend ---> body ---> finalSuspend
+//                          |             ^
+//                          +-------------+  cir.co_return
+//   } cleanup all {
+//     destroy                               coro.free
+//   }
+//   exit                                    coro.end + cir.return
+//
+// Leaving the scope body always runs the destroy region. In particular, the
+// cir.yield emitted for an await's destroy case leaves through it. A
+// suspended coroutine (cir.coro.exit) must keep its frame, so it branches
+// straight to the exit region.
+class CIRCoroutineFlattening : public mlir::OpRewritePattern<cir::CoroutineOp> {
+  void collectCoroExitsOps(
+      cir::CoroutineOp &coroOp,
+      llvm::SmallVectorImpl<cir::CoroExitOp> &coroExitsOps) const {
+    // cir.coro.exit is produced by flattening a cir.await, which can only
+    // appear in the initialSuspend, body and finalSuspend regions.
+    auto collectCoroExit = [&](cir::CoroExitOp op) {
+      coroExitsOps.push_back(op);
+    };
+    coroOp.getInitialSuspend().walk(collectCoroExit);
+    coroOp.getBody().walk(collectCoroExit);
+    coroOp.getFinalSuspend().walk(collectCoroExit);
+  }
+
+public:
+  using OpRewritePattern<cir::CoroutineOp>::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(cir::CoroutineOp coroOp,
+                  mlir::PatternRewriter &rewriter) const override {
+
+    assert(!cir::MissingFeatures::coroutineUnwindRegion());
----------------
Andres-Salamanca wrote:

Done! I’ve refactored the function into smaller helpers. Let me know what you think of the result.

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