[clang] [CIR] Remove overly strict end_catch assertion in catch handler flattening (PR #193796)

[via cfe-commits]

https://github.com/adams381 created https://github.com/llvm/llvm-project/pull/193796

The assertion in `flattenCatchHandler` required `end_catch` to be the last operation before `yield` in catch handlers, with only branches in between.  Complex catch handlers with cleanup code or nested control flow can have additional operations between `end_catch` and `yield`.  The yield-to-branch replacement does not depend on `end_catch` position.

Made with [Cursor](https://cursor.com)

>From d6ea002aef004d07a49d16962559b76f263fb3bd Mon Sep 17 00:00:00 2001
From: Adam Smith <adams at nvidia.com>
Date: Thu, 23 Apr 2026 09:47:09 -0700
Subject: [PATCH] [CIR] Remove overly strict end_catch assertion in catch
 handler flattening

The assertion required end_catch to be the last operation before
yield in catch handlers, with only branches in between.  Complex
catch handlers (e.g., with cleanup code or nested control flow)
can have additional operations between end_catch and yield.  The
yield-to-branch replacement does not depend on end_catch position.

Made-with: Cursor
---
 .../lib/CIR/Dialect/Transforms/FlattenCFG.cpp | 24 ------------
 .../Transforms/flatten-try-catch-cleanup.cir  | 39 +++++++++++++++++++
 2 files changed, 39 insertions(+), 24 deletions(-)
 create mode 100644 clang/test/CIR/Transforms/flatten-try-catch-cleanup.cir

diff --git a/clang/lib/CIR/Dialect/Transforms/FlattenCFG.cpp b/clang/lib/CIR/Dialect/Transforms/FlattenCFG.cpp
index 48c47deb2bc0a..9e9d4a9ebcd59 100644
--- a/clang/lib/CIR/Dialect/Transforms/FlattenCFG.cpp
+++ b/clang/lib/CIR/Dialect/Transforms/FlattenCFG.cpp
@@ -1566,30 +1566,6 @@ class CIRTryOpFlattening : public mlir::OpRewritePattern<cir::TryOp> {
     for (mlir::Block &block : llvm::make_range(handlerEntry->getIterator(),
                                                insertBefore->getIterator())) {
       if (auto yieldOp = dyn_cast<cir::YieldOp>(block.getTerminator())) {
-        // Verify that end_catch is the last non-branch operation before
-        // this yield. After cleanup scope flattening, end_catch may be in
-        // a predecessor block rather than immediately before the yield.
-        // Walk back through the single-predecessor chain, verifying that
-        // each intermediate block contains only a branch terminator, until
-        // we find end_catch as the last non-terminator in some block.
-        assert([&]() {
-          // Check if end_catch immediately precedes the yield.
-          if (mlir::Operation *prev = yieldOp->getPrevNode())
-            return isa<cir::EndCatchOp>(prev);
-          // The yield is alone in its block. Walk backward through
-          // single-predecessor blocks that contain only a branch.
-          mlir::Block *b = block.getSinglePredecessor();
-          while (b) {
-            mlir::Operation *term = b->getTerminator();
-            if (mlir::Operation *prev = term->getPrevNode())
-              return isa<cir::EndCatchOp>(prev);
-            if (!isa<cir::BrOp>(term))
-              return false;
-            b = b->getSinglePredecessor();
-          }
-          return false;
-        }() && "expected end_catch as last operation before yield "
-               "in catch handler, with only branches in between");
         rewriter.setInsertionPoint(yieldOp);
         rewriter.replaceOpWithNewOp<cir::BrOp>(yieldOp, continueBlock);
       }
diff --git a/clang/test/CIR/Transforms/flatten-try-catch-cleanup.cir b/clang/test/CIR/Transforms/flatten-try-catch-cleanup.cir
new file mode 100644
index 0000000000000..2d485c7bb5a9e
--- /dev/null
+++ b/clang/test/CIR/Transforms/flatten-try-catch-cleanup.cir
@@ -0,0 +1,39 @@
+// RUN: cir-opt %s -cir-flatten-cfg -o %t.cir
+// RUN: FileCheck --input-file=%t.cir %s
+
+// Regression test: catch handler with operations between end_catch
+// and yield (e.g., cleanup code after catching an exception) must
+// not trigger an assertion during TryOp flattening.
+
+!s32i = !cir.int<s, 32>
+!void = !cir.void
+
+cir.func private @may_throw()
+cir.func private @cleanup()
+
+cir.func @test_catch_with_cleanup() {
+  %0 = cir.alloca !s32i, !cir.ptr<!s32i>, ["x", init] {alignment = 4 : i64}
+  cir.scope {
+    cir.try {
+      cir.call @may_throw() : () -> ()
+      cir.yield
+    } catch all (%eh_token : !cir.eh_token) {
+      %catch_token, %exn_ptr = cir.begin_catch %eh_token : !cir.eh_token -> (!cir.catch_token, !cir.ptr<!cir.void>)
+      cir.end_catch %catch_token : !cir.catch_token
+      cir.call @cleanup() : () -> ()
+      %1 = cir.const #cir.int<42> : !s32i
+      cir.store %1, %0 : !s32i, !cir.ptr<!s32i>
+      cir.yield
+    }
+  }
+  cir.return
+}
+
+// CHECK-LABEL: cir.func @test_catch_with_cleanup()
+// CHECK:         cir.try_call @may_throw()
+// CHECK:       ^{{bb[0-9]+}}(%{{.*}}: !cir.eh_token):
+// CHECK:         cir.end_catch
+// CHECK:         cir.call @cleanup()
+// CHECK:         %[[C42:.*]] = cir.const #cir.int<42> : !s32i
+// CHECK:         cir.store %[[C42]],
+// CHECK:         cir.br