[Mlir-commits] [mlir] [MLIR] Prevent invalid IR from being passed outside of RemoveDeadValues (PR #121079)

[Andrzej Warzyński]

================
@@ -588,22 +648,60 @@ static void cleanBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la) {
       operandValues.push_back(successorOperands[operandIdx]);
     }
 
-    BitVector successorLiveOperands = markLives(operandValues, la);
-
     // Do (3)
-    for (int argIdx = successorLiveOperands.size() - 1; argIdx >= 0; --argIdx) {
-      if (!successorLiveOperands[argIdx]) {
-        if (successorBlock->getNumArguments() < successorOperands.size()) {
-          // if block was cleaned through a different code path
-          // we only need to remove operands from the invokation
-          successorOperands.erase(argIdx);
-          continue;
-        }
+    BitVector successorNonLive =
+        markLives(operandValues, deletionSet, la).flip();
+    updateDeletionSet(deletionSet, successorBlock->getArguments(),
+                      successorNonLive);
+    cl.blocks.push_back({successorBlock, successorNonLive});
+    cl.successorOperands.push_back({branchOp, succIdx, successorNonLive});
+  }
+}
+
+void cleanup(CleanupList &cl) {
+  for (auto &op : cl.operations) {
----------------
banach-space wrote:

[nit] Same suggestion for other `for` loops in this method. The code is quite self-explanatory, hence leaving this as a [nit].
```suggestion
  // 1. Operations
  for (auto &op : cl.operations) {
```

https://github.com/llvm/llvm-project/pull/121079