shared_ptrify ownershp of PoolEntries in PBQP's CostPool

Leveraging both intrusive shared_ptr-ing (std::enable_shared_from_this)
and shared_ptr<T>-owning-U (to allow external users to hold
std::shared_ptr<CostT> while keeping the underlying PoolEntry alive).

The intrusiveness could be removed if we had a weak_set that implicitly
removed items from the set when their underlying data went away.

This /might/ fix an existing memory leak reported by LeakSanitizer in
r217504.

llvm-svn: 217563

1 files changed, 10 insertions, 46 deletions

diff --git a/llvm/include/llvm/CodeGen/PBQP/CostAllocator.h b/llvm/include/llvm/CodeGen/PBQP/CostAllocator.h
index 892e5b4557b..67ae5ec0a1f 100644
--- a/llvm/include/llvm/CodeGen/PBQP/CostAllocator.h
+++ b/llvm/include/llvm/CodeGen/PBQP/CostAllocator.h
@@ -18,6 +18,7 @@
 #ifndef LLVM_CODEGEN_PBQP_COSTALLOCATOR_H
 #define LLVM_CODEGEN_PBQP_COSTALLOCATOR_H
 
+#include <memory>
 #include <set>
 #include <type_traits>
 
@@ -27,57 +28,20 @@ template <typename CostT,
           typename CostKeyTComparator>
 class CostPool {
 public:
-
-  class PoolEntry {
+  class PoolEntry : public std::enable_shared_from_this<PoolEntry> {
   public:
     template <typename CostKeyT>
     PoolEntry(CostPool &pool, CostKeyT cost)
-      : pool(pool), cost(std::move(cost)), refCount(0) {}
+        : pool(pool), cost(std::move(cost)) {}
     ~PoolEntry() { pool.removeEntry(this); }
-    void incRef() { ++refCount; }
-    bool decRef() { --refCount; return (refCount == 0); }
     CostT& getCost() { return cost; }
     const CostT& getCost() const { return cost; }
   private:
     CostPool &pool;
     CostT cost;
-    std::size_t refCount;
   };
 
-  class PoolRef {
-  public:
-    PoolRef(PoolEntry *entry) : entry(entry) {
-      this->entry->incRef();
-    }
-    PoolRef(const PoolRef &r) {
-      entry = r.entry;
-      entry->incRef();
-    }
-    PoolRef& operator=(const PoolRef &r) {
-      assert(entry != nullptr && "entry should not be null.");
-      PoolEntry *temp = r.entry;
-      temp->incRef();
-      entry->decRef();
-      entry = temp;
-      return *this;
-    }
-
-    ~PoolRef() {
-      if (entry->decRef())
-        delete entry;
-    }
-    void reset(PoolEntry *entry) {
-      entry->incRef();
-      this->entry->decRef();
-      this->entry = entry;
-    }
-    CostT& operator*() { return entry->getCost(); }
-    const CostT& operator*() const { return entry->getCost(); }
-    CostT* operator->() { return &entry->getCost(); }
-    const CostT* operator->() const { return &entry->getCost(); }
-  private:
-    PoolEntry *entry;
-  };
+  typedef std::shared_ptr<CostT> PoolRef;
 
 private:
   class EntryComparator {
@@ -104,19 +68,19 @@ private:
   void removeEntry(PoolEntry *p) { entrySet.erase(p); }
 
 public:
-
   template <typename CostKeyT>
-  PoolRef getCost(CostKeyT costKey) {
+  std::shared_ptr<CostT> getCost(CostKeyT costKey) {
     typename EntrySet::iterator itr =
       std::lower_bound(entrySet.begin(), entrySet.end(), costKey,
                        EntryComparator());
 
     if (itr != entrySet.end() && costKey == (*itr)->getCost())
-      return PoolRef(*itr);
+      return std::shared_ptr<CostT>((*itr)->shared_from_this(),
+                                    &(*itr)->getCost());
 
-    PoolEntry *p = new PoolEntry(*this, std::move(costKey));
-    entrySet.insert(itr, p);
-    return PoolRef(p);
+    auto p = std::make_shared<PoolEntry>(*this, std::move(costKey));
+    entrySet.insert(itr, p.get());
+    return std::shared_ptr<CostT>(std::move(p), &p->getCost());
   }
 };