[BOLT] Fix CFG in BinaryFunction::eraseInvalidBBs()

Summary:
When we erase invalid/unreachable basic blocks, we have to remove them
from a list of predecessors of regular blocks, otherwise the CFG will be
left in a broken state containing references to removed basic blocks.

(cherry picked from FBD7464292)

bolt/BinaryBasicBlock.cpp

@@ -271,6 +271,14 @@ void BinaryBasicBlock::replaceSuccessor(BinaryBasicBlock *Succ,
   NewSucc->addPredecessor(this);
 }
 
+void BinaryBasicBlock::removeAllSuccessors() {
+  for (auto *SuccessorBB : successors()) {
+    SuccessorBB->removePredecessor(this);
+  }
+  Successors.clear();
+  BranchInfo.clear();
+}
+
 void BinaryBasicBlock::removeSuccessor(BinaryBasicBlock *Succ) {
   Succ->removePredecessor(this);
   auto I = succ_begin();
@@ -292,9 +300,17 @@ void BinaryBasicBlock::addPredecessor(BinaryBasicBlock *Pred) {
 }
 
 void BinaryBasicBlock::removePredecessor(BinaryBasicBlock *Pred) {
-  auto I = std::find(pred_begin(), pred_end(), Pred);
-  assert(I != pred_end() && "Pred is not a predecessor of this block!");
-  Predecessors.erase(I);
+  // Note: the predecessor could be listed multiple times.
+  bool Erased{false};
+  for (auto PredI = Predecessors.begin(); PredI != Predecessors.end(); ) {
+    if (*PredI == Pred) {
+      Erased = true;
+      PredI = Predecessors.erase(PredI);
+    } else {
+      ++PredI;
+    }
+  }
+  assert(Erased && "Pred is not a predecessor of this block!");
 }
 
 void BinaryBasicBlock::removeDuplicateConditionalSuccessor(MCInst *CondBranch) {

bolt/BinaryBasicBlock.h

@@ -544,13 +544,9 @@ public:
   /// list of predecessors of /p Succ and update branch info.
   void removeSuccessor(BinaryBasicBlock *Succ);
 
-  /// Remove a range of successor blocks.
-  template <typename Itr>
-  void removeSuccessors(Itr Begin, Itr End) {
-    while (Begin != End) {
-      removeSuccessor(*Begin++);
-    }
-  }
+  /// Remove all successors of the basic block, and remove the block
+  /// from respective lists of predecessors.
+  void removeAllSuccessors();
 
   /// Remove useless duplicate successors.  When the conditional
   /// successor is the same as the unconditional successor, we can

bolt/BinaryFunction.cpp

@@ -294,7 +294,7 @@ BinaryFunction::getBasicBlockContainingOffset(uint64_t Offset) {
   return (Offset < BB->getOffset() + BB->getOriginalSize()) ? BB : nullptr;
 }
 
-void BinaryFunction::markUnreachable() {
+void BinaryFunction::markUnreachableBlocks() {
   std::stack<BinaryBasicBlock *> Stack;
 
   for (auto *BB : layout()) {
@@ -342,10 +342,13 @@ std::pair<unsigned, uint64_t> BinaryFunction::eraseInvalidBBs() {
 
   BasicBlockListType NewBasicBlocks;
   for (auto I = BasicBlocks.begin(), E = BasicBlocks.end(); I != E; ++I) {
-    if ((*I)->isValid()) {
-      NewBasicBlocks.push_back(*I);
+    auto *BB = *I;
+    if (BB->isValid()) {
+      NewBasicBlocks.push_back(BB);
     } else {
-      DeletedBasicBlocks.push_back(*I);
+      // Make sure the block is removed from the list of predecessors.
+      BB->removeAllSuccessors();
+      DeletedBasicBlocks.push_back(BB);
     }
   }
   BasicBlocks = std::move(NewBasicBlocks);
@@ -2602,6 +2605,33 @@ bool BinaryFunction::validateCFG() const {
   if (!Valid)
     return Valid;
 
+  // Make sure all blocks in CFG are valid.
+  auto validateBlock = [this](const BinaryBasicBlock *BB, StringRef Desc) {
+    if (!BB->isValid()) {
+      errs() << "BOLT-ERROR: deleted " << Desc << " " << BB->getName()
+             << " detected in:\n";
+      this->dump();
+      return false;
+    }
+    return true;
+  };
+  for (const auto *BB : BasicBlocks) {
+    if (!validateBlock(BB, "block"))
+      return false;
+    for (const auto *PredBB : BB->predecessors())
+      if (!validateBlock(PredBB, "predecessor"))
+        return false;
+    for (const auto *SuccBB: BB->successors())
+      if (!validateBlock(SuccBB, "successor"))
+        return false;
+    for (const auto *LP: BB->landing_pads())
+      if (!validateBlock(LP, "landing pad"))
+        return false;
+    for (const auto *Thrower: BB->throwers())
+      if (!validateBlock(Thrower, "thrower"))
+        return false;
+  }
+
   for (const auto *BB : BasicBlocks) {
     std::unordered_set<const BinaryBasicBlock *> BBLandingPads;
     for (const auto *LP : BB->landing_pads()) {

bolt/BinaryFunction.h

@@ -1367,7 +1367,7 @@ public:
 
   /// Mark all blocks that are unreachable from a root (entry point
   /// or landing pad) as invalid.
-  void markUnreachable();
+  void markUnreachableBlocks();
 
   /// Rebuilds BBs layout, ignoring dead BBs. Returns the number of removed
   /// BBs and the removed number of bytes of code.

bolt/Passes/BinaryPasses.cpp

@@ -326,7 +326,7 @@ void EliminateUnreachableBlocks::runOnFunction(BinaryFunction& Function) {
   if (Function.layout_size() > 0) {
     unsigned Count;
     uint64_t Bytes;
-    Function.markUnreachable();
+    Function.markUnreachableBlocks();
     DEBUG({
       for (auto *BB : Function.layout()) {
         if (!BB->isValid()) {
@@ -339,14 +339,16 @@ void EliminateUnreachableBlocks::runOnFunction(BinaryFunction& Function) {
     std::tie(Count, Bytes) = Function.eraseInvalidBBs();
     DeletedBlocks += Count;
     DeletedBytes += Bytes;
-    if (Count && opts::Verbosity > 0) {
+    if (Count) {
       Modified.insert(&Function);
+      if (opts::Verbosity > 0) {
         outs() << "BOLT-INFO: Removed " << Count
                << " dead basic block(s) accounting for " << Bytes
                << " bytes in function " << Function << '\n';
       }
     }
   }
+}
 
 void EliminateUnreachableBlocks::runOnFunctions(
   BinaryContext&,
@@ -765,7 +767,6 @@ uint64_t fixDoubleJumps(BinaryContext &BC,
                        BB.isLandingPad() ||
                        BB.isEntryPoint());
         }
-        assert(Function.validateCFG());
       }
     }
   }
@@ -819,7 +820,7 @@ uint64_t SimplifyConditionalTailCalls::fixTailCalls(BinaryContext &BC,
                                                     BinaryFunction &BF) {
   // Need updated indices to correctly detect branch' direction.
   BF.updateLayoutIndices();
-  BF.markUnreachable();
+  BF.markUnreachableBlocks();
 
   auto &MIB = BC.MIB;
   uint64_t NumLocalCTCCandidates = 0;
@@ -986,6 +987,8 @@ uint64_t SimplifyConditionalTailCalls::fixTailCalls(BinaryContext &BC,
     const auto Stats = BF.eraseInvalidBBs();
     DeletedBlocks += Stats.first;
     DeletedBytes += Stats.second;
+
+    assert(BF.validateCFG());
   }
 
   DEBUG(dbgs() << "BOLT: created " << NumLocalCTCs
@@ -1120,6 +1123,7 @@ void Peepholes::runOnFunctions(BinaryContext &BC,
         addTailcallTraps(BC, Function);
       if (Opts & opts::PEEP_USELESS_BRANCHES)
         removeUselessCondBranches(BC, Function);
+      assert(Function.validateCFG());
     }
   }
   outs() << "BOLT-INFO: Peephole: " << NumShortened
@@ -1275,6 +1279,8 @@ void IdenticalCodeFolding::runOnFunctions(BinaryContext &BC,
     BF.hash(/*Recompute=*/true, /*UseDFS=*/UseDFS);
 
     CongruentBuckets[&BF].emplace(&BF);
+
+    dbgs() << BF.getPrintName() << " : " << BF.getKnownExecutionCount() << "\n";
   }
 
   // We repeat the pass until no new modifications happen.

bolt/Passes/IndirectCallPromotion.cpp

@@ -736,20 +736,12 @@ BinaryBasicBlock *IndirectCallPromotion::fixCFG(
   BinaryBasicBlock *MergeBlock = nullptr;
 
   auto moveSuccessors = [](BinaryBasicBlock *Old, BinaryBasicBlock *New) {
-    std::vector<BinaryBasicBlock*> OldSucc(Old->successors().begin(),
-                                           Old->successors().end());
-    std::vector<BinaryBranchInfo> BranchInfo(Old->branch_info_begin(),
+    // Copy over successors to the new block.
+    New->addSuccessors(Old->successors().begin(),
+                       Old->successors().end(),
+                       Old->branch_info_begin(),
                        Old->branch_info_end());
-
-    // Remove all successors from the old block.
-    Old->removeSuccessors(OldSucc.begin(), OldSucc.end());
-    assert(Old->succ_empty());
-
-    // Move them to the new block.
-    New->addSuccessors(OldSucc.begin(),
-                       OldSucc.end(),
-                       BranchInfo.begin(),
-                       BranchInfo.end());
+    Old->removeAllSuccessors();
 
     // Update the execution count on the new block.
     New->setExecutionCount(Old->getExecutionCount());

bolt/Passes/Inliner.cpp

@@ -362,7 +362,7 @@ InlineSmallFunctions::inlineCall(
   // If the inlining cannot happen as a simple instruction insertion into
   // CallerBB, we remove the outgoing CFG edges of the caller block.
   if (InlinedInstance.size() > 1 || !CanMergeFirstInlinedBlock) {
-    CallerBB->removeSuccessors(CallerBB->succ_begin(), CallerBB->succ_end());
+    CallerBB->removeAllSuccessors();
     if (!ShouldSplitCallerBB) {
       // Update the after-inlined point.
       AfterInlinedBB = CallerFunction.getBasicBlockAfter(CallerBB);
@@ -402,7 +402,7 @@ InlineSmallFunctions::inlineCall(
             FirstBB->succ_begin(),
             FirstBB->succ_end());
       }
-      FirstBB->removeSuccessors(FirstBB->succ_begin(), FirstBB->succ_end());
+      FirstBB->removeAllSuccessors();
     }
     InlinedInstance.erase(InlinedInstance.begin());
   } else {

bolt/Passes/LongJmp.cpp

@@ -328,7 +328,7 @@ void LongJmpPass::removeStubRef(const BinaryContext &BC,
 
   if (StubRefCount[StubBB] == 0) {
     // Remove the block from CFG
-    StubBB->removeSuccessors(StubBB->succ_begin(), StubBB->succ_end());
+    StubBB->removeAllSuccessors();
     StubBB->markValid(false);
     StubBB->setEntryPoint(false);
   }