[MLIR] Fix walk() after PostOrderTraversal change

[aengelke]

make_early_inc_range doesn't keep the range alive, only the iterators.
This breaks with the recent PostOrderTraversal change, which no longer
stores the state in the iterators. Store the range in a variable to keep
it alive for the entire loop.

Fixup of #191047 / 691a130.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Alexis Engelke (aengelke)

Changes

make_early_inc_range doesn't keep the range alive, only the iterators.
This breaks with the recent PostOrderTraversal change, which no longer
stores the state in the iterators. Store the range in a variable to keep
it alive for the entire loop.

Fixup of #191047 / 691a130.

---

Full diff: https://github.com/llvm/llvm-project/pull/191357.diff

1 Files Affected:

• (modified) mlir/include/mlir/IR/Visitors.h (+6-4)

diff --git a/mlir/include/mlir/IR/Visitors.h b/mlir/include/mlir/IR/Visitors.h
index 893f66ae33deb..907f470c0d248 100644
--- a/mlir/include/mlir/IR/Visitors.h
+++ b/mlir/include/mlir/IR/Visitors.h
@@ -120,8 +120,9 @@ void walk(Operation *op, function_ref<void(Block *)> callback,
           WalkOrder order) {
   for (auto &region : Iterator::makeIterable(*op)) {
     // Early increment here in the case where the block is erased.
-    for (auto &block :
-         llvm::make_early_inc_range(Iterator::makeIterable(region))) {
+    // PostOrderTraversal keeps state outside of iterators, so store it here.
+    auto &&It = Iterator::makeIterable(region);
+    for (auto &block : llvm::make_early_inc_range(It)) {
       if (order == WalkOrder::PreOrder)
         callback(&block);
       for (auto &nestedOp : Iterator::makeIterable(block))
@@ -195,8 +196,9 @@ WalkResult walk(Operation *op, function_ref<WalkResult(Block *)> callback,
                 WalkOrder order) {
   for (auto &region : Iterator::makeIterable(*op)) {
     // Early increment here in the case where the block is erased.
-    for (auto &block :
-         llvm::make_early_inc_range(Iterator::makeIterable(region))) {
+    // PostOrderTraversal keeps state outside of iterators, so store it here.
+    auto &&It = Iterator::makeIterable(region);
+    for (auto &block : llvm::make_early_inc_range(It)) {
       if (order == WalkOrder::PreOrder) {
         WalkResult result = callback(&block);
         if (result.wasSkipped())

[aengelke]

(For reference: using C++23 (P2718R0) would fix this problem as well, but it'll be a long time until we can use C++23...)