From 014575932fc3acbb0f1c5b46ff2bfcebf69e6e62 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph@gmail.com>
Date: Sat, 27 Feb 2021 19:18:09 +0000
Subject: [PATCH] Fix Block::eraseArguments: keep track the first removed
 element while removing

Not only this is likely more efficient than BitVector::find_first(), but
also if the BitVector is empty find_first() returns -1, which
llvm::drop_begin isn't robust against.
---
 mlir/lib/IR/Block.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/IR/Block.cpp b/mlir/lib/IR/Block.cpp
index 0758d1ecb647..53797aa49fd2 100644
--- a/mlir/lib/IR/Block.cpp
+++ b/mlir/lib/IR/Block.cpp
@@ -192,15 +192,17 @@ void Block::eraseArguments(llvm::BitVector eraseIndices) {
   // We do this in reverse so that we erase later indices before earlier
   // indices, to avoid shifting the later indices.
   unsigned originalNumArgs = getNumArguments();
+  int64_t firstErased = originalNumArgs;
   for (unsigned i = 0; i < originalNumArgs; ++i) {
     int64_t currentPos = originalNumArgs - i - 1;
     if (eraseIndices.test(currentPos)) {
       arguments[currentPos].destroy();
       arguments.erase(arguments.begin() + currentPos);
+      firstErased = currentPos;
     }
   }
   // Update the cached position for the arguments after the first erased one.
-  int64_t index = eraseIndices.find_first();
+  int64_t index = firstErased;
   for (BlockArgument arg : llvm::drop_begin(arguments, index))
     arg.setArgNumber(index++);
 }