[BOLT] Allocate FunctionFragment on heap

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.

Reviewed By: rafauler

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

Author: Fabian Parzefall

bolt/include/bolt/Core/BinaryEmitter.h

@@ -35,7 +35,7 @@ void emitBinaryContext(MCStreamer &Streamer, BinaryContext &BC,
 /// 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

bolt/include/bolt/Core/FunctionLayout.h

@@ -24,6 +24,8 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/ADT/iterator_range.h"
+#include <iterator>
+#include <utility>
 
 namespace llvm {
 namespace bolt {
@@ -69,28 +71,41 @@ class FunctionFragment {
   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
@@ -103,102 +118,84 @@ class FunctionFragment {
 /// 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
@@ -209,12 +206,25 @@ class FunctionLayout {
   /// 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 *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.
-  BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *BB,
-                                       bool IgnoreSplits = true) const;
+  const BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *BB,
+                                             bool IgnoreSplits = true) const;
 
   /// True if the layout contains at least two non-empty fragments.
   bool isSplit() const;
@@ -230,29 +240,49 @@ class FunctionLayout {
   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()};

bolt/lib/Core/BinaryContext.cpp

@@ -2199,7 +2199,7 @@ BinaryContext::calculateEmittedSize(BinaryFunction &BF, bool FixBranches) {
 
   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);

bolt/lib/Core/BinaryEmitter.cpp

@@ -129,15 +129,15 @@ class BinaryEmitter {
 
   /// 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)
@@ -234,7 +234,7 @@ void BinaryEmitter::emitFunctions() {
           !Function->hasValidProfile())
         Streamer.setAllowAutoPadding(false);
 
-      const FunctionLayout &Layout = Function->getLayout();
+      FunctionLayout &Layout = Function->getLayout();
       Emitted |= emitFunction(*Function, Layout.getMainFragment());
 
       if (Function->isSplit()) {
@@ -243,7 +243,7 @@ void BinaryEmitter::emitFunctions() {
 
         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())
@@ -280,7 +280,7 @@ void BinaryEmitter::emitFunctions() {
 }
 
 bool BinaryEmitter::emitFunction(BinaryFunction &Function,
-                                 const FunctionFragment &FF) {
+                                 FunctionFragment &FF) {
   if (Function.size() == 0 && !Function.hasIslandsInfo())
     return false;
 
@@ -402,8 +402,7 @@ bool BinaryEmitter::emitFunction(BinaryFunction &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() &&
@@ -1160,7 +1159,7 @@ void emitBinaryContext(MCStreamer &Streamer, BinaryContext &BC,
 }
 
 void emitFunctionBody(MCStreamer &Streamer, BinaryFunction &BF,
-                      const FunctionFragment &FF, bool EmitCodeOnly) {
+                      FunctionFragment &FF, bool EmitCodeOnly) {
   BinaryEmitter(Streamer, BF.getBinaryContext())
       .emitFunctionBody(BF, FF, EmitCodeOnly);
 }