[CIR][FlattenCFG] Fix use after free when flattening terminator (#843)

Per the operation walking documentation [1]:

> A callback on a block or operation is allowed to erase that block or
> operation if either:
>   * the walk is in post-order, or
> * the walk is in pre-order and the walk is skipped after the erasure.

We were doing neither when erasing terminator operations and replacing
them with a branch, leading to a use after free and ASAN errors. This
fixes the following tests with ASAN:

```
  Clang :: CIR/CodeGen/switch-gnurange.cpp
  Clang :: CIR/Lowering/atomic-runtime.cpp
  Clang :: CIR/Lowering/loop.cir
  Clang :: CIR/Lowering/loops-with-break.cir
  Clang :: CIR/Lowering/loops-with-continue.cir
  Clang :: CIR/Lowering/switch.cir
  Clang :: CIR/Transforms/Target/x86_64/x86_64-call-conv-lowering-pass.cpp
  Clang :: CIR/Transforms/loop.cir
  Clang :: CIR/Transforms/switch.cir
```

These two tests still fail with ASAN after this, which I'm looking into:
```
  Clang :: CIR/CodeGen/pointer-arith-ext.c
  Clang :: CIR/Transforms/Target/x86_64/x86_64-call-conv-lowering-pass.cpp
```

`CIR/CodeGen/global-new.cpp` is failing even on a non-ASAN Release build
for me on the parent commit, so it's unrelated.

[1]
https://github.com/llvm/llvm-project/blob/0c55ad11ab3857056bb3917fdf087c4aa811b790/mlir/include/mlir/IR/Operation.h#L767-L770

Author: smeenai

clang/include/clang/CIR/Interfaces/CIRLoopOpInterface.td

@@ -71,14 +71,13 @@ def LoopOpInterface : OpInterface<"LoopOpInterface", [
       }],
       /*retTy=*/"mlir::WalkResult",
       /*methodName=*/"walkBodySkippingNestedLoops",
-      /*args=*/(ins "::llvm::function_ref<void (Operation *)>":$callback),
+      /*args=*/(ins "::llvm::function_ref<mlir::WalkResult (Operation *)>":$callback),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.getBody().template walk<WalkOrder::PreOrder>([&](Operation *op) {
           if (isa<LoopOpInterface>(op))
             return mlir::WalkResult::skip();
-          callback(op);
-          return mlir::WalkResult::advance();
+          return callback(op);
         });
       }]
     >

clang/lib/CIR/Dialect/Transforms/FlattenCFG.cpp

@@ -36,13 +36,13 @@ void lowerTerminator(mlir::Operation *op, mlir::Block *dest,
 /// Walks a region while skipping operations of type `Ops`. This ensures the
 /// callback is not applied to said operations and its children.
 template <typename... Ops>
-void walkRegionSkipping(mlir::Region &region,
-                        mlir::function_ref<void(mlir::Operation *)> callback) {
+void walkRegionSkipping(
+    mlir::Region &region,
+    mlir::function_ref<mlir::WalkResult(mlir::Operation *)> callback) {
   region.walk<mlir::WalkOrder::PreOrder>([&](mlir::Operation *op) {
     if (isa<Ops...>(op))
       return mlir::WalkResult::skip();
-    callback(op);
-    return mlir::WalkResult::advance();
+    return callback(op);
   });
 }
 
@@ -541,15 +541,21 @@ class CIRLoopOpInterfaceFlattening
     // Lower continue statements.
     mlir::Block *dest = (step ? step : cond);
     op.walkBodySkippingNestedLoops([&](mlir::Operation *op) {
-      if (isa<mlir::cir::ContinueOp>(op))
-        lowerTerminator(op, dest, rewriter);
+      if (!isa<mlir::cir::ContinueOp>(op))
+        return mlir::WalkResult::advance();
+
+      lowerTerminator(op, dest, rewriter);
+      return mlir::WalkResult::skip();
     });
 
     // Lower break statements.
     walkRegionSkipping<mlir::cir::LoopOpInterface, mlir::cir::SwitchOp>(
         op.getBody(), [&](mlir::Operation *op) {
-          if (isa<mlir::cir::BreakOp>(op))
-            lowerTerminator(op, exit, rewriter);
+          if (!isa<mlir::cir::BreakOp>(op))
+            return mlir::WalkResult::advance();
+
+          lowerTerminator(op, exit, rewriter);
+          return mlir::WalkResult::skip();
         });
 
     // Lower optional body region yield.
@@ -705,8 +711,11 @@ class CIRSwitchOpFlattening
       // Handle break statements.
       walkRegionSkipping<mlir::cir::LoopOpInterface, mlir::cir::SwitchOp>(
           region, [&](mlir::Operation *op) {
-            if (isa<mlir::cir::BreakOp>(op))
-              lowerTerminator(op, exitBlock, rewriter);
+            if (!isa<mlir::cir::BreakOp>(op))
+              return mlir::WalkResult::advance();
+
+            lowerTerminator(op, exitBlock, rewriter);
+            return mlir::WalkResult::skip();
           });
 
       // Extract region contents before erasing the switch op.