[mlir][DialectConversion] Remove usage of std::distance to track position.

[mravishankar]

Summary

Remove use of iterator::difference_type to know where to insert a
moved or erased block during undo actions.

Depends On D84795

[rriddle]

mlir/lib/Transforms/DialectConversion.cpp

992–996

nit: Can you assign insertAfterBlock to a temporary variable? Would likely help the formatting here.

1018–1019

Same here.

1198

block->getPrevNode()?

[mravishankar]

Addressing comments

mlir/lib/Transforms/DialectConversion.cpp

1198

Oh thanks!

[mravishankar]

Rebase

mlir/lib/Transforms/DialectConversion.cpp

enum class BlockActionKind {
   Create,
   Erase,
   Merge,
   Move,
   Split,
   TypeConversion
 };
 
-/// Original position of the given block in its parent region.  We cannot use
-/// a region iterator because it could have been invalidated by other region
-/// operations since the position was stored.
+/// Original position of the given block in its parent region. During undo
+/// actions, the block needs to be placed after `insertAfterBlock`.
 struct BlockPosition {
   Region *region;
-  Region::iterator::difference_type position;
+  Block *insertAfterBlock;
 };
 
 /// Information needed to undo the merge actions.
 /// - the source block, and
 /// - the Operation that was the last operation in the dest block before the
 ///   merge (could be null if the dest block was empty).
 struct MergeInfo {
   Block *sourceBlock;
   Operation *destBlockLastInst;
 };
 
 /// The storage class for an undoable block action (one of BlockActionKind),
 /// contains the information necessary to undo this action.
 struct BlockAction {
   static BlockAction getCreate(Block *block) {
     return {BlockActionKind::Create, block, {}};
   }
-  static BlockAction getErase(Block *block, BlockPosition originalPos) {
-    return {BlockActionKind::Erase, block, {originalPos}};
+  static BlockAction getErase(Block *block, BlockPosition originalPosition) {
+    return {BlockActionKind::Erase, block, {originalPosition}};
   }
   static BlockAction getMerge(Block *block, Block *sourceBlock) {
     BlockAction action{BlockActionKind::Merge, block, {}};
     action.mergeInfo = {sourceBlock, block->begin() == block->end()
                                          ? nullptr
                                          : &(block->back())};
     return action;
   }
-  static BlockAction getMove(Block *block, BlockPosition originalPos) {
-    return {BlockActionKind::Move, block, {originalPos}};
+  static BlockAction getMove(Block *block, BlockPosition originalPosition) {
+    return {BlockActionKind::Move, block, {originalPosition}};
   }
   static BlockAction getSplit(Block *block, Block *originalBlock) {
     BlockAction action{BlockActionKind::Split, block, {}};
     action.originalBlock = originalBlock;
     return action;
   }
   static BlockAction getTypeConversion(Block *block) {
     return BlockAction{BlockActionKind::TypeConversion, block, {}};
case BlockActionKind::Create: {
         blockOps.remove(blockOps.begin());
       action.block->dropAllDefinedValueUses();
       action.block->erase();
       break;
     }
     // Put the block (owned by action) back into its original position.
     case BlockActionKind::Erase: {
       auto &blockList = action.originalPosition.region->getBlocks();
-      blockList.insert(
-          std::next(blockList.begin(), action.originalPosition.position),
+      Block *insertAfterBlock = action.originalPosition.insertAfterBlock;
+      blockList.insert((insertAfterBlock
+                            ? std::next(Region::iterator(insertAfterBlock))
+                            : blockList.end()),
                        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;
       Block::iterator splitPoint =
           (action.mergeInfo.destBlockLastInst
                ? ++Block::iterator(action.mergeInfo.destBlockLastInst)
                : action.block->begin());
       sourceBlock->getOperations().splice(sourceBlock->begin(),
                                           action.block->getOperations(),
                                           splitPoint, action.block->end());
       break;
     }
     // Move the block back to its original position.
     case BlockActionKind::Move: {
       Region *originalRegion = action.originalPosition.region;
+      Block *insertAfterBlock = action.originalPosition.insertAfterBlock;
       originalRegion->getBlocks().splice(
-          std::next(originalRegion->begin(), action.originalPosition.position),
+          (insertAfterBlock ? std::next(Region::iterator(insertAfterBlock))
+                            : originalRegion->end()),
           action.block->getParent()->getBlocks(), action.block);
       break;
     }
     // Merge back the block that was split out.
     case BlockActionKind::Split: {
       action.originalBlock->getOperations().splice(
           action.originalBlock->end(), action.block->getOperations());
       action.block->dropAllDefinedValueUses();
void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,
 
   // Mark this operation as recursively ignored so that we don't need to
   // convert any nested operations.
   markNestedOpsIgnored(op);
 }
 
 void ConversionPatternRewriterImpl::notifyBlockIsBeingErased(Block *block) {
   Region *region = block->getParent();
-  auto position = std::distance(region->begin(), Region::iterator(block));
-  blockActions.push_back(BlockAction::getErase(block, {region, position}));
+  Block *origPrevBlock = block->getPrevNode();
+  blockActions.push_back(BlockAction::getErase(block, {region, origPrevBlock}));
 }
 
 void ConversionPatternRewriterImpl::notifyCreatedBlock(Block *block) {
   blockActions.push_back(BlockAction::getCreate(block));
 }
 
 void ConversionPatternRewriterImpl::notifySplitBlock(Block *block,
                                                      Block *continuation) {
   blockActions.push_back(BlockAction::getSplit(continuation, block));
 }
 
 void ConversionPatternRewriterImpl::notifyBlocksBeingMerged(Block *block,
                                                             Block *srcBlock) {
   blockActions.push_back(BlockAction::getMerge(block, srcBlock));
 }
 
 void ConversionPatternRewriterImpl::notifyRegionIsBeingInlinedBefore(
     Region &region, Region &parent, Region::iterator before) {
+  Block *origPrevBlock = nullptr;
   for (auto &pair : llvm::enumerate(region)) {
     Block &block = pair.value();
-    Region::iterator::difference_type position = pair.index();
-    blockActions.push_back(BlockAction::getMove(&block, {&region, position}));
+    blockActions.push_back(
+        BlockAction::getMove(&block, {&region, origPrevBlock}));
+    origPrevBlock = █
   }
 }
 
 void ConversionPatternRewriterImpl::notifyRegionWasClonedBefore(
     iterator_range<Region::iterator> &blocks, Location origRegionLoc) {
   for (Block &block : blocks)
     blockActions.push_back(BlockAction::getCreate(&block));