diff --git a/llvm/include/llvm/IR/Function.h b/llvm/include/llvm/IR/Function.h index ce0fb45a65d2d..995f797e2ec3a 100644 --- a/llvm/include/llvm/IR/Function.h +++ b/llvm/include/llvm/IR/Function.h @@ -715,6 +715,10 @@ class LLVM_EXTERNAL_VISIBILITY Function : public GlobalObject, Function::iterator FromBeginIt, Function::iterator FromEndIt); + /// Erases a range of BasicBlocks from \p FromIt to (not including) \p ToIt. + /// \Returns \p ToIt. + Function::iterator erase(Function::iterator FromIt, Function::iterator ToIt); + /// Get the underlying elements of the Function... the basic block list is /// empty for external functions. /// diff --git a/llvm/lib/IR/Function.cpp b/llvm/lib/IR/Function.cpp index 06826a80c614b..aaef0891bc61f 100644 --- a/llvm/lib/IR/Function.cpp +++ b/llvm/lib/IR/Function.cpp @@ -380,6 +380,11 @@ void Function::splice(Function::iterator ToIt, Function *FromF, BasicBlocks.splice(ToIt, FromF->BasicBlocks, FromBeginIt, FromEndIt); } +Function::iterator Function::erase(Function::iterator FromIt, + Function::iterator ToIt) { + return BasicBlocks.erase(FromIt, ToIt); +} + //===----------------------------------------------------------------------===// // Function Implementation //===----------------------------------------------------------------------===// diff --git a/llvm/unittests/IR/FunctionTest.cpp b/llvm/unittests/IR/FunctionTest.cpp index 8d8a2feb6e35a..d53578b38186c 100644 --- a/llvm/unittests/IR/FunctionTest.cpp +++ b/llvm/unittests/IR/FunctionTest.cpp @@ -443,4 +443,47 @@ TEST(FunctionTest, SpliceEndBeforeBegin) { "FromBeginIt not before FromEndIt!"); } #endif //EXPENSIVE_CHECKS + +TEST(FunctionTest, EraseBBs) { + 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->getIterator(), std::next(BB2->getIterator())); + 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 all BBs. + It = F->erase(F->begin(), F->end()); + EXPECT_EQ(F->size(), 0u); +} } // end namespace