diff --git a/llvm/include/llvm/IR/Function.h b/llvm/include/llvm/IR/Function.h --- a/llvm/include/llvm/IR/Function.h +++ b/llvm/include/llvm/IR/Function.h @@ -719,10 +719,16 @@ /// \Returns \p ToIt. Function::iterator erase(Function::iterator FromIt, Function::iterator ToIt); + /// Erases \p BB fromt he basic block list. \Returns the iterator of the basic + /// block following \p BB. + Function::iterator erase(BasicBlock *BB) { + auto BBIt = BB->getIterator(); + return erase(BBIt, std::next(BBIt)); + } + private: // These need access to the underlying BB list. friend void BasicBlock::removeFromParent(); - friend iplist::iterator BasicBlock::eraseFromParent(); template friend class InstIterator; friend class llvm::SymbolTableListTraits; diff --git a/llvm/lib/IR/BasicBlock.cpp b/llvm/lib/IR/BasicBlock.cpp --- a/llvm/lib/IR/BasicBlock.cpp +++ b/llvm/lib/IR/BasicBlock.cpp @@ -130,7 +130,7 @@ } iplist::iterator BasicBlock::eraseFromParent() { - return getParent()->getBasicBlockList().erase(getIterator()); + return getParent()->erase(this); } void BasicBlock::moveBefore(BasicBlock *MovePos) { diff --git a/llvm/unittests/IR/FunctionTest.cpp b/llvm/unittests/IR/FunctionTest.cpp --- a/llvm/unittests/IR/FunctionTest.cpp +++ b/llvm/unittests/IR/FunctionTest.cpp @@ -486,4 +486,54 @@ It = F->erase(F->begin(), F->end()); EXPECT_EQ(F->size(), 0u); } + +TEST(FunctionTest, EraseSingleBB) { + LLVMContext Ctx; + std::unique_ptr M = parseIR(Ctx, R"( + define void @foo() { + bb1: + br label %bb2 + bb2: + br label %bb3 + bb3: + br label %bb4 + bb4: + br label %bb5 + bb5: + ret void + } +)"); + + Function *F = M->getFunction("foo"); + BasicBlock *BB1 = getBBWithName(F, "bb1"); + BasicBlock *BB2 = getBBWithName(F, "bb2"); + BasicBlock *BB3 = getBBWithName(F, "bb3"); + BasicBlock *BB4 = getBBWithName(F, "bb4"); + BasicBlock *BB5 = getBBWithName(F, "bb5"); + EXPECT_EQ(F->size(), 5u); + + // Erase BB2. + BB1->getTerminator()->eraseFromParent(); + auto It = F->erase(BB2); + EXPECT_EQ(F->size(), 4u); + // Check that the iterator returned matches the node after the erased one. + EXPECT_EQ(It, BB3->getIterator()); + + It = F->begin(); + EXPECT_EQ(&*It++, BB1); + EXPECT_EQ(&*It++, BB3); + EXPECT_EQ(&*It++, BB4); + EXPECT_EQ(&*It++, BB5); + + // Erase BB5. + BB4->getTerminator()->eraseFromParent(); + It = F->erase(BB5); + EXPECT_EQ(F->size(), 3u); + EXPECT_EQ(It, F->end()); + It = F->begin(); + EXPECT_EQ(&*It++, BB1); + EXPECT_EQ(&*It++, BB3); + EXPECT_EQ(&*It++, BB4); +} + } // end namespace