diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h --- a/bolt/include/bolt/Core/BinaryFunction.h +++ b/bolt/include/bolt/Core/BinaryFunction.h @@ -888,13 +888,11 @@ 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 diff --git a/bolt/include/bolt/Passes/ValidateInternalCalls.h b/bolt/include/bolt/Passes/ValidateInternalCalls.h --- a/bolt/include/bolt/Passes/ValidateInternalCalls.h +++ b/bolt/include/bolt/Passes/ValidateInternalCalls.h @@ -61,7 +61,7 @@ /// 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 diff --git a/bolt/lib/Passes/ValidateInternalCalls.cpp b/bolt/lib/Passes/ValidateInternalCalls.cpp --- a/bolt/lib/Passes/ValidateInternalCalls.cpp +++ b/bolt/lib/Passes/ValidateInternalCalls.cpp @@ -11,9 +11,12 @@ //===----------------------------------------------------------------------===// #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 #define DEBUG_TYPE "bolt-internalcalls" @@ -90,34 +93,43 @@ } // end anonymous namespace -bool 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())) - continue; +void ValidateInternalCalls::fixCFGForPIC(BinaryFunction &Function) const { + std::queue Work; + for (BinaryBasicBlock &BB : Function) + Work.emplace(&BB); - BC.MIB->addAnnotation(Inst, getProcessedICTag(), 0U); - InstructionListType MovedInsts = BB.splitInstructions(&Inst); - if (!MovedInsts.empty()) { - // Split this block at the call instruction. Create an unreachable - // block. - std::vector> 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)); - } - // Update successors - BB.removeAllSuccessors(); - BB.addSuccessor(Target, BB.getExecutionCount(), 0ULL); - return true; + while (!Work.empty()) { + BinaryBasicBlock &BB = *Work.front(); + Work.pop(); + + // Search for the next internal call. + const BinaryBasicBlock::iterator InternalCall = + llvm::find_if(BB, [&](const MCInst &Inst) { + return getInternalCallTarget(Function, Inst) != nullptr; + }); + + // No internal call? Done with this block. + if (InternalCall == BB.end()) + continue; + + BinaryBasicBlock *Target = getInternalCallTarget(Function, *InternalCall); + InstructionListType MovedInsts = BB.splitInstructions(&*InternalCall); + if (!MovedInsts.empty()) { + // Split this block at the call instruction. + std::unique_ptr NewBB = Function.createBasicBlock(); + NewBB->setOffset(0); + NewBB->addInstructions(MovedInsts.begin(), MovedInsts.end()); + BB.moveAllSuccessorsTo(NewBB.get()); + + Work.emplace(NewBB.get()); + std::vector> NewBBs; + NewBBs.emplace_back(std::move(NewBB)); + Function.insertBasicBlocks(&BB, std::move(NewBBs)); } + // Update successors + BB.removeAllSuccessors(); + BB.addSuccessor(Target, BB.getExecutionCount(), 0ULL); } - return false; } bool ValidateInternalCalls::fixCFGForIC(BinaryFunction &Function) const { @@ -194,9 +206,7 @@ } bool ValidateInternalCalls::analyzeFunction(BinaryFunction &Function) const { - while (fixCFGForPIC(Function)) { - } - clearAnnotations(Function); + fixCFGForPIC(Function); while (fixCFGForIC(Function)) { }