[Mlir-commits] [mlir] e84aa59 - Fix Block::eraseArgument when block arg is also a successor operand.

[Sean Silva]

Author: Sean Silva
Date: 2020-02-19T09:25:06-08:00
New Revision: e84aa5922b1cc8e40f60a560c3803122d820c53d

URL: https://github.com/llvm/llvm-project/commit/e84aa5922b1cc8e40f60a560c3803122d820c53d
DIFF: https://github.com/llvm/llvm-project/commit/e84aa5922b1cc8e40f60a560c3803122d820c53d.diff

LOG: Fix Block::eraseArgument when block arg is also a successor operand.

Summary:
This could trigger an assertion due to the block argument being used by
this block's own successor operands.

Reviewers: rriddle!

Subscribers: mehdi_amini, jpienaar, burmako, shauheen, antiagainst, nicolasvasilache, arpith-jacob, mgester, lucyrfox, liufengdb, Joonsoo, llvm-commits

Tags: #llvm

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

Added: 
    

Modified: 
    mlir/lib/IR/Block.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/Block.cpp b/mlir/lib/IR/Block.cpp
index e001c1000886..3e150614fb9e 100644
--- a/mlir/lib/IR/Block.cpp
+++ b/mlir/lib/IR/Block.cpp
@@ -170,20 +170,21 @@ BlockArgument Block::insertArgument(unsigned index, Type type) {
 void Block::eraseArgument(unsigned index, bool updatePredTerms) {
   assert(index < arguments.size());
 
+  // If requested, update predecessors. We do this first since this block might
+  // be a predecessor of itself and use this block argument as a successor
+  // operand.
+  if (updatePredTerms) {
+    // Erase this argument from each of the predecessor's terminator.
+    for (auto predIt = pred_begin(), predE = pred_end(); predIt != predE;
+         ++predIt) {
+      auto *predTerminator = (*predIt)->getTerminator();
+      predTerminator->eraseSuccessorOperand(predIt.getSuccessorIndex(), index);
+    }
+  }
+
   // Delete the argument.
   arguments[index].destroy();
   arguments.erase(arguments.begin() + index);
-
-  // If we aren't updating predecessors, there is nothing left to do.
-  if (!updatePredTerms)
-    return;
-
-  // Erase this argument from each of the predecessor's terminator.
-  for (auto predIt = pred_begin(), predE = pred_end(); predIt != predE;
-       ++predIt) {
-    auto *predTerminator = (*predIt)->getTerminator();
-    predTerminator->eraseSuccessorOperand(predIt.getSuccessorIndex(), index);
-  }
 }
 
 /// Insert one value to the given position of the argument list. The existing