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

Reviewed By: Amir

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

Author: maksfb

bolt/include/bolt/Core/BinaryBasicBlock.h

@@ -150,11 +150,9 @@ class BinaryBasicBlock {
   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.
@@ -561,6 +559,12 @@ class BinaryBasicBlock {
   /// 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; }
 
@@ -787,6 +791,9 @@ class BinaryBasicBlock {
   /// 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;

bolt/include/bolt/Core/BinaryFunction.h

@@ -701,6 +701,28 @@ class BinaryFunction {
     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();
@@ -1531,62 +1553,34 @@ class BinaryFunction {
     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);
+    BB->setLayoutIndex(layout_size());
+    BasicBlocksLayout.emplace_back(BB);
 
     return BB;
   }

bolt/lib/Core/BinaryBasicBlock.cpp

@@ -608,7 +608,7 @@ BinaryBasicBlock::getBranchInfo(const MCSymbol *Label) {
 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);

bolt/lib/Core/BinaryFunction.cpp

@@ -1957,8 +1957,10 @@ bool BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
     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);
     }
@@ -1997,8 +1999,9 @@ bool BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
           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);
       }
     }
@@ -2255,8 +2258,9 @@ void BinaryFunction::removeConditionalTailCalls() {
     // 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);
 
@@ -3833,8 +3837,8 @@ BinaryBasicBlock *BinaryFunction::splitEdge(BinaryBasicBlock *From,
   // 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

bolt/lib/Passes/BinaryPasses.cpp

@@ -1760,8 +1760,8 @@ void SpecializeMemcpy1::runOnFunctions(BinaryContext &BC) {
           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());

bolt/lib/Passes/IndirectCallPromotion.cpp

@@ -764,8 +764,8 @@ IndirectCallPromotion::rewriteCall(
     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");

bolt/lib/Passes/Inliner.cpp

@@ -289,7 +289,7 @@ Inliner::inlineCall(BinaryBasicBlock &CallerBB,
   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())

bolt/lib/Passes/Instrumentation.cpp

@@ -597,8 +597,7 @@ void Instrumentation::createAuxiliaryFunctions(BinaryContext &BC) {
     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),

bolt/lib/Passes/LongJmp.cpp

@@ -77,7 +77,7 @@ LongJmpPass::createNewStub(BinaryBasicBlock &SourceBB, const MCSymbol *TgtSym,
   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)

bolt/lib/Passes/PatchEntries.cpp

@@ -118,7 +118,7 @@ void PatchEntries::runOnFunctions(BinaryContext &BC) {
 
       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;

bolt/lib/Passes/RetpolineInsertion.cpp

@@ -91,8 +91,7 @@ BinaryFunction *createNewRetpoline(BinaryContext &BC,
   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);
   }
 

bolt/lib/Passes/TailDuplication.cpp

@@ -532,7 +532,7 @@ std::vector<BinaryBasicBlock *> TailDuplication::duplicateBlocks(
 
   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());

bolt/lib/Passes/ValidateInternalCalls.cpp

@@ -105,7 +105,7 @@ bool ValidateInternalCalls::fixCFGForPIC(BinaryFunction &Function) const {
         // 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));

bolt/unittests/Core/MCPlusBuilder.cpp

@@ -115,7 +115,7 @@ 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));