[BOLT] Process each block only once in fixCFGForPIC

[FPar]

Summary

Rather than iterating over the whole function from the start until no
internal calls are found, process each block only once and continue
processing after splitting. This version of the function also does not
seemingly invalidate iterators from within the loop.

[maksfb]

Thanks for improving the efficiency of fixCFGForPIC(). The change looks fine to me. Do you think the implementation will be simpler if we make Blocks a queue and add newly added blocks to the end? This way we can get rid of the inner loop and dealing with CurrentBB/NextBB.

[FPar]

Eliminate inner loop through queue

[FPar]

In D130436#3674716, @maksfb wrote:

Thanks for improving the efficiency of fixCFGForPIC(). The change looks fine to me. Do you think the implementation will be simpler if we make Blocks a queue and add newly added blocks to the end? This way we can get rid of the inner loop and dealing with CurrentBB/NextBB.

That makes totally sense. It somehow slipped my mind.

As far as I can tell, there is also no need for setting annotation in this function anymore. Is that correct?

[maksfb]

LGTM. Good point, no need for the annotation.

bolt/include/bolt/Core/BinaryFunction.h

public:
 
   /// Return BinaryContext for the function.
   BinaryContext &getBinaryContext() { return BC; }
 
   /// Attempt to validate CFG invariants.
   bool validateCFG() const;
 
   BinaryBasicBlock *getBasicBlockForLabel(const MCSymbol *Label) {
-    auto I = LabelToBB.find(Label);
-    return I == LabelToBB.end() ? nullptr : I->second;
+    return LabelToBB.lookup(Label);
   }
 
   const BinaryBasicBlock *getBasicBlockForLabel(const MCSymbol *Label) const {
-    auto I = LabelToBB.find(Label);
-    return I == LabelToBB.end() ? nullptr : I->second;
+    return LabelToBB.lookup(Label);
   }
 
   /// Retrieve the landing pad BB associated with invoke instruction \p Invoke
   /// that is in \p BB. Return nullptr if none exists
   BinaryBasicBlock *getLandingPadBBFor(const BinaryBasicBlock &BB,
                                        const MCInst &InvokeInst) const {
     assert(BC.MIB->isInvoke(InvokeInst) && "must be invoke instruction");
     const Optional<MCPlus::MCLandingPad> LP = BC.MIB->getEHInfo(InvokeInst);

bolt/include/bolt/Passes/ValidateInternalCalls.h

public:
 
   void runOnFunctions(BinaryContext &BC) override;
 
 private:
   /// Fix the CFG to take into consideration internal calls that do not
   /// return, but are only used as a trick to perform Position Independent
   /// Code (PIC) computations. This will change internal calls to be treated
   /// as unconditional jumps.
-  bool fixCFGForPIC(BinaryFunction &Function) const;
+  void fixCFGForPIC(BinaryFunction &Function) const;
 
   /// Fix the CFG to take into consideration real internal calls (whole
   /// functions that got inlined inside its caller, but the CALL instruction
   /// wasn't removed).
   bool fixCFGForIC(BinaryFunction &Function) const;
 
   /// Detect tail calls in the range of the PIC access and fail to validate if
   /// one is detected. Tail calls are dangerous because they may be emitted

bolt/lib/Passes/ValidateInternalCalls.cpp

 //===- bolt/Passes/ValidateInternalCalls.cpp ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the ValidateInternalCalls class.
 //
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/ValidateInternalCalls.h"
+#include "bolt/Core/BinaryBasicBlock.h"
 #include "bolt/Passes/DataflowInfoManager.h"
 #include "bolt/Passes/FrameAnalysis.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/MC/MCInstPrinter.h"
+#include <functional>
+#include <memory>
+#include <utility>
+#include <vector>
 
 #define DEBUG_TYPE "bolt-internalcalls"
 
 namespace llvm {
 namespace bolt {
 
 namespace {
 
public:
 
   void run() {
     StackPointerTrackingBase<StackPointerTrackingForInternalCalls>::run();
   }
 };
 
 } // end anonymous namespace
 
-bool ValidateInternalCalls::fixCFGForPIC(BinaryFunction &Function) const {
+void ValidateInternalCalls::fixCFGForPIC(BinaryFunction &Function) const {
   const BinaryContext &BC = Function.getBinaryContext();
-  for (BinaryBasicBlock &BB : Function) {
-    for (auto II = BB.begin(); II != BB.end(); ++II) {
-      MCInst &Inst = *II;
-      BinaryBasicBlock *Target = getInternalCallTarget(Function, Inst);
-      if (!Target || BC.MIB->hasAnnotation(Inst, getProcessedICTag()))
+
+  // Make a copy of the blocks, because splitting and adding new blocks would invalidate iterators.
+  SmallVector<std::reference_wrapper<BinaryBasicBlock>, 0> Blocks(
+      Function.blocks());
+  for (BinaryBasicBlock &BB : Blocks) {
+    // When a block is split, current is updated to point to the newly created
+    // block.
+    BinaryBasicBlock *CurrentBB = &BB;
+    while (CurrentBB) {
+      // Search for the next internal call.
+      const BinaryBasicBlock::iterator InternalCall =
+          llvm::find_if(*CurrentBB, [&](const MCInst &Inst) {
+            const BinaryBasicBlock *const Target =
+                getInternalCallTarget(Function, Inst);
+            return Target && BC.MIB->hasAnnotation(Inst, getProcessedICTag());
+          });
+
+      // No internal call? Done with this block.
+      if (InternalCall == CurrentBB->end()) {
+        CurrentBB = nullptr;
         continue;
+      }
 
-      BC.MIB->addAnnotation(Inst, getProcessedICTag(), 0U);
-      InstructionListType MovedInsts = BB.splitInstructions(&Inst);
+      BinaryBasicBlock* NextBB = nullptr;
+
+      BinaryBasicBlock *Target = getInternalCallTarget(Function, *InternalCall);
+      BC.MIB->addAnnotation(*InternalCall, getProcessedICTag(), 0U);
+      InstructionListType MovedInsts =
+          CurrentBB->splitInstructions(&*InternalCall);
       if (!MovedInsts.empty()) {
-        // Split this block at the call instruction. Create an unreachable
-        // block.
+        // Split this block at the call instruction.
+        std::unique_ptr<BinaryBasicBlock> NewBB = Function.createBasicBlock();
+        NewBB->setOffset(0);
+        NewBB->addInstructions(MovedInsts.begin(), MovedInsts.end());
+        CurrentBB->moveAllSuccessorsTo(NewBB.get());
+
+        NextBB = NewBB.get();
         std::vector<std::unique_ptr<BinaryBasicBlock>> NewBBs;
-        NewBBs.emplace_back(Function.createBasicBlock());
-        NewBBs.back()->setOffset(0);
-        NewBBs.back()->addInstructions(MovedInsts.begin(), MovedInsts.end());
-        BB.moveAllSuccessorsTo(NewBBs.back().get());
-        Function.insertBasicBlocks(&BB, std::move(NewBBs));
+        NewBBs.emplace_back(std::move(NewBB));
+        Function.insertBasicBlocks(CurrentBB, std::move(NewBBs));
       }
       // Update successors
-      BB.removeAllSuccessors();
-      BB.addSuccessor(Target, BB.getExecutionCount(), 0ULL);
-      return true;
+      CurrentBB->removeAllSuccessors();
+      CurrentBB->addSuccessor(Target, BB.getExecutionCount(), 0ULL);
+
+      CurrentBB = NextBB;
     }
   }
-  return false;
 }
 
 bool ValidateInternalCalls::fixCFGForIC(BinaryFunction &Function) const {
   const BinaryContext &BC = Function.getBinaryContext();
   // Track SP value
   StackPointerTrackingForInternalCalls SPTIC(Function);
   SPTIC.run();
 
bool ValidateInternalCalls::hasTailCallsInRange(
   for (BinaryBasicBlock &BB : Function)
     for (MCInst &Inst : BB)
       if (BC.MIB->isTailCall(Inst))
         return true;
   return false;
 }
 
 bool ValidateInternalCalls::analyzeFunction(BinaryFunction &Function) const {
-  while (fixCFGForPIC(Function)) {
-  }
+  fixCFGForPIC(Function);
   clearAnnotations(Function);
   while (fixCFGForIC(Function)) {
   }
 
   BinaryContext &BC = Function.getBinaryContext();
   RegAnalysis RA = RegAnalysis(BC, nullptr, nullptr);
   RA.setConservativeStrategy(RegAnalysis::ConservativeStrategy::CLOBBERS_NONE);
   bool HasTailCalls = hasTailCallsInRange(Function);