[BOLT] Allocate FunctionFragment on heap

[FPar]

Summary

This changes FunctionFragment from being used as a temporary proxy
object to access basic block ranges to a heap-allocated object that can
store fragment-specific information.

[rafauler]

LGTM

[FPar]

Fix iterator conversion build failure

bolt/include/bolt/Core/BinaryEmitter.h

 /// contained in \p BC. This is done to distinguish them from sections emitted
 /// by LLVM backend.
 void emitBinaryContext(MCStreamer &Streamer, BinaryContext &BC,
                        StringRef OrgSecPrefix = "");
 
 /// Emit \p BF function code. The caller is responsible for emitting function
 /// symbol(s) and setting the section to emit the code to.
 void emitFunctionBody(MCStreamer &Streamer, BinaryFunction &BF,
-                      const FunctionFragment &FF, bool EmitCodeOnly);
+                      FunctionFragment &FF, bool EmitCodeOnly);
 
 } // namespace bolt
 } // namespace llvm
 
 #endif

bolt/include/bolt/Core/FunctionLayout.h

 #ifndef BOLT_CORE_FUNCTION_LAYOUT_H
 #define BOLT_CORE_FUNCTION_LAYOUT_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/ADT/iterator_range.h"
+#include <iterator>
+#include <utility>
 
 namespace llvm {
 namespace bolt {
 
 class BinaryFunction;
 class BinaryBasicBlock;
 class FunctionLayout;
 

 };
 
 /// A freestanding subset of contiguous blocks of a function.
 class FunctionFragment {
   using BasicBlockListType = SmallVector<BinaryBasicBlock *, 0>;
   using FragmentListType = SmallVector<unsigned, 0>;
 
 public:
-  using const_iterator = BasicBlockListType::const_iterator;
+  using iterator = raw_pointer_iterator<BasicBlockListType::const_iterator,
+                                        BinaryBasicBlock>;
+  using const_iterator =
+      raw_pointer_iterator<BasicBlockListType::const_iterator,
+                           const BinaryBasicBlock>;
 
 private:
+  FunctionLayout *Layout;
   FragmentNum Num;
-  const FunctionLayout &Layout;
+  unsigned StartIndex;
+  unsigned Size = 0;
 
-  FunctionFragment(FragmentNum Num, const FunctionLayout &Layout)
-      : Num(Num), Layout(Layout) {}
+  FunctionFragment(FunctionLayout &Layout, FragmentNum Num);
+  FunctionFragment(const FunctionFragment &) = default;
+  FunctionFragment(FunctionFragment &&) = default;
+  FunctionFragment &operator=(const FunctionFragment &) = default;
+  FunctionFragment &operator=(FunctionFragment &&) = default;
+  ~FunctionFragment() = default;
 
 public:
   FragmentNum getFragmentNum() const { return Num; }
-  bool isMainFragment() const { return Num.get() == 0; }
-  bool isSplitFragment() const { return Num.get() > 0; }
+  bool isMainFragment() const {
+    return getFragmentNum() == FragmentNum::main();
+  }
+  bool isSplitFragment() const { return !isMainFragment(); }
 
-  unsigned size() const;
-  bool empty() const;
+  unsigned size() const { return Size; };
+  bool empty() const { return size() == 0; };
+  iterator begin();
   const_iterator begin() const;
+  iterator end();
   const_iterator end() const;
-  BinaryBasicBlock *front() const;
+  const BinaryBasicBlock *front() const;
 
   friend class FunctionLayout;
-  friend class FragmentIterator;
 };
 
 /// The function layout represents the fragments we split a function into and
 /// the order of basic blocks within each fragment.
 ///
 /// Internally, the function layout stores blocks across fragments contiguously.
 /// This is necessary to retain compatibility with existing code and tests that
 /// iterate  over all blocks of the layout and depend on that order. When
 /// writing new code, avoid iterating using FunctionLayout::blocks() by
 /// iterating either over fragments or over BinaryFunction::begin()..end().
 class FunctionLayout {
 private:
+  using FragmentListType = SmallVector<FunctionFragment *, 0>;
   using BasicBlockListType = SmallVector<BinaryBasicBlock *, 0>;
   using block_iterator = BasicBlockListType::iterator;
-  using FragmentListType = SmallVector<unsigned, 0>;
-
-public:
-  class FragmentIterator;
-
-  class FragmentIterator
-      : public iterator_facade_base<
-            FragmentIterator, std::bidirectional_iterator_tag, FunctionFragment,
-            std::ptrdiff_t, FunctionFragment *, FunctionFragment> {
-    FragmentNum Num;
-    const FunctionLayout *Layout;
-
-    FragmentIterator(FragmentNum Num, const FunctionLayout *Layout)
-        : Num(Num), Layout(Layout) {
-      assert(Num.get() <= Layout->fragment_size() &&
-             "Initializing iterator out of bounds");
-    }
 
-  public:
-    bool operator==(const FragmentIterator &Other) const {
-      return Num == Other.Num;
-    }
-
-    FunctionFragment operator*() const {
-      assert(Num.get() < Layout->fragment_size() &&
-             "Dereferencing end() iterator (or past it)");
-      return FunctionFragment(Num, *Layout);
-    }
-
-    FragmentIterator &operator++() {
-      assert(Num.get() < Layout->fragment_size() &&
-             "Incrementing iterator past end()");
-      Num = FragmentNum(Num.get() + 1);
-      return *this;
-    }
-
-    FragmentIterator &operator--() {
-      assert(Num.get() > 0 && "Decrementing iterator past begin()");
-      Num = FragmentNum(Num.get() - 1);
-      return *this;
-    }
-
-    friend class FunctionLayout;
-  };
-
-  using const_iterator = FragmentIterator;
-  using block_const_iterator = BasicBlockListType::const_iterator;
+  using fragment_iterator = pointee_iterator<FragmentListType::const_iterator>;
+  using fragment_const_iterator =
+      pointee_iterator<FragmentListType::const_iterator,
+                       const FunctionFragment>;
+  using block_const_iterator =
+      raw_pointer_iterator<BasicBlockListType::const_iterator,
+                           const BinaryBasicBlock>;
+  using block_reverse_iterator = std::reverse_iterator<block_iterator>;
   using block_const_reverse_iterator =
-      BasicBlockListType::const_reverse_iterator;
+      std::reverse_iterator<block_const_iterator>;
 
 private:
+  FragmentListType Fragments;
   BasicBlockListType Blocks;
-  /// List of indices dividing block list into fragments. To simplify iteration,
-  /// we have `Fragments.back()` equals `Blocks.size()`. Hence,
-  /// `Fragments.size()` equals `this->size() + 1`. Always contains at least one
-  /// fragment.
-  FragmentListType Fragments = {0, 0};
 
 public:
+  FunctionLayout();
+  FunctionLayout(const FunctionLayout &Other);
+  FunctionLayout(FunctionLayout &&Other);
+  FunctionLayout &operator=(const FunctionLayout &Other);
+  FunctionLayout &operator=(FunctionLayout &&Other);
+  ~FunctionLayout();
+
   /// Add an empty fragment.
-  FunctionFragment addFragment();
+  FunctionFragment &addFragment();
+
+  /// Return the fragment identified by Num.
+  FunctionFragment &getFragment(FragmentNum Num);
 
   /// Return the fragment identified by Num.
-  FunctionFragment getFragment(FragmentNum Num) const;
+  const FunctionFragment &getFragment(FragmentNum Num) const;
 
   /// Get the fragment that contains all entry blocks and other blocks that
   /// cannot be split.
-  FunctionFragment getMainFragment() const {
+  FunctionFragment &getMainFragment() {
     return getFragment(FragmentNum::main());
   }
 
   /// Get the fragment that contains all entry blocks and other blocks that
   /// cannot be split.
-  iterator_range<const_iterator> getSplitFragments() const {
+  const FunctionFragment &getMainFragment() const {
+    return getFragment(FragmentNum::main());
+  }
+
+  /// Get the fragment that contains all entry blocks and other blocks that
+  /// cannot be split.
+  iterator_range<fragment_iterator> getSplitFragments() {
+    return {++fragment_begin(), fragment_end()};
+  }
+
+  /// Get the fragment that contains all entry blocks and other blocks that
+  /// cannot be split.
+  iterator_range<fragment_const_iterator> getSplitFragments() const {
     return {++fragment_begin(), fragment_end()};
   }
 
   /// Find the fragment that contains BB.
-  FunctionFragment findFragment(const BinaryBasicBlock *BB) const;
+  const FunctionFragment &findFragment(const BinaryBasicBlock *BB) const;
 
   /// Add BB to the end of the last fragment.
   void addBasicBlock(BinaryBasicBlock *BB);
 
   /// Insert range of basic blocks after InsertAfter. If InsertAfter is nullptr,
   /// the blocks will be inserted at the start of the function.
-  void insertBasicBlocks(BinaryBasicBlock *InsertAfter,
+  void insertBasicBlocks(const BinaryBasicBlock *InsertAfter,
                          ArrayRef<BinaryBasicBlock *> NewBlocks);
 
   /// Erase all blocks from the layout that are in ToErase. If this method
   /// erases all blocks of a fragment, it will be removed as well.
   void eraseBasicBlocks(const DenseSet<const BinaryBasicBlock *> ToErase);
 
   /// Make sure fragments' and basic blocks' indices match the current layout.
-  void updateLayoutIndices() const;
+  void updateLayoutIndices();
 
   /// Replace the current layout with NewLayout. Uses the block's
   /// self-identifying fragment number to assign blocks to infer function
   /// fragments. Returns `true` if the new layout is different from the current
   /// layout.
   bool update(ArrayRef<BinaryBasicBlock *> NewLayout);
 
   /// Clear layout releasing memory.
   void clear();
 
-  BinaryBasicBlock *getBlock(unsigned Index) const { return Blocks[Index]; }
+  BinaryBasicBlock *getBlock(unsigned Index) { return Blocks[Index]; }
+
+  const BinaryBasicBlock *getBlock(unsigned Index) const {
+    return Blocks[Index];
+  }
 
   /// Returns the basic block after the given basic block in the layout or
   /// nullptr if the last basic block is given.
-  BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *BB,
+  BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *const BB,
+                                       const bool IgnoreSplits = true) {
+    return const_cast<BinaryBasicBlock *>(
+        static_cast<const FunctionLayout &>(*this).getBasicBlockAfter(
+            BB, IgnoreSplits));
+  }
+
+  /// Returns the basic block after the given basic block in the layout or
+  /// nullptr if the last basic block is given.
+  const BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *BB,
                                              bool IgnoreSplits = true) const;
 
   /// True if the layout contains at least two non-empty fragments.
   bool isSplit() const;
 
   /// Get the edit distance of the new layout with respect to the previous
   /// layout after basic block reordering.
   uint64_t
   getEditDistance(ArrayRef<const BinaryBasicBlock *> OldBlockOrder) const;
 
   /// True if the function is split into at most 2 fragments. Mostly used for
   /// checking whether a function can be processed in places that do not support
   /// multiple fragments yet.
   bool isHotColdSplit() const { return fragment_size() <= 2; }
 
   size_t fragment_size() const {
-    assert(Fragments.size() >= 2 &&
+    assert(Fragments.size() >= 1 &&
            "Layout should have at least one fragment.");
-    return Fragments.size() - 1;
+    return Fragments.size();
   }
-  bool fragment_empty() const { return Fragments.size() == 1; }
-  const_iterator fragment_begin() const { return {FragmentNum(0), this}; }
-  const_iterator fragment_end() const {
-    return {FragmentNum(fragment_size()), this};
+  bool fragment_empty() const { return fragment_size() == 0; }
+
+  fragment_iterator fragment_begin() { return Fragments.begin(); }
+  fragment_const_iterator fragment_begin() const { return Fragments.begin(); }
+  fragment_iterator fragment_end() { return Fragments.end(); }
+  fragment_const_iterator fragment_end() const { return Fragments.end(); }
+  iterator_range<fragment_iterator> fragments() {
+    return {fragment_begin(), fragment_end()};
   }
-  iterator_range<const_iterator> fragments() const {
+  iterator_range<fragment_const_iterator> fragments() const {
     return {fragment_begin(), fragment_end()};
   }
 
   size_t block_size() const { return Blocks.size(); }
   bool block_empty() const { return Blocks.empty(); }
+
+  /// Required to return non-const qualified `BinaryBasicBlock *` for graph
+  /// traits.
   BinaryBasicBlock *block_front() const { return Blocks.front(); }
-  BinaryBasicBlock *block_back() const { return Blocks.back(); }
-  block_const_iterator block_begin() const { return Blocks.begin(); }
-  block_const_iterator block_end() const { return Blocks.end(); }
+  const BinaryBasicBlock *block_back() const { return Blocks.back(); }
+
+  block_iterator block_begin() { return Blocks.begin(); }
+  block_const_iterator block_begin() const {
+    return block_const_iterator(Blocks.begin());
+  }
+  block_iterator block_end() { return Blocks.end(); }
+  block_const_iterator block_end() const {
+    return block_const_iterator(Blocks.end());
+  }
+  iterator_range<block_iterator> blocks() {
+    return {block_begin(), block_end()};
+  }
   iterator_range<block_const_iterator> blocks() const {
     return {block_begin(), block_end()};
   }
+
+  block_reverse_iterator block_rbegin() { return Blocks.rbegin(); }
   block_const_reverse_iterator block_rbegin() const { return Blocks.rbegin(); }
+  block_reverse_iterator block_rend() { return Blocks.rend(); }
   block_const_reverse_iterator block_rend() const { return Blocks.rend(); }
   iterator_range<block_const_reverse_iterator> rblocks() const {
     return {block_rbegin(), block_rend()};
   }
 
 private:
   block_const_iterator findBasicBlockPos(const BinaryBasicBlock *BB) const;
   block_iterator findBasicBlockPos(const BinaryBasicBlock *BB);
 
   friend class FunctionFragment;
 };
 
 } // namespace bolt
 } // namespace llvm
 
 #endif

bolt/lib/Core/BinaryContext.cpp

BinaryContext::calculateEmittedSize(BinaryFunction &BF, bool FixBranches) {
   Streamer->switchSection(Section);
   Streamer->emitLabel(StartLabel);
   emitFunctionBody(*Streamer, BF, BF.getLayout().getMainFragment(),
                    /*EmitCodeOnly=*/true);
   Streamer->emitLabel(EndLabel);
 
   using LabelRange = std::pair<const MCSymbol *, const MCSymbol *>;
   SmallVector<LabelRange> SplitLabels;
-  for (const FunctionFragment FF : BF.getLayout().getSplitFragments()) {
+  for (FunctionFragment &FF : BF.getLayout().getSplitFragments()) {
     MCSymbol *const SplitStartLabel = LocalCtx->createTempSymbol();
     MCSymbol *const SplitEndLabel = LocalCtx->createTempSymbol();
     SplitLabels.emplace_back(SplitStartLabel, SplitEndLabel);
 
     MCSectionELF *const SplitSection = LocalCtx->getELFSection(
         BF.getCodeSectionName(FF.getFragmentNum()), ELF::SHT_PROGBITS,
         ELF::SHF_EXECINSTR | ELF::SHF_ALLOC);
     SplitSection->setHasInstructions(true);

bolt/lib/Core/BinaryEmitter.cpp

public:
   BinaryEmitter(MCStreamer &Streamer, BinaryContext &BC)
       : Streamer(Streamer), BC(BC) {}
 
   /// Emit all code and data.
   void emitAll(StringRef OrgSecPrefix);
 
   /// Emit function code. The caller is responsible for emitting function
   /// symbol(s) and setting the section to emit the code to.
-  void emitFunctionBody(BinaryFunction &BF, const FunctionFragment &FF,
+  void emitFunctionBody(BinaryFunction &BF, FunctionFragment &FF,
                         bool EmitCodeOnly = false);
 
 private:
   /// Emit function code.
   void emitFunctions();
 
   /// Emit a single function.
-  bool emitFunction(BinaryFunction &BF, const FunctionFragment &FF);
+  bool emitFunction(BinaryFunction &BF, FunctionFragment &FF);
 
   /// Helper for emitFunctionBody to write data inside a function
   /// (used for AArch64)
   void emitConstantIslands(BinaryFunction &BF, bool EmitColdPart,
                            BinaryFunction *OnBehalfOf = nullptr);
 
   /// Emit jump tables for the function.
   void emitJumpTables(const BinaryFunction &BF);
for (BinaryFunction *Function : Functions) {
       // Was any part of the function emitted.
       bool Emitted = false;
 
       // Turn off Intel JCC Erratum mitigation for cold code if requested
       if (HasProfile && opts::X86AlignBranchBoundaryHotOnly &&
           !Function->hasValidProfile())
         Streamer.setAllowAutoPadding(false);
 
-      const FunctionLayout &Layout = Function->getLayout();
+      FunctionLayout &Layout = Function->getLayout();
       Emitted |= emitFunction(*Function, Layout.getMainFragment());
 
       if (Function->isSplit()) {
         if (opts::X86AlignBranchBoundaryHotOnly)
           Streamer.setAllowAutoPadding(false);
 
         assert((Layout.fragment_size() == 1 || Function->isSimple()) &&
                "Only simple functions can have fragments");
-        for (const FunctionFragment FF : Layout.getSplitFragments()) {
+        for (FunctionFragment &FF : Layout.getSplitFragments()) {
           // Skip empty fragments so no symbols and sections for empty fragments
           // are generated
           if (FF.empty() && !Function->hasConstantIsland())
             continue;
           Emitted |= emitFunction(*Function, FF);
         }
       }
 
void BinaryEmitter::emitFunctions() {
   // Mark the end of hot text.
   if (opts::HotText) {
     Streamer.switchSection(BC.getTextSection());
     Streamer.emitLabel(BC.getHotTextEndSymbol());
   }
 }
 
 bool BinaryEmitter::emitFunction(BinaryFunction &Function,
-                                 const FunctionFragment &FF) {
+                                 FunctionFragment &FF) {
   if (Function.size() == 0 && !Function.hasIslandsInfo())
     return false;
 
   if (Function.getState() == BinaryFunction::State::Empty)
     return false;
 
   MCSection *Section =
       BC.getCodeSection(Function.getCodeSectionName(FF.getFragmentNum()));
bool BinaryEmitter::emitFunction(BinaryFunction &Function,
   emitLSDA(Function, FF.isSplitFragment());
 
   if (FF.isMainFragment() && opts::JumpTables > JTS_NONE)
     emitJumpTables(Function);
 
   return true;
 }
 
-void BinaryEmitter::emitFunctionBody(BinaryFunction &BF,
-                                     const FunctionFragment &FF,
+void BinaryEmitter::emitFunctionBody(BinaryFunction &BF, FunctionFragment &FF,
                                      bool EmitCodeOnly) {
   if (!EmitCodeOnly && FF.isSplitFragment() && BF.hasConstantIsland()) {
     assert(BF.getLayout().isHotColdSplit() &&
            "Constant island support only with hot/cold split");
     BF.duplicateConstantIslands();
   }
 
   // Track the first emitted instruction with debug info.

 namespace bolt {
 
 void emitBinaryContext(MCStreamer &Streamer, BinaryContext &BC,
                        StringRef OrgSecPrefix) {
   BinaryEmitter(Streamer, BC).emitAll(OrgSecPrefix);
 }
 
 void emitFunctionBody(MCStreamer &Streamer, BinaryFunction &BF,
-                      const FunctionFragment &FF, bool EmitCodeOnly) {
+                      FunctionFragment &FF, bool EmitCodeOnly) {
   BinaryEmitter(Streamer, BF.getBinaryContext())
       .emitFunctionBody(BF, FF, EmitCodeOnly);
 }
 
 } // namespace bolt
 } // namespace llvm

bolt/lib/Core/BinaryFunction.cpp

for (const std::pair<const uint32_t, MCInst> &II : Instructions) {
         OS << LI->second->getName() << ":\n";
       }
 
       BC.printInstruction(OS, II.second, Offset, this);
     }
   }
 
   StringRef SplitPointMsg = "";
-  for (const FunctionFragment FF : Layout.fragments()) {
+  for (const FunctionFragment &FF : Layout.fragments()) {
     OS << SplitPointMsg;
     SplitPointMsg = "-------   HOT-COLD SPLIT POINT   -------\n\n";
     for (const BinaryBasicBlock *BB : FF) {
       OS << BB->getName() << " (" << BB->size()
          << " instructions, align : " << BB->getAlignment() << ")\n";
 
       if (isEntryPoint(*BB)) {
         if (MCSymbol *EntrySymbol = getSecondaryEntryPointSymbol(*BB))

 bool BinaryFunction::finalizeCFIState() {
   LLVM_DEBUG(
       dbgs() << "Trying to fix CFI states for each BB after reordering.\n");
   LLVM_DEBUG(dbgs() << "This is the list of CFI states for each BB of " << *this
                     << ": ");
 
   const char *Sep = "";
   (void)Sep;
-  for (const FunctionFragment FF : Layout.fragments()) {
+  for (FunctionFragment &FF : Layout.fragments()) {
     // Hot-cold border: at start of each region (with a different FDE) we need
     // to reset the CFI state.
     int32_t State = 0;
 
     for (BinaryBasicBlock *BB : FF) {
       const int32_t CFIStateAtExit = BB->getCFIStateAtExit();
 
       // We need to recover the correct state if it doesn't match expected
void BinaryFunction::updateOutputValues(const MCAsmLayout &Layout) {
 
   assert((getLayout().isHotColdSplit() ||
           (cold().getAddress() == 0 && cold().getImageSize() == 0 &&
            BC.HasRelocations)) &&
          "Function must be split two ways or cold fragment must have no "
          "address (only in relocation mode)");
 
   BinaryBasicBlock *PrevBB = nullptr;
-  for (const FunctionFragment &FF : getLayout().fragments()) {
+  for (FunctionFragment &FF : getLayout().fragments()) {
     const uint64_t FragmentBaseAddress =
         getCodeSection(isSimple() ? FF.getFragmentNum() : FragmentNum::main())
             ->getOutputAddress();
     for (BinaryBasicBlock *const BB : FF) {
       assert(BB->getLabel()->isDefined() && "symbol should be defined");
       if (!BC.HasRelocations) {
         if (BB->isSplit()) {
           assert(FragmentBaseAddress == cold().getAddress());

bolt/lib/Core/Exceptions.cpp

void BinaryFunction::updateEHRanges() {
   assert(CurrentState == State::CFG_Finalized && "unexpected state");
 
   // Build call sites table.
   struct EHInfo {
     const MCSymbol *LP; // landing pad
     uint64_t Action;
   };
 
-  for (const FunctionFragment FF : getLayout().fragments()) {
+  for (FunctionFragment &FF : getLayout().fragments()) {
     // Sites to update - either regular or cold.
     CallSitesType &Sites = FF.isMainFragment() ? CallSites : ColdCallSites;
 
     // If previous call can throw, this is its exception handler.
     EHInfo PreviousEH = {nullptr, 0};
 
     // Marker for the beginning of exceptions range.
     const MCSymbol *StartRange = nullptr;

bolt/lib/Core/FunctionLayout.cpp

 #include "bolt/Core/FunctionLayout.h"
 #include "bolt/Core/BinaryFunction.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/edit_distance.h"
 #include <algorithm>
 #include <cstddef>
 #include <functional>
+#include <iterator>
+#include <memory>
 
 using namespace llvm;
 using namespace bolt;
 
-unsigned FunctionFragment::size() const { return end() - begin(); }
-bool FunctionFragment::empty() const { return end() == begin(); }
+FunctionFragment::FunctionFragment(FunctionLayout &Layout,
+                                   const FragmentNum Num)
+    : Layout(&Layout), Num(Num), StartIndex(Layout.block_size()) {}
+
+FunctionFragment::iterator FunctionFragment::begin() {
+  return iterator(Layout->block_begin() + StartIndex);
+}
 FunctionFragment::const_iterator FunctionFragment::begin() const {
-  return Layout.block_begin() + Layout.Fragments[Num.get()];
+  return const_iterator(Layout->block_begin() + StartIndex);
+}
+FunctionFragment::iterator FunctionFragment::end() {
+  return iterator(Layout->block_begin() + StartIndex + Size);
 }
 FunctionFragment::const_iterator FunctionFragment::end() const {
-  return Layout.block_begin() + Layout.Fragments[Num.get() + 1];
+  return const_iterator(Layout->block_begin() + StartIndex + Size);
+}
+
+const BinaryBasicBlock *FunctionFragment::front() const { return *begin(); }
+
+FunctionLayout::FunctionLayout() { addFragment(); }
+
+FunctionLayout::FunctionLayout(const FunctionLayout &Other)
+    : Blocks(Other.Blocks) {
+  for (FunctionFragment *const FF : Other.Fragments) {
+    auto *Copy = new FunctionFragment(*FF);
+    Copy->Layout = this;
+    Fragments.emplace_back(Copy);
+  }
+}
+
+FunctionLayout::FunctionLayout(FunctionLayout &&Other)
+    : Fragments(std::move(Other.Fragments)), Blocks(std::move(Other.Blocks)) {
+  for (FunctionFragment *const F : Fragments)
+    F->Layout = this;
+}
+
+FunctionLayout &FunctionLayout::operator=(const FunctionLayout &Other) {
+  Blocks = Other.Blocks;
+  for (FunctionFragment *const FF : Other.Fragments) {
+    auto *const Copy = new FunctionFragment(*FF);
+    Copy->Layout = this;
+    Fragments.emplace_back(Copy);
+  }
+  return *this;
+}
+
+FunctionLayout &FunctionLayout::operator=(FunctionLayout &&Other) {
+  Fragments = std::move(Other.Fragments);
+  Blocks = std::move(Other.Blocks);
+  for (FunctionFragment *const FF : Fragments)
+    FF->Layout = this;
+  return *this;
+}
+
+FunctionLayout::~FunctionLayout() {
+  for (FunctionFragment *const F : Fragments) {
+    delete F;
+  }
+}
+
+FunctionFragment &FunctionLayout::addFragment() {
+  FunctionFragment *const FF =
+      new FunctionFragment(*this, FragmentNum(Fragments.size()));
+  Fragments.emplace_back(FF);
+  return *FF;
 }
-BinaryBasicBlock *FunctionFragment::front() const { return *begin(); }
 
-FunctionFragment FunctionLayout::addFragment() {
-  Fragments.emplace_back(Blocks.size());
-  return getFragment(FragmentNum(Blocks.size() - 1));
+FunctionFragment &FunctionLayout::getFragment(FragmentNum Num) {
+  return *Fragments[Num.get()];
 }
 
-FunctionFragment FunctionLayout::getFragment(FragmentNum Num) const {
-  return FunctionFragment(Num, *this);
+const FunctionFragment &FunctionLayout::getFragment(FragmentNum Num) const {
+  return *Fragments[Num.get()];
 }
 
-FunctionFragment
-FunctionLayout::findFragment(const BinaryBasicBlock *BB) const {
+const FunctionFragment &
+FunctionLayout::findFragment(const BinaryBasicBlock *const BB) const {
   return getFragment(BB->getFragmentNum());
 }
 
-void FunctionLayout::addBasicBlock(BinaryBasicBlock *BB) {
+void FunctionLayout::addBasicBlock(BinaryBasicBlock *const BB) {
   BB->setLayoutIndex(Blocks.size());
   Blocks.emplace_back(BB);
-  ++Fragments.back();
-  assert(Fragments.back() == Blocks.size());
+  Fragments.back()->Size++;
 }
 
-void FunctionLayout::insertBasicBlocks(BinaryBasicBlock *InsertAfter,
-                                       ArrayRef<BinaryBasicBlock *> NewBlocks) {
-  const block_iterator InsertBeforePos =
-      InsertAfter ? std::next(findBasicBlockPos(InsertAfter)) : Blocks.begin();
-  Blocks.insert(InsertBeforePos, NewBlocks.begin(), NewBlocks.end());
-
-  unsigned FragmentUpdateStart =
-      InsertAfter ? InsertAfter->getFragmentNum().get() + 1 : 1;
-  std::for_each(
-      Fragments.begin() + FragmentUpdateStart, Fragments.end(),
-      [&](unsigned &FragmentOffset) { FragmentOffset += NewBlocks.size(); });
+void FunctionLayout::insertBasicBlocks(
+    const BinaryBasicBlock *const InsertAfter,
+    const ArrayRef<BinaryBasicBlock *> NewBlocks) {
+  block_iterator InsertBeforePos = Blocks.begin();
+  FragmentNum InsertFragmentNum = FragmentNum::main();
+  unsigned LayoutIndex = 0;
+
+  if (InsertAfter) {
+    InsertBeforePos = std::next(findBasicBlockPos(InsertAfter));
+    InsertFragmentNum = InsertAfter->getFragmentNum();
+    LayoutIndex = InsertAfter->getLayoutIndex();
+  }
+
+  llvm::copy(NewBlocks, std::inserter(Blocks, InsertBeforePos));
+
+  for (BinaryBasicBlock *const BB : NewBlocks) {
+    BB->setFragmentNum(InsertFragmentNum);
+    BB->setLayoutIndex(LayoutIndex++);
+  }
+
+  const fragment_iterator InsertFragment =
+      fragment_begin() + InsertFragmentNum.get();
+  InsertFragment->Size += NewBlocks.size();
+
+  const fragment_iterator TailBegin = std::next(InsertFragment);
+  auto const UpdateFragment = [&](FunctionFragment &FF) {
+    FF.StartIndex += NewBlocks.size();
+    for (BinaryBasicBlock *const BB : FF)
+      BB->setLayoutIndex(LayoutIndex++);
+  };
+  std::for_each(TailBegin, fragment_end(), UpdateFragment);
 }
 
 void FunctionLayout::eraseBasicBlocks(
     const DenseSet<const BinaryBasicBlock *> ToErase) {
-  auto IsErased = [&](const BinaryBasicBlock *const BB) {
+  const auto IsErased = [&](const BinaryBasicBlock *const BB) {
     return ToErase.contains(BB);
   };
-  FragmentListType NewFragments;
-  NewFragments.emplace_back(0);
-  for (const FunctionFragment FF : fragments()) {
-    unsigned ErasedBlocks = count_if(FF, IsErased);
-    // Only add the fragment if it is non-empty after removing blocks.
-    unsigned NewFragment = NewFragments.back() + FF.size() - ErasedBlocks;
-    NewFragments.emplace_back(NewFragment);
+
+  unsigned TotalErased = 0;
+  for (FunctionFragment &FF : fragments()) {
+    unsigned Erased = count_if(FF, IsErased);
+    FF.Size -= Erased;
+    FF.StartIndex -= TotalErased;
+    TotalErased += Erased;
   }
   llvm::erase_if(Blocks, IsErased);
-  Fragments = std::move(NewFragments);
 
   // Remove empty fragments at the end
-  const_iterator EmptyTailBegin =
-      llvm::find_if_not(reverse(fragments()), [](const FunctionFragment &FF) {
-        return FF.empty();
-      }).base();
-  if (EmptyTailBegin != fragment_end()) {
-    // Add +1 for one-past-the-end entry
-    const FunctionFragment TailBegin = *EmptyTailBegin;
-    unsigned NewFragmentSize = TailBegin.getFragmentNum().get() + 1;
-    Fragments.resize(NewFragmentSize);
-  }
+  const auto IsEmpty = [](const FunctionFragment *const FF) {
+    return FF->empty();
+  };
+  const FragmentListType::iterator EmptyTailBegin =
+      llvm::find_if_not(reverse(Fragments), IsEmpty).base();
+  std::for_each(EmptyTailBegin, Fragments.end(),
+                [](FunctionFragment *const FF) { delete FF; });
+  Fragments.erase(EmptyTailBegin, Fragments.end());
 
   updateLayoutIndices();
 }
 
-void FunctionLayout::updateLayoutIndices() const {
+void FunctionLayout::updateLayoutIndices() {
   unsigned BlockIndex = 0;
-  for (const FunctionFragment FF : fragments()) {
+  for (FunctionFragment &FF : fragments()) {
     for (BinaryBasicBlock *const BB : FF) {
       BB->setLayoutIndex(BlockIndex++);
       BB->setFragmentNum(FF.getFragmentNum());
     }
   }
 }
 
 bool FunctionLayout::update(const ArrayRef<BinaryBasicBlock *> NewLayout) {
   const bool EqualBlockOrder = llvm::equal(Blocks, NewLayout);
   if (EqualBlockOrder) {
     const bool EqualPartitioning =
-        llvm::all_of(fragments(), [](const FunctionFragment FF) {
+        llvm::all_of(fragments(), [](const FunctionFragment &FF) {
           return llvm::all_of(FF, [&](const BinaryBasicBlock *const BB) {
             return FF.Num == BB->getFragmentNum();
           });
         });
     if (EqualPartitioning)
       return false;
   }
 
-  Blocks = BasicBlockListType(NewLayout.begin(), NewLayout.end());
-  Fragments = {0, 0};
+  clear();
 
   // Generate fragments
-  for (const auto &BB : enumerate(Blocks)) {
-    unsigned FragmentNum = BB.value()->getFragmentNum().get();
+  for (BinaryBasicBlock *const BB : NewLayout) {
+    FragmentNum Num = BB->getFragmentNum();
 
-    assert(FragmentNum >= fragment_size() - 1 &&
+    assert(Num >= Fragments.back()->getFragmentNum() &&
            "Blocks must be arranged such that fragments are monotonically "
            "increasing.");
 
     // Add empty fragments if necessary
-    for (unsigned I = fragment_size(); I <= FragmentNum; ++I) {
+    while (Fragments.back()->getFragmentNum() < Num)
       addFragment();
-      Fragments[I] = BB.index();
-    }
 
     // Set the next fragment to point one past the current BB
-    Fragments[FragmentNum + 1] = BB.index() + 1;
+    addBasicBlock(BB);
   }
 
   return true;
 }
 
 void FunctionLayout::clear() {
-  Blocks = {};
-  Fragments = {0, 0};
+  Blocks = BasicBlockListType();
+  for (FunctionFragment *const FF : Fragments)
+    delete FF;
+  Fragments = FragmentListType();
+  addFragment();
 }
 
-BinaryBasicBlock *FunctionLayout::getBasicBlockAfter(const BinaryBasicBlock *BB,
-                                                     bool IgnoreSplits) const {
-  const block_const_iterator BBPos = find(Blocks, BB);
-  if (BBPos == Blocks.end())
+const BinaryBasicBlock *
+FunctionLayout::getBasicBlockAfter(const BinaryBasicBlock *BB,
+  const block_const_iterator BBPos = find(blocks(), BB);
+  if (BBPos == block_end())
     return nullptr;
 
   const block_const_iterator BlockAfter = std::next(BBPos);
-  if (BlockAfter == Blocks.end())
+  if (BlockAfter == block_end())
     return nullptr;
 
   if (!IgnoreSplits)
     if (BlockAfter == getFragment(BB->getFragmentNum()).end())
       return nullptr;
 
   return *BlockAfter;
 }
 
 bool FunctionLayout::isSplit() const {
-  unsigned NonEmptyFragCount = llvm::count_if(
+  const unsigned NonEmptyFragCount = llvm::count_if(
       fragments(), [](const FunctionFragment &FF) { return !FF.empty(); });
   return NonEmptyFragCount >= 2;
 }
 
 uint64_t FunctionLayout::getEditDistance(
     const ArrayRef<const BinaryBasicBlock *> OldBlockOrder) const {
   return ComputeEditDistance<const BinaryBasicBlock *>(OldBlockOrder, Blocks);
 }
 
 FunctionLayout::block_const_iterator
 FunctionLayout::findBasicBlockPos(const BinaryBasicBlock *BB) const {
-  return find(Blocks, BB);
+  return block_const_iterator(find(Blocks, BB));
 }
 
 FunctionLayout::block_iterator
 FunctionLayout::findBasicBlockPos(const BinaryBasicBlock *BB) {
   return find(Blocks, BB);
 }

bolt/lib/Passes/BinaryPasses.cpp

 }
 
 void LowerAnnotations::runOnFunctions(BinaryContext &BC) {
   std::vector<std::pair<MCInst *, uint32_t>> PreservedOffsetAnnotations;
 
   for (auto &It : BC.getBinaryFunctions()) {
     BinaryFunction &BF = It.second;
 
-    for (const FunctionFragment FF : BF.getLayout().fragments()) {
+    for (FunctionFragment &FF : BF.getLayout().fragments()) {
       int64_t CurrentGnuArgsSize = 0;
 
       for (BinaryBasicBlock *const BB : FF) {
         // First convert GnuArgsSize annotations into CFIs. This may change
         // instr pointers, so do it before recording ptrs for preserved
         // annotations
         if (BF.usesGnuArgsSize()) {
           for (auto II = BB->begin(); II != BB->end(); ++II) {

bolt/lib/Rewrite/RewriteInstance.cpp

if (Section &&
       continue;
 
     // Restore origin section for functions that were emitted or supposed to
     // be emitted to patch sections.
     if (Section)
       BC->deregisterSection(*Section);
     assert(Function->getOriginSectionName() && "expected origin section");
     Function->CodeSectionName = Function->getOriginSectionName()->str();
-    for (const FunctionFragment FF :
+    for (const FunctionFragment &FF :
          Function->getLayout().getSplitFragments()) {
       if (ErrorOr<BinarySection &> ColdSection =
               Function->getCodeSection(FF.getFragmentNum()))
         BC->deregisterSection(*ColdSection);
     }
     if (Function->getLayout().isSplit())
       Function->setColdCodeSectionName(getBOLTTextSectionName());
   }
if (opts::JumpTables == JTS_BASIC) {
                           << '\n');
         RTDyld.reassignSectionAddress(Section.getSectionID(), JT->getAddress());
       }
     }
 
     if (!Function.isSplit())
       continue;
 
-    for (const FunctionFragment FF : Function.getLayout().getSplitFragments()) {
+    for (const FunctionFragment &FF :
+         Function.getLayout().getSplitFragments()) {
       ErrorOr<BinarySection &> ColdSection =
           Function.getCodeSection(FF.getFragmentNum());
       assert(ColdSection && "cannot find section for cold part");
       // Cold fragments are aligned at 16 bytes.
       NextAvailableAddress = alignTo(NextAvailableAddress, 16);
       BinaryFunction::FragmentInfo &ColdPart = Function.cold();
       if (TooLarge) {
         // The corresponding FDE will refer to address 0.
if (Function.isFolded()) {
                           .concat(".icf.0")
                           .toStringRef(Buf));
       ICFSymbol.st_value = ICFParent->getOutputAddress();
       ICFSymbol.st_size = ICFParent->getOutputSize();
       ICFSymbol.st_shndx = ICFParent->getCodeSection()->getIndex();
       Symbols.emplace_back(ICFSymbol);
     }
     if (Function.isSplit() && Function.cold().getAddress()) {
-      for (const FunctionFragment FF :
+      for (const FunctionFragment &FF :
            Function.getLayout().getSplitFragments()) {
         ELFSymTy NewColdSym = FunctionSymbol;
         const SmallString<256> SymbolName = formatv(
             "{0}.cold.{1}", cantFail(FunctionSymbol.getName(StringSection)),
             FF.getFragmentNum().get() - 1);
         NewColdSym.st_name = AddToStrTab(SymbolName);
         NewColdSym.st_shndx =
             Function.getCodeSection(FF.getFragmentNum())->getIndex();