[BOLT][NFCI] Refactor interface for adding basic blocks

[Amir]

LGTM

bolt/include/bolt/Core/BinaryBasicBlock.h

 
 private:
   BinaryBasicBlock() = delete;
   BinaryBasicBlock(const BinaryBasicBlock &) = delete;
   BinaryBasicBlock(const BinaryBasicBlock &&) = delete;
   BinaryBasicBlock &operator=(const BinaryBasicBlock &) = delete;
   BinaryBasicBlock &operator=(const BinaryBasicBlock &&) = delete;
 
-  explicit BinaryBasicBlock(BinaryFunction *Function, MCSymbol *Label,
-                            uint32_t Offset = INVALID_OFFSET)
+  explicit BinaryBasicBlock(BinaryFunction *Function, MCSymbol *Label)
       : Function(Function), Label(Label) {
     assert(Function && "Function must be non-null");
-    InputRange.first = Offset;
   }
 
   // Exclusively managed by BinaryFunction.
   friend class BinaryFunction;
   friend bool operator<(const BinaryBasicBlock &LHS,
                         const BinaryBasicBlock &RHS);
 
   /// Assign new label to the basic block.
public:
   /// Calculate and return CFI state after execution of this basic block.
   /// The state depends on CFI state at entry and CFI instructions inside the
   /// basic block.
   int32_t getCFIStateAtExit() const { return getCFIStateAtInstr(nullptr); }
 
   /// Set minimum alignment for the basic block.
   void setAlignment(uint32_t Align) { Alignment = Align; }
 
+  /// Set alignment of the block based on the alignment of its offset.
+  void setDerivedAlignment() {
+    const uint64_t DerivedAlignment = getOffset() & (1 + ~getOffset());
+    Alignment = std::min(DerivedAlignment, uint64_t(32));
+  }
+
   /// Return required alignment for the block.
   uint32_t getAlignment() const { return Alignment; }
 
   /// Set the maximum number of bytes to use for the block alignment.
   void setAlignmentMaxBytes(uint32_t Value) { AlignmentMaxBytes = Value; }
 
   /// Return the maximum number of bytes to use for the block alignment.
   uint32_t getAlignmentMaxBytes() const { return AlignmentMaxBytes; }
public:
 
   /// Split basic block at the instruction pointed to by II.
   /// All iterators pointing after II get invalidated.
   ///
   /// Return the new basic block that starts with the instruction
   /// at the split point.
   BinaryBasicBlock *splitAt(iterator II);
 
+  /// Set start offset of this basic block in the input binary.
+  void setOffset(uint32_t Offset) { InputRange.first = Offset; };
+
   /// Sets address of the basic block in the output.
   void setOutputStartAddress(uint64_t Address) {
     OutputAddressRange.first = Address;
   }
 
   /// Sets address past the end of the basic block in the output.
   void setOutputEndAddress(uint64_t Address) {
     OutputAddressRange.second = Address;

bolt/include/bolt/Core/BinaryFunction.h

BinaryFunction(const std::string &Name, BinaryContext &BC, bool IsSimple)
       : Address(0), Size(0), BC(BC), IsSimple(IsSimple),
         CodeSectionName(buildCodeSectionName(Name, BC)),
         ColdCodeSectionName(buildColdCodeSectionName(Name, BC)),
         FunctionNumber(++Count) {
     Symbols.push_back(BC.Ctx->getOrCreateSymbol(Name));
     IsInjected = true;
   }
 
+  /// Create a basic block at a given \p Offset in the function and append it
+  /// to the end of list of blocks. Used during CFG construction only.
+  BinaryBasicBlock *addBasicBlockAt(uint64_t Offset, MCSymbol *Label) {
+    assert(CurrentState == State::Disassembled &&
+           "Cannot add block with an offset in non-disassembled state.");
+    assert(!getBasicBlockAtOffset(Offset) &&
+           "Basic block already exists at the offset.");
+
+    BasicBlocks.emplace_back(createBasicBlock(Label).release());
+    BinaryBasicBlock *BB = BasicBlocks.back();
+
+    BB->setIndex(BasicBlocks.size() - 1);
+    BB->setOffset(Offset);
+
+    BasicBlockOffsets.emplace_back(Offset, BB);
+    assert(std::is_sorted(BasicBlockOffsets.begin(), BasicBlockOffsets.end(),
+                          CompareBasicBlockOffsets()) &&
+           std::is_sorted(begin(), end()));
+
+    return BB;
+  }
+
   /// Clear state of the function that could not be disassembled or if its
   /// disassembled state was later invalidated.
   void clearDisasmState();
 
   /// Release memory allocated for CFG and instructions.
   /// We still keep basic blocks for address translation/mapping purposes.
   void releaseCFG() {
     for (BinaryBasicBlock *BB : BasicBlocks)
public:
 
   /// Return true if the given address \p PC is inside the function body.
   bool containsAddress(uint64_t PC, bool UseMaxSize = false) const {
     if (UseMaxSize)
       return Address <= PC && PC < Address + MaxSize;
     return Address <= PC && PC < Address + Size;
   }
 
-  /// Create a basic block at a given \p Offset in the
-  /// function.
-  /// If \p DeriveAlignment is true, set the alignment of the block based
-  /// on the alignment of the existing offset.
-  /// The new block is not inserted into the CFG.  The client must
-  /// use insertBasicBlocks to add any new blocks to the CFG.
+  /// Create a basic block in the function. The new block is *NOT* inserted
+  /// into the CFG. The caller must use insertBasicBlocks() to add any new
+  /// blocks to the CFG.
   std::unique_ptr<BinaryBasicBlock>
-  createBasicBlock(uint64_t Offset, MCSymbol *Label = nullptr,
-                   bool DeriveAlignment = false) {
-    assert(BC.Ctx && "cannot be called with empty context");
+  createBasicBlock(MCSymbol *Label = nullptr) {
     if (!Label) {
       std::unique_lock<std::shared_timed_mutex> Lock(BC.CtxMutex);
       Label = BC.Ctx->createNamedTempSymbol("BB");
     }
-    auto BB = std::unique_ptr<BinaryBasicBlock>(
-        new BinaryBasicBlock(this, Label, Offset));
-
-    if (DeriveAlignment) {
-      uint64_t DerivedAlignment = Offset & (1 + ~Offset);
-      BB->setAlignment(std::min(DerivedAlignment, uint64_t(32)));
-    }
+    auto BB =
+        std::unique_ptr<BinaryBasicBlock>(new BinaryBasicBlock(this, Label));
 
     LabelToBB[Label] = BB.get();
 
     return BB;
   }
 
-  /// Create a basic block at a given \p Offset in the
-  /// function and append it to the end of list of blocks.
-  /// If \p DeriveAlignment is true, set the alignment of the block based
-  /// on the alignment of the existing offset.
-  ///
-  /// Returns NULL if basic block already exists at the \p Offset.
-  BinaryBasicBlock *addBasicBlock(uint64_t Offset, MCSymbol *Label = nullptr,
-                                  bool DeriveAlignment = false) {
-    assert((CurrentState == State::CFG || !getBasicBlockAtOffset(Offset)) &&
-           "basic block already exists in pre-CFG state");
-
-    std::unique_ptr<BinaryBasicBlock> BBPtr =
-        createBasicBlock(Offset, Label, DeriveAlignment);
-    BasicBlocks.emplace_back(BBPtr.release());
+  /// Create a new basic block with an optional \p Label and add it to the list
+  /// of basic blocks of this function.
+  BinaryBasicBlock *addBasicBlock(MCSymbol *Label = nullptr) {
+    assert(CurrentState == State::CFG && "Can only add blocks in CFG state");
 
+    BasicBlocks.emplace_back(createBasicBlock(Label).release());
     BinaryBasicBlock *BB = BasicBlocks.back();
-    BB->setIndex(BasicBlocks.size() - 1);
 
-    if (CurrentState == State::Disassembled) {
-      BasicBlockOffsets.emplace_back(Offset, BB);
-    } else if (CurrentState == State::CFG) {
-      BB->setLayoutIndex(layout_size());
-      BasicBlocksLayout.emplace_back(BB);
-    }
-
-    assert(CurrentState == State::CFG ||
-           (std::is_sorted(BasicBlockOffsets.begin(), BasicBlockOffsets.end(),
-                           CompareBasicBlockOffsets()) &&
-            std::is_sorted(begin(), end())));
+    BB->setIndex(BasicBlocks.size() - 1);
 
     return BB;
   }
 
   /// Add basic block \BB as an entry point to the function. Return global
   /// symbol associated with the entry.
   MCSymbol *addEntryPoint(const BinaryBasicBlock &BB);
 

bolt/lib/Core/BinaryBasicBlock.cpp

BinaryBasicBlock::getBranchInfo(const MCSymbol *Label) {
 
   llvm_unreachable("Invalid successor");
   return *BI;
 }
 
 BinaryBasicBlock *BinaryBasicBlock::splitAt(iterator II) {
   assert(II != end() && "expected iterator pointing to instruction");
 
-  BinaryBasicBlock *NewBlock = getFunction()->addBasicBlock(0);
+  BinaryBasicBlock *NewBlock = getFunction()->addBasicBlock();
 
   // Adjust successors/predecessors and propagate the execution count.
   moveAllSuccessorsTo(NewBlock);
   addSuccessor(NewBlock, getExecutionCount(), 0);
 
   // Set correct CFI state for the new block.
   NewBlock->setCFIState(getCFIStateAtInstr(&*II));
 

bolt/lib/Core/BinaryFunction.cpp

bool BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
   for (auto I = Instructions.begin(), E = Instructions.end(); I != E; ++I) {
     const uint32_t Offset = I->first;
     MCInst &Instr = I->second;
 
     auto LI = Labels.find(Offset);
     if (LI != Labels.end()) {
       // Always create new BB at branch destination.
       PrevBB = InsertBB ? InsertBB : PrevBB;
-      InsertBB = addBasicBlock(LI->first, LI->second,
-                               opts::PreserveBlocksAlignment && IsLastInstrNop);
+      InsertBB = addBasicBlockAt(LI->first, LI->second);
+      if (opts::PreserveBlocksAlignment && IsLastInstrNop)
+        InsertBB->setDerivedAlignment();
+
       if (PrevBB)
         updateOffset(LastInstrOffset);
     }
 
     const uint64_t InstrInputAddr = I->first + Address;
     bool IsSDTMarker =
         MIB->isNoop(Instr) && BC.SDTMarkers.count(InstrInputAddr);
     bool IsLKMarker = BC.LKMarkers.count(InstrInputAddr);
if (!InsertBB) {
         // Temporarily restore inserter basic block.
         InsertBB = PrevBB;
       } else {
         MCSymbol *Label;
         {
           auto L = BC.scopeLock();
           Label = BC.Ctx->createNamedTempSymbol("FT");
         }
-        InsertBB = addBasicBlock(
-            Offset, Label, opts::PreserveBlocksAlignment && IsLastInstrNop);
+        InsertBB = addBasicBlockAt(Offset, Label);
+        if (opts::PreserveBlocksAlignment && IsLastInstrNop)
+          InsertBB->setDerivedAlignment();
         updateOffset(LastInstrOffset);
       }
     }
     if (Offset == getFirstInstructionOffset()) {
       // Add associated CFI pseudos in the first offset
       addCFIPlaceholders(Offset, InsertBB);
     }
 
for (auto BBI = begin(); BBI != end(); ++BBI) {
     // the same destination.
     const MCSymbol *CTCTargetLabel = BC.MIB->getTargetSymbol(*CTCInstr);
     assert(CTCTargetLabel && "symbol expected for conditional tail call");
     MCInst TailCallInstr;
     BC.MIB->createTailCall(TailCallInstr, CTCTargetLabel, BC.Ctx.get());
     // Link new BBs to the original input offset of the BB where the CTC
     // is, so we can map samples recorded in new BBs back to the original BB
     // seem in the input binary (if using BAT)
-    std::unique_ptr<BinaryBasicBlock> TailCallBB = createBasicBlock(
-        BB.getInputOffset(), BC.Ctx->createNamedTempSymbol("TC"));
+    std::unique_ptr<BinaryBasicBlock> TailCallBB =
+        createBasicBlock(BC.Ctx->createNamedTempSymbol("TC"));
+    TailCallBB->setOffset(BB.getInputOffset());
     TailCallBB->addInstruction(TailCallInstr);
     TailCallBB->setCFIState(CFIStateBeforeCTC);
 
     // Add CFG edge with profile info from BB to TailCallBB.
     BB.addSuccessor(TailCallBB.get(), CTCTakenCount, CTCMispredCount);
 
     // Add execution count for the block.
     TailCallBB->setExecutionCount(CTCTakenCount);
BinaryBasicBlock *BinaryFunction::splitEdge(BinaryBasicBlock *From,
   MCSymbol *Tmp;
   {
     auto L = BC.scopeLock();
     Tmp = BC.Ctx->createNamedTempSymbol("SplitEdge");
   }
   // Link new BBs to the original input offset of the From BB, so we can map
   // samples recorded in new BBs back to the original BB seem in the input
   // binary (if using BAT)
-  std::unique_ptr<BinaryBasicBlock> NewBB =
-      createBasicBlock(From->getInputOffset(), Tmp);
+  std::unique_ptr<BinaryBasicBlock> NewBB = createBasicBlock(Tmp);
+  NewBB->setOffset(From->getInputOffset());
   BinaryBasicBlock *NewBBPtr = NewBB.get();
 
   // Update "From" BB
   auto I = From->succ_begin();
   auto BI = From->branch_info_begin();
   for (; I != From->succ_end(); ++I) {
     if (*I == To)
       break;

bolt/lib/Passes/BinaryPasses.cpp

for (BinaryBasicBlock *CurBB : Blocks) {
           NextBB = OneByteMemcpyBB->splitAt(OneByteMemcpyBB->begin());
           NextBB->eraseInstruction(NextBB->begin());
         } else {
           NextBB = OneByteMemcpyBB->getSuccessor();
           OneByteMemcpyBB->eraseInstruction(OneByteMemcpyBB->begin());
           assert(NextBB && "unexpected call to memcpy() with no return");
         }
 
-        BinaryBasicBlock *MemcpyBB =
-            Function.addBasicBlock(CurBB->getInputOffset());
+        BinaryBasicBlock *MemcpyBB = Function.addBasicBlock();
+        MemcpyBB->setOffset(CurBB->getInputOffset());
         InstructionListType CmpJCC =
             BC.MIB->createCmpJE(BC.MIB->getIntArgRegister(2), 1,
                                 OneByteMemcpyBB->getLabel(), BC.Ctx.get());
         CurBB->addInstructions(CmpJCC);
         CurBB->addSuccessor(MemcpyBB);
 
         MemcpyBB->addInstruction(std::move(MemcpyInstr));
         MemcpyBB->addSuccessor(NextBB);

bolt/lib/Passes/IndirectCallPromotion.cpp

else
     IndCallBlock.replaceInstruction(std::prev(IndCallBlock.end()),
                                     ICPcode.front().second);
   IndCallBlock.addInstructions(TailInsts.begin(), TailInsts.end());
 
   for (auto Itr = ICPcode.begin() + 1; Itr != ICPcode.end(); ++Itr) {
     MCSymbol *&Sym = Itr->first;
     InstructionListType &Insts = Itr->second;
     assert(Sym);
-    std::unique_ptr<BinaryBasicBlock> TBB =
-        Function.createBasicBlock(OrigOffset, Sym);
+    std::unique_ptr<BinaryBasicBlock> TBB = Function.createBasicBlock(Sym);
+    TBB->setOffset(OrigOffset);
     for (MCInst &Inst : Insts) // sanitize new instructions.
       if (MIB->isCall(Inst))
         MIB->removeAnnotation(Inst, "CallProfile");
     TBB->addInstructions(Insts.begin(), Insts.end());
     NewBBs.emplace_back(std::move(TBB));
   }
 
   // Move tail of instructions from after the original call to

bolt/lib/Passes/Inliner.cpp

Inliner::inlineCall(BinaryBasicBlock &CallerBB,
   // Save execution count of the first block as we don't want it to change
   // later due to profile adjustment rounding errors.
   const uint64_t FirstInlinedBBCount = FirstInlinedBB->getKnownExecutionCount();
 
   // Copy basic blocks and maintain a map from their origin.
   std::unordered_map<const BinaryBasicBlock *, BinaryBasicBlock *> InlinedBBMap;
   InlinedBBMap[&Callee.front()] = FirstInlinedBB;
   for (auto BBI = std::next(Callee.begin()); BBI != Callee.end(); ++BBI) {
-    BinaryBasicBlock *InlinedBB = CallerFunction.addBasicBlock(0);
+    BinaryBasicBlock *InlinedBB = CallerFunction.addBasicBlock();
     InlinedBBMap[&*BBI] = InlinedBB;
     InlinedBB->setCFIState(FirstInlinedBB->getCFIState());
     if (Callee.hasValidProfile())
       InlinedBB->setExecutionCount(BBI->getKnownExecutionCount());
     else
       InlinedBB->setExecutionCount(FirstInlinedBBCount);
   }
 

bolt/lib/Passes/Instrumentation.cpp

 }
 
 void Instrumentation::createAuxiliaryFunctions(BinaryContext &BC) {
   auto createSimpleFunction =
       [&](StringRef Title, InstructionListType Instrs) -> BinaryFunction * {
     BinaryFunction *Func = BC.createInjectedBinaryFunction(std::string(Title));
 
     std::vector<std::unique_ptr<BinaryBasicBlock>> BBs;
-    BBs.emplace_back(
-        Func->createBasicBlock(BinaryBasicBlock::INVALID_OFFSET, nullptr));
+    BBs.emplace_back(Func->createBasicBlock());
     BBs.back()->addInstructions(Instrs.begin(), Instrs.end());
     BBs.back()->setCFIState(0);
     Func->insertBasicBlocks(nullptr, std::move(BBs),
                             /*UpdateLayout=*/true,
                             /*UpdateCFIState=*/false);
     Func->updateState(BinaryFunction::State::CFG_Finalized);
     return Func;
   };

bolt/lib/Passes/LongJmp.cpp

 
 std::pair<std::unique_ptr<BinaryBasicBlock>, MCSymbol *>
 LongJmpPass::createNewStub(BinaryBasicBlock &SourceBB, const MCSymbol *TgtSym,
                            bool TgtIsFunc, uint64_t AtAddress) {
   BinaryFunction &Func = *SourceBB.getFunction();
   const BinaryContext &BC = Func.getBinaryContext();
   const bool IsCold = SourceBB.isCold();
   MCSymbol *StubSym = BC.Ctx->createNamedTempSymbol("Stub");
-  std::unique_ptr<BinaryBasicBlock> StubBB = Func.createBasicBlock(0, StubSym);
+  std::unique_ptr<BinaryBasicBlock> StubBB = Func.createBasicBlock(StubSym);
   MCInst Inst;
   BC.MIB->createUncondBranch(Inst, TgtSym, BC.Ctx.get());
   if (TgtIsFunc)
     BC.MIB->convertJmpToTailCall(Inst);
   StubBB->addInstruction(Inst);
   StubBB->setExecutionCount(0);
 
   // Register this in stubs maps

bolt/lib/Passes/PatchEntries.cpp

for (Patch &Patch : PendingPatches) {
           NameResolver::append(Patch.Symbol->getName(), ".org.0"));
       // Force the function to be emitted at the given address.
       PatchFunction->setOutputAddress(Patch.Address);
       PatchFunction->setFileOffset(Patch.FileOffset);
       PatchFunction->setOriginSection(Patch.Section);
 
       InstructionListType Seq;
       BC.MIB->createLongTailCall(Seq, Patch.Symbol, BC.Ctx.get());
-      PatchFunction->addBasicBlock(0)->addInstructions(Seq);
+      PatchFunction->addBasicBlock()->addInstructions(Seq);
 
       // Verify the size requirements.
       uint64_t HotSize, ColdSize;
       std::tie(HotSize, ColdSize) = BC.calculateEmittedSize(*PatchFunction);
       assert(!ColdSize && "unexpected cold code");
       assert(HotSize <= PatchSize && "max patch size exceeded");
     }
 
     Function.setIsPatched(true);
   }
 }
 
 } // end namespace bolt
 } // end namespace llvm

bolt/lib/Passes/RetpolineInsertion.cpp

LLVM_DEBUG(dbgs() << "BOLT-DEBUG: Creating a new retpoline function["
                     << RetpolineTag << "]\n");
 
   BinaryFunction *NewRetpoline =
       BC.createInjectedBinaryFunction(RetpolineTag, true);
   std::vector<std::unique_ptr<BinaryBasicBlock>> NewBlocks(3);
   for (int I = 0; I < 3; I++) {
     MCSymbol *Symbol =
         Ctx.createNamedTempSymbol(Twine(RetpolineTag + "_BB" + to_string(I)));
-    NewBlocks[I] = NewRetpoline->createBasicBlock(
-        BinaryBasicBlock::INVALID_OFFSET, Symbol);
+    NewBlocks[I] = NewRetpoline->createBasicBlock(Symbol);
     NewBlocks[I].get()->setCFIState(0);
   }
 
   BinaryBasicBlock &BB0 = *NewBlocks[0].get();
   BinaryBasicBlock &BB1 = *NewBlocks[1].get();
   BinaryBasicBlock &BB2 = *NewBlocks[2].get();
 
   BB0.addSuccessor(&BB2, 0, 0);

bolt/lib/Passes/TailDuplication.cpp

assert(LastDuplicatedBB->succ_size() == 1 &&
          "tail duplication cannot act on a block with more than 1 successor");
   LastDuplicatedBB->removeSuccessor(LastDuplicatedBB->getSuccessor());
 
   std::vector<std::unique_ptr<BinaryBasicBlock>> DuplicatedBlocks;
   std::vector<BinaryBasicBlock *> DuplicatedBlocksToReturn;
 
   for (BinaryBasicBlock *CurBB : BlocksToDuplicate) {
     DuplicatedBlocks.emplace_back(
-        BF->createBasicBlock(0, (BC.Ctx)->createNamedTempSymbol("tail-dup")));
+        BF->createBasicBlock((BC.Ctx)->createNamedTempSymbol("tail-dup")));
     BinaryBasicBlock *NewBB = DuplicatedBlocks.back().get();
 
     NewBB->addInstructions(CurBB->begin(), CurBB->end());
     // Set execution count as if it was just a copy of the original
     NewBB->setExecutionCount(CurBB->getExecutionCount());
     NewBB->setIsCold(CurBB->isCold());
     LastDuplicatedBB->addSuccessor(NewBB, LastBI);
 

bolt/lib/Passes/ValidateInternalCalls.cpp

for (auto II = BB.begin(); II != BB.end(); ++II) {
         continue;
 
       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<std::unique_ptr<BinaryBasicBlock>> NewBBs;
-        NewBBs.emplace_back(Function.createBasicBlock(0));
+        NewBBs.emplace_back(Function.createBasicBlock());
         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;

bolt/unittests/Core/MCPlusBuilder.cpp

TEST_P(MCPlusBuilderTester, AliasSmallerAX) {
   size_t AliasesAXCount = sizeof(AliasesAX) / sizeof(*AliasesAX);
   testRegAliases(Triple::x86_64, X86::AX, AliasesAX, AliasesAXCount, true);
 }
 
 TEST_P(MCPlusBuilderTester, ReplaceRegWithImm) {
   if (GetParam() != Triple::x86_64)
     GTEST_SKIP();
   BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
-  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock(0);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
   MCInst Inst; // cmpl    %eax, %ebx
   Inst.setOpcode(X86::CMP32rr);
   Inst.addOperand(MCOperand::createReg(X86::EAX));
   Inst.addOperand(MCOperand::createReg(X86::EBX));
   auto II = BB->addInstruction(Inst);
   bool Replaced = BC->MIB->replaceRegWithImm(*II, X86::EBX, 1);
   ASSERT_TRUE(Replaced);
   ASSERT_EQ(II->getOpcode(), X86::CMP32ri8);