[mlir] Rework DialectConversion inlineRegionBefore

The previous logic for inlining a region A with N blocks into region B
would produce incorrect results on rollback for N greater than 1. This
rollback logic would leave blocks 1..N in region B and only move block 0
to region A.

The new inlining action recording stores the block move actions from N-1
to 0. Now on roll back, block 0 is moved to region A and then 1..N is
appended to the list of blocks in region A.

Differential Revision: https://reviews.llvm.org/D91185

mlir/lib/Transforms/Utils/DialectConversion.cpp

@@ -1211,13 +1211,15 @@ void ConversionPatternRewriterImpl::notifyBlocksBeingMerged(Block *block,
 
 void ConversionPatternRewriterImpl::notifyRegionIsBeingInlinedBefore(
     Region &region, Region &parent, Region::iterator before) {
-  Block *origPrevBlock = nullptr;
-  for (auto &pair : llvm::enumerate(region)) {
-    Block &block = pair.value();
+  if (region.empty())
+    return;
+  Block *laterBlock = &region.back();
+  for (auto &earlierBlock : llvm::drop_begin(llvm::reverse(region), 1)) {
     blockActions.push_back(
-        BlockAction::getMove(&block, {&region, origPrevBlock}));
-    origPrevBlock = █
+        BlockAction::getMove(laterBlock, {&region, &earlierBlock}));
+    laterBlock = &earlierBlock;
   }
+  blockActions.push_back(BlockAction::getMove(laterBlock, {&region, nullptr}));
 }
 
 void ConversionPatternRewriterImpl::notifyRegionWasClonedBefore(

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

@@ -69,3 +69,18 @@ func @test_unknown_dynamically_legal() {
   "foo.unknown_op"() {} : () -> ()
   "test.return"() : () -> ()
 }
+
+// -----
+
+// Test that region inlining can be properly undone.
+func @test_undo_region_inline() {
+  "test.region"() ({
+    ^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"() : () -> ()
+}