2 files changed, 20 insertions, 18 deletions

diff --git a/include/llvm/Analysis/BlockFrequencyInfoImpl.h b/include/llvm/Analysis/BlockFrequencyInfoImpl.h
index bea9f79aeb..ca98a2e1d2 100644
--- a/include/llvm/Analysis/BlockFrequencyInfoImpl.h
+++ b/include/llvm/Analysis/BlockFrequencyInfoImpl.h
@@ -945,28 +945,30 @@ public:
     typedef SmallVector<std::pair<BlockNode, BlockMass>, 4> ExitMap;
     typedef SmallVector<BlockNode, 4> MemberList;
     BlockNode Header;       ///< Header.
+    bool IsPackaged;        ///< Whether this has been packaged.
     ExitMap Exits;          ///< Successor edges (and weights).
     MemberList Members;     ///< Members of the loop.
     BlockMass BackedgeMass; ///< Mass returned to loop header.
     BlockMass Mass;
     Float Scale;
 
-    LoopData(const BlockNode &Header) : Header(Header) {}
+    LoopData(const BlockNode &Header) : Header(Header), IsPackaged(false) {}
   };
 
   /// \brief Index of loop information.
   struct WorkingData {
+    LoopData *Loop;           ///< The loop this block is the header of.
     BlockNode ContainingLoop; ///< The block whose loop this block is inside.
-    uint32_t LoopIndex;       ///< Index into PackagedLoops.
     bool IsPackaged;          ///< Has ContainingLoop been packaged up?
-    bool IsAPackage;          ///< Has this block's loop been packaged up?
     BlockMass Mass;           ///< Mass distribution from the entry block.
 
-    WorkingData()
-        : LoopIndex(UINT32_MAX), IsPackaged(false), IsAPackage(false) {}
+    WorkingData() : Loop(nullptr), IsPackaged(false) {}
 
     bool hasLoopHeader() const { return ContainingLoop.isValid(); }
-    bool isLoopHeader() const { return LoopIndex != UINT32_MAX; }
+    bool isLoopHeader() const { return Loop; }
+
+    /// \brief Has this block's loop been packaged up?
+    bool isAPackage() const { return Loop && Loop->IsPackaged; }
   };
 
   /// \brief Unscaled probability weight.
@@ -1040,7 +1042,7 @@ public:
   std::vector<WorkingData> Working;
 
   /// \brief Indexed information about packaged loops.
-  std::vector<LoopData> PackagedLoops;
+  std::vector<std::unique_ptr<LoopData>> PackagedLoops;
 
   /// \brief Add all edges out of a packaged loop to the distribution.
   ///
@@ -1061,9 +1063,8 @@ public:
 
   LoopData &getLoopPackage(const BlockNode &Head) {
     assert(Head.Index < Working.size());
-    size_t Index = Working[Head.Index].LoopIndex;
-    assert(Index < PackagedLoops.size());
-    return PackagedLoops[Index];
+    assert(Working[Head.Index].Loop != nullptr);
+    return *Working[Head.Index].Loop;
   }
 
   /// \brief Distribute mass according to a distribution.
@@ -1428,8 +1429,8 @@ template <class BT> void BlockFrequencyInfoImpl<BT>::initializeLoops() {
     BlockNode Header = getNode(Loop->getHeader());
     assert(Header.isValid());
 
-    Working[Header.Index].LoopIndex = PackagedLoops.size();
-    PackagedLoops.emplace_back(Header);
+    PackagedLoops.emplace_back(new LoopData(Header));
+    Working[Header.Index].Loop = PackagedLoops.back().get();
     DEBUG(dbgs() << " - loop = " << getBlockName(Header) << "\n");
   }
 
@@ -1452,7 +1453,7 @@ template <class BT> void BlockFrequencyInfoImpl<BT>::initializeLoops() {
     assert(HeaderData.isLoopHeader());
 
     Working[Index].ContainingLoop = Header;
-    PackagedLoops[HeaderData.LoopIndex].Members.push_back(Index);
+    HeaderData.Loop->Members.push_back(Index);
     DEBUG(dbgs() << " - loop = " << getBlockName(Header)
                  << ": member = " << getBlockName(Index) << "\n");
   }
@@ -1460,7 +1461,7 @@ template <class BT> void BlockFrequencyInfoImpl<BT>::initializeLoops() {
 
 template <class BT> void BlockFrequencyInfoImpl<BT>::computeMassInLoops() {
   // Visit loops with the deepest first, and the top-level loops last.
-  for (auto L = PackagedLoops.rbegin(), LE = PackagedLoops.rend(); L != LE; ++L)
+  for (const auto &L : make_range(PackagedLoops.rbegin(), PackagedLoops.rend()))
     computeMassInLoop(L->Header);
 }
 
diff --git a/lib/Analysis/BlockFrequencyInfoImpl.cpp b/lib/Analysis/BlockFrequencyInfoImpl.cpp
index bc3722eb9d..909786d366 100644
--- a/lib/Analysis/BlockFrequencyInfoImpl.cpp
+++ b/lib/Analysis/BlockFrequencyInfoImpl.cpp
@@ -602,7 +602,7 @@ void BlockFrequencyInfoImplBase::clear() {
   // does not actually clear heap storage.
   std::vector<FrequencyData>().swap(Freqs);
   std::vector<WorkingData>().swap(Working);
-  std::vector<LoopData>().swap(PackagedLoops);
+  std::vector<std::unique_ptr<LoopData>>().swap(PackagedLoops);
 }
 
 /// \brief Clear all memory not needed downstream.
@@ -646,7 +646,7 @@ static BlockMass &getPackageMass(BlockFrequencyInfoImplBase &BFI,
                                  const BlockNode &Node) {
   assert(Node.isValid());
   assert(!BFI.Working[Node.Index].IsPackaged);
-  if (!BFI.Working[Node.Index].IsAPackage)
+  if (!BFI.Working[Node.Index].isAPackage())
     return BFI.Working[Node.Index].Mass;
 
   return BFI.getLoopPackage(Node).Mass;
@@ -744,8 +744,9 @@ void BlockFrequencyInfoImplBase::computeLoopScale(const BlockNode &LoopHead) {
 /// \brief Package up a loop.
 void BlockFrequencyInfoImplBase::packageLoop(const BlockNode &LoopHead) {
   DEBUG(dbgs() << "packaging-loop: " << getBlockName(LoopHead) << "\n");
-  Working[LoopHead.Index].IsAPackage = true;
-  for (const BlockNode &M : getLoopPackage(LoopHead).Members) {
+  auto &PackagedLoop = getLoopPackage(LoopHead);
+  PackagedLoop.IsPackaged = true;
+  for (const BlockNode &M : PackagedLoop.Members) {
     DEBUG(dbgs() << " - node: " << getBlockName(M.Index) << "\n");
     Working[M.Index].IsPackaged = true;
   }