Fix rollback of first block in a region.

[tpopp]

Summary

Previously the first block might end up at the end of the region on rollback.

[rriddle]

mlir/test/lib/Dialect/Test/TestPatterns.cpp

212

nit: Add braces around the if here.

mlir/lib/Transforms/Utils/DialectConversion.cpp

case BlockActionKind::Create: {
       break;
     }
     // Put the block (owned by action) back into its original position.
     case BlockActionKind::Erase: {
       auto &blockList = action.originalPosition.region->getBlocks();
       Block *insertAfterBlock = action.originalPosition.insertAfterBlock;
       blockList.insert((insertAfterBlock
                             ? std::next(Region::iterator(insertAfterBlock))
-                            : blockList.end()),
+                            : blockList.begin()),
                        action.block);
       break;
     }
     // Split the block at the position which was originally the end of the
     // destination block (owned by action), and put the instructions back into
     // the block used before the merge.
     case BlockActionKind::Merge: {
       Block *sourceBlock = action.mergeInfo.sourceBlock;

mlir/test/Transforms/test-legalizer-full.mlir

^bb1(%i0: i64):
        // expected-error@+1 {{failed to legalize operation 'std.br'}}
        br ^bb2(%i0 : i64)
     ^bb2(%i1: i64):
       "test.invalid"(%i1) : (i64) -> ()
   }) {} : () -> ()
 
   "test.return"() : () -> ()
 }
+
+// -----
+
+// Test that multiple block erases can be properly undone.
+func @test_undo_block_erase() {
+   // expected-error@+1 {{failed to legalize operation 'test.region'}}
+  "test.region"() ({
+    ^bb1(%i0: i64):
+       br ^bb2(%i0 : i64)
+    ^bb2(%i1: i64):
+      "test.invalid"(%i1) : (i64) -> ()
+  }) {legalizer.should_clone, legalizer.erase_old_blocks} : () -> ()
+
+  "test.return"() : () -> ()
+}

mlir/test/lib/Dialect/Test/TestPatterns.cpp

struct TestRegionRewriteBlockMovement : public ConversionPattern {
   TestRegionRewriteBlockMovement(MLIRContext *ctx)
       : ConversionPattern("test.region", 1, ctx) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     // Inline this region into the parent region.
     auto &parentRegion = *op->getParentRegion();
+    auto &opRegion = op->getRegion(0);
     if (op->getAttr("legalizer.should_clone"))
-      rewriter.cloneRegionBefore(op->getRegion(0), parentRegion,
-                                 parentRegion.end());
+      rewriter.cloneRegionBefore(opRegion, parentRegion, parentRegion.end());
     else
-      rewriter.inlineRegionBefore(op->getRegion(0), parentRegion,
-                                  parentRegion.end());
+      rewriter.inlineRegionBefore(opRegion, parentRegion, parentRegion.end());
+
+    if (op->getAttr("legalizer.erase_old_blocks")) {
+      while (!opRegion.empty())
+        rewriter.eraseBlock(&opRegion.front());
+    }
 
     // Drop this operation.
     rewriter.eraseOp(op);
     return success();
   }
 };
 /// This pattern is a simple pattern that generates a region containing an
 /// illegal operation.