[llvm] c2355b3 - [IR] Adds Function::erase() for erasing a range of basic blocks

[Vasileios Porpodas via llvm-commits]

Author: Vasileios Porpodas
Date: 2022-12-14T16:23:26-08:00
New Revision: c2355b361a0eb2432e1fe9017e2c589818fffb84

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

LOG: [IR] Adds Function::erase() for erasing a range of basic blocks

This is part of a series of patches that aim at making Function::getBasicBlockList() private.

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

Added: 
    

Modified: 
    llvm/include/llvm/IR/Function.h
    llvm/lib/IR/Function.cpp
    llvm/unittests/IR/FunctionTest.cpp

Removed: 
    


################################################################################
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<Module> 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