[BOLT] Fix memory leak in dataflowanalysis

Do this by making sure that the underlying BumpPtr allocator can be freed if
the DataflowAnalysis is only instantiated for a "short" period of time, i.e.,
does get destructed after analyzing one BinaryFunction and before analyzing
the next BinaryFunction.

Author: kbeyls

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -351,34 +351,71 @@ class MCPlusBuilder {
     return nullptr;
   }
 
-  /// Initialize a new annotation allocator and return its id
-  AllocatorIdTy initializeNewAnnotationAllocator() {
+private:
+  std::vector<AllocatorIdTy> AvailablePrivateAllocIds;
+  std::vector<AllocatorIdTy> AvailableSharedAllocIds;
+  llvm::sys::RWMutex AllocIdMgmtMutex;
+
+  AllocatorIdTy _initializeNewAnnotationAllocator() {
     AnnotationAllocators.emplace(MaxAllocatorId, AnnotationAllocator());
+    assert(MaxAllocatorId < std::numeric_limits<AllocatorIdTy>::max());
     return MaxAllocatorId++;
   }
 
+public:
+  AllocatorIdTy getPrivateAllocatorId() {
+    std::unique_lock<llvm::sys::RWMutex> Lock(AllocIdMgmtMutex);
+    if (AvailablePrivateAllocIds.size() > 0) {
+      AllocatorIdTy AllocId = AvailablePrivateAllocIds.back();
+      AvailablePrivateAllocIds.pop_back();
+      AnnotationAllocators.emplace(AllocId, AnnotationAllocator());
+      return AllocId;
+    }
+    AllocatorIdTy AllocId = _initializeNewAnnotationAllocator();
+    return AllocId;
+  }
+
+  std::vector<AllocatorIdTy> getSharedAllocIds(unsigned nrAllocs) {
+    std::unique_lock<llvm::sys::RWMutex> Lock(AllocIdMgmtMutex);
+    for (unsigned I = AvailableSharedAllocIds.size(); I <= nrAllocs; ++I) {
+      // allocate new shared alloc Ids
+      AllocatorIdTy AllocId = _initializeNewAnnotationAllocator();
+      AvailableSharedAllocIds.push_back(AllocId);
+    }
+    std::vector<AllocatorIdTy> Result;
+    for (unsigned I = 0; I < nrAllocs; ++I)
+      Result.push_back(AvailableSharedAllocIds[I]);
+    return Result;
+  }
+
+  void freePrivateAllocatorId(AllocatorIdTy AllocId) {
+    std::unique_lock<llvm::sys::RWMutex> Lock(AllocIdMgmtMutex);
+    assert(AnnotationAllocators.count(AllocId) != 0);
+    AnnotationAllocator &Allocator = AnnotationAllocators.find(AllocId)->second;
+    for (MCPlus::MCAnnotation *Annotation : Allocator.AnnotationPool)
+      Annotation->~MCAnnotation();
+
+    Allocator.AnnotationPool
+        .clear(); // FIXME: wouldn't this call the destructor anyway?
+    Allocator.ValueAllocator.Reset();
+    AnnotationAllocators.erase(AllocId);
+    AvailablePrivateAllocIds.push_back(AllocId);
+  }
+
   /// Return the annotation allocator of a given id
   AnnotationAllocator &getAnnotationAllocator(AllocatorIdTy AllocatorId) {
+    std::shared_lock<llvm::sys::RWMutex> Lock(AllocIdMgmtMutex);
     assert(AnnotationAllocators.count(AllocatorId) &&
            "allocator not initialized");
     return AnnotationAllocators.find(AllocatorId)->second;
   }
 
   // Check if an annotation allocator with the given id exists
   bool checkAllocatorExists(AllocatorIdTy AllocatorId) {
+    std::shared_lock<llvm::sys::RWMutex> Lock(AllocIdMgmtMutex);
     return AnnotationAllocators.count(AllocatorId);
   }
 
-  /// Free the values allocator within the annotation allocator
-  void freeValuesAllocator(AllocatorIdTy AllocatorId) {
-    AnnotationAllocator &Allocator = getAnnotationAllocator(AllocatorId);
-    for (MCPlus::MCAnnotation *Annotation : Allocator.AnnotationPool)
-      Annotation->~MCAnnotation();
-
-    Allocator.AnnotationPool.clear();
-    Allocator.ValueAllocator.Reset();
-  }
-
   virtual ~MCPlusBuilder() { freeAnnotations(); }
 
   /// Free all memory allocated for annotations

bolt/include/bolt/Passes/DataflowAnalysis.h

@@ -163,6 +163,8 @@ class DataflowAnalysis {
   /// The id of the annotation allocator to be used
   MCPlusBuilder::AllocatorIdTy AllocatorId = 0;
 
+  bool IsPrivateAllocatorId;
+
   /// Tracks the state at basic block start (end) if direction of the dataflow
   /// is forward (backward).
   std::unordered_map<const BinaryBasicBlock *, StateTy> StateAtBBEntry;
@@ -243,9 +245,6 @@ class DataflowAnalysis {
   }
 
 public:
-  /// Return the allocator id
-  unsigned getAllocatorId() { return AllocatorId; }
-
   /// If the direction of the dataflow is forward, operates on the last
   /// instruction of all predecessors when performing an iteration of the
   /// dataflow equation for the start of this BB.  If backwards, operates on
@@ -259,11 +258,40 @@ class DataflowAnalysis {
 
   /// We need the current binary context and the function that will be processed
   /// in this dataflow analysis.
-  DataflowAnalysis(BinaryFunction &BF,
-                   MCPlusBuilder::AllocatorIdTy AllocatorId = 0)
-      : BC(BF.getBinaryContext()), Func(BF), AllocatorId(AllocatorId) {}
-
-  virtual ~DataflowAnalysis() { cleanAnnotations(); }
+  ///
+  /// There are 2 constructors:
+  /// * One version takes an AllocId, where it is assumed that:
+  ///   * If multiple threads are used, each different thread will get a different
+  ///     AllocId.
+  ///   * AllocIds may be reused by other functions/objects/DataflowAnalyses that
+  ///     take an AllocId of where to allocate MCAnnotations.
+  ///     As a result, the MCAnnotations this dataflow analysis allocates will
+  ///     not get freed before the end of the program run.
+  /// * The second version takes a bool argument UsePrivateAllocators, which
+  ///   must be set to true.
+  ///   * When constructed this way, the dataflowanalysis will construct its own
+  ///     allocator, only used for MCAnnotations allocated by this DataFlowAnalysis.
+  ///   * When the DataflowAnalysis is destructed, that allocator gets freed,
+  ///     so that the memory used by this DataflowAnalysis gets completely freed.
+  ///   * Note however, that there cannot be a huge number of allocators all live
+  ///     at the same time. Therefore, only use this function when you're going
+  ///     to destruct this DataFlowAnalysis after analyzing one BinaryFunction,
+  ///     before creating a new DataFlowAnalysis to analyze another BinaryFunction.
+  DataflowAnalysis(BinaryFunction &BF, const bool UsePrivateAllocators)
+      : BC(BF.getBinaryContext()), Func(BF),
+        AllocatorId(BC.MIB->getPrivateAllocatorId()),
+        IsPrivateAllocatorId(true) {
+    assert(UsePrivateAllocators == true);
+  }
+  DataflowAnalysis(BinaryFunction &BF, MCPlusBuilder::AllocatorIdTy AllocId)
+      : BC(BF.getBinaryContext()), Func(BF), AllocatorId(AllocId),
+        IsPrivateAllocatorId(false) {}
+
+  virtual ~DataflowAnalysis() {
+    cleanAnnotations();
+    if (IsPrivateAllocatorId)
+      BC.MIB->freePrivateAllocatorId(AllocatorId);
+  }
 
   /// Track the state at basic block start (end) if direction of the dataflow
   /// is forward (backward).
@@ -529,8 +557,11 @@ class InstrsDataflowAnalysis
     return count(*Expressions[PointIdx], Expr);
   }
 
+  InstrsDataflowAnalysis(BinaryFunction &BF, const bool UsePrivateAllocators)
+      : DataflowAnalysis<Derived, BitVector, Backward, StatePrinterTy>(
+            BF, UsePrivateAllocators) {}
   InstrsDataflowAnalysis(BinaryFunction &BF,
-                         MCPlusBuilder::AllocatorIdTy AllocId = 0)
+                         MCPlusBuilder::AllocatorIdTy AllocId)
       : DataflowAnalysis<Derived, BitVector, Backward, StatePrinterTy>(
             BF, AllocId) {}
   virtual ~InstrsDataflowAnalysis() {}

bolt/include/bolt/Passes/DataflowInfoManager.h

@@ -44,16 +44,11 @@ class DataflowInfoManager {
   std::unique_ptr<std::unordered_map<const MCInst *, BinaryBasicBlock *>>
       InsnToBB;
 
-  // Id of the allocator to be used for annotations added by any of the managed
-  // analysis
-  MCPlusBuilder::AllocatorIdTy AllocatorId;
-
 public:
   DataflowInfoManager(BinaryFunction &BF, const RegAnalysis *RA,
-                      const FrameAnalysis *FA,
-                      MCPlusBuilder::AllocatorIdTy AllocId = 0)
-      : RA(RA), FA(FA), BC(BF.getBinaryContext()), BF(BF),
-        AllocatorId(AllocId){};
+                      const FrameAnalysis *FA)
+      : RA(RA), FA(FA), BC(BF.getBinaryContext()), BF(BF)
+      {};
 
   /// Helper function to fetch the parent BB associated with a program point
   /// If PP is a BB itself, then return itself (cast to a BinaryBasicBlock)

bolt/include/bolt/Passes/DominatorAnalysis.h

@@ -32,6 +32,9 @@ class DominatorAnalysis
                                 Backward>;
 
 public:
+  DominatorAnalysis(BinaryFunction &BF)
+      : InstrsDataflowAnalysis<DominatorAnalysis<Backward>, Backward>(BF) {
+  }
   DominatorAnalysis(BinaryFunction &BF, MCPlusBuilder::AllocatorIdTy AllocId)
       : InstrsDataflowAnalysis<DominatorAnalysis<Backward>, Backward>(BF,
                                                                       AllocId) {

bolt/include/bolt/Passes/FrameAnalysis.h

@@ -208,7 +208,7 @@ class FrameAnalysis {
   /// Get or create an SPT object and run the analysis
   StackPointerTracking &getSPT(BinaryFunction &BF) {
     if (!SPTMap.count(&BF)) {
-      SPTMap.emplace(&BF, std::make_unique<StackPointerTracking>(BF));
+      SPTMap.emplace(&BF, std::make_unique<StackPointerTracking>(BF, true));
       auto Iter = SPTMap.find(&BF);
       assert(Iter != SPTMap.end() && "item should exist");
       Iter->second->run();

bolt/include/bolt/Passes/LivenessAnalysis.h

@@ -31,8 +31,8 @@ class LivenessAnalysis : public DataflowAnalysis<LivenessAnalysis, BitVector,
 
 public:
   LivenessAnalysis(const RegAnalysis &RA, BinaryFunction &BF,
-                   MCPlusBuilder::AllocatorIdTy AllocId)
-      : Parent(BF, AllocId), RA(RA),
+                   const bool UsePrivateAllocators)
+      : Parent(BF, UsePrivateAllocators), RA(RA),
         NumRegs(BF.getBinaryContext().MRI->getNumRegs()) {}
   virtual ~LivenessAnalysis();
 

bolt/include/bolt/Passes/ReachingDefOrUse.h

@@ -34,9 +34,10 @@ class ReachingDefOrUse
 
 public:
   ReachingDefOrUse(const RegAnalysis &RA, BinaryFunction &BF,
-                   std::optional<MCPhysReg> TrackingReg = std::nullopt,
-                   MCPlusBuilder::AllocatorIdTy AllocId = 0)
-      : InstrsDataflowAnalysis<ReachingDefOrUse<Def>, !Def>(BF, AllocId),
+                   const bool UsePrivateAllocators,
+                   std::optional<MCPhysReg> TrackingReg = std::nullopt)
+      : InstrsDataflowAnalysis<ReachingDefOrUse<Def>, !Def>(
+            BF, UsePrivateAllocators),
         RA(RA), TrackingReg(TrackingReg) {}
   virtual ~ReachingDefOrUse() {}
 

bolt/include/bolt/Passes/ReachingInsns.h

@@ -26,8 +26,9 @@ class ReachingInsns
   friend class DataflowAnalysis<ReachingInsns<Backward>, BitVector, Backward>;
 
 public:
-  ReachingInsns(BinaryFunction &BF, MCPlusBuilder::AllocatorIdTy AllocId = 0)
-      : InstrsDataflowAnalysis<ReachingInsns, Backward>(BF, AllocId) {}
+  ReachingInsns(BinaryFunction &BF, const bool UsePrivateAllocators)
+      : InstrsDataflowAnalysis<ReachingInsns, Backward>(BF,
+                                                        UsePrivateAllocators) {}
   virtual ~ReachingInsns() {}
 
   bool isInLoop(const BinaryBasicBlock &BB) {

bolt/include/bolt/Passes/StackAllocationAnalysis.h

@@ -31,8 +31,9 @@ class StackAllocationAnalysis
 
 public:
   StackAllocationAnalysis(BinaryFunction &BF, StackPointerTracking &SPT,
-                          MCPlusBuilder::AllocatorIdTy AllocId)
-      : InstrsDataflowAnalysis<StackAllocationAnalysis, false>(BF, AllocId),
+                          const bool UsePrivateAllocators)
+      : InstrsDataflowAnalysis<StackAllocationAnalysis, false>(
+            BF, UsePrivateAllocators),
         SPT(SPT) {}
   virtual ~StackAllocationAnalysis() {}
 

bolt/include/bolt/Passes/StackAvailableExpressions.h

@@ -28,7 +28,8 @@ class StackAvailableExpressions
 
 public:
   StackAvailableExpressions(const RegAnalysis &RA, const FrameAnalysis &FA,
-                            BinaryFunction &BF);
+                            BinaryFunction &BF,
+                            const bool UsePrivateAllocators);
   virtual ~StackAvailableExpressions() {}
 
   void run() { InstrsDataflowAnalysis<StackAvailableExpressions>::run(); }