From 56bf709cbbc16a77b055c84a15e4c696d8298b04 Mon Sep 17 00:00:00 2001
From: AmrDeveloper <amr96@programmer.net>
Date: Sat, 25 Jan 2025 21:37:46 +0100
Subject: [PATCH] [CIR] Cleanup cir.scopes with a single cir.yield operation

---
 clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp | 11 ++++++++++-
 clang/test/CIR/Transforms/merge-cleanups.cir         |  9 +++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp b/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp
index 555cb20408a5..f923ac69dbeb 100644
--- a/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp
+++ b/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp
@@ -62,7 +62,16 @@ struct RemoveEmptyScope : public OpRewritePattern<ScopeOp> {
   LogicalResult match(ScopeOp op) const final {
     // TODO: Remove this logic once CIR uses MLIR infrastructure to remove
     // trivially dead operations
-    return success(op.isEmpty());
+    if (op.isEmpty()) {
+      return success();
+    }
+
+    Region *region = op.getRegions().front();
+    if (region && region->getBlocks().front().getOperations().size() == 1) {
+      return success(isa<YieldOp>(region->getBlocks().front().front()));
+    }
+
+    return failure();
   }
 
   void rewrite(ScopeOp op, PatternRewriter &rewriter) const final {
diff --git a/clang/test/CIR/Transforms/merge-cleanups.cir b/clang/test/CIR/Transforms/merge-cleanups.cir
index 715c7525b94d..a111fec0b0b4 100644
--- a/clang/test/CIR/Transforms/merge-cleanups.cir
+++ b/clang/test/CIR/Transforms/merge-cleanups.cir
@@ -138,4 +138,13 @@ module  {
     cir.return %0 : !cir.ptr<!s32i, addrspace(target<2>)>
   }
 
+  // Should remove scope with only yield
+  cir.func @removeBlockWithScopeYeild(%arg0: !s32i) {
+    cir.scope {
+      cir.yield
+    }
+    cir.return
+  }
+  //      CHECK: cir.func @removeBlockWithScopeYeild
+  // CHECK-NEXT: cir.return
 }