[Mlir-commits] [llvm] [mlir] [ADT] Make DenseMap/DenseSet more resilient agains OOM situations (PR #107251)

Marc Auberer llvmlistbot at llvm.org
Fri Sep 13 02:13:47 PDT 2024


https://github.com/marcauberer updated https://github.com/llvm/llvm-project/pull/107251

>From 7f437b15244eee1278c39952b5bf39c39331cff3 Mon Sep 17 00:00:00 2001
From: Marc Auberer <marc.auberer at sap.com>
Date: Wed, 4 Sep 2024 14:07:01 +0000
Subject: [PATCH] [ADT] Make DenseMap/DenseSet more resilient agains OOM
 situations

---
 llvm/include/llvm/ADT/DenseMap.h      | 423 +++++++++++++++++++-------
 llvm/include/llvm/ADT/DenseSet.h      |  68 ++++-
 llvm/include/llvm/TextAPI/SymbolSet.h |  14 +-
 mlir/include/mlir/Support/LLVM.h      |  10 +-
 4 files changed, 385 insertions(+), 130 deletions(-)

diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
index 68498a3de993a1..3a092431efe89f 100644
--- a/llvm/include/llvm/ADT/DenseMap.h
+++ b/llvm/include/llvm/ADT/DenseMap.h
@@ -16,6 +16,7 @@
 
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/EpochTracker.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/AlignOf.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/MathExtras.h"
@@ -36,8 +37,6 @@ namespace llvm {
 
 namespace detail {
 
-// We extend a pair to allow users to override the bucket type with their own
-// implementation without requiring two members.
 template <typename KeyT, typename ValueT>
 struct DenseMapPair : public std::pair<KeyT, ValueT> {
   using std::pair<KeyT, ValueT>::pair;
@@ -48,16 +47,72 @@ struct DenseMapPair : public std::pair<KeyT, ValueT> {
   const ValueT &getSecond() const { return std::pair<KeyT, ValueT>::second; }
 };
 
+// OOM safe DenseMap bucket
+// flags to indicate if key and value are constructed
+// this is necessary since assigning empty key or tombstone key throws for some
+// key types
+
+template <typename KeyT, typename ValueT, typename Enable = void>
+struct DenseMapPairImpl;
+
+// part I: helpers to decide which specialization of DenseMapPairImpl to use
+template <typename KeyT, typename ValueT>
+using EnableIfTriviallyDestructibleMap =
+    typename std::enable_if_t<std::is_trivially_destructible_v<KeyT> &&
+                              std::is_trivially_destructible_v<ValueT>>;
+template <typename KeyT, typename ValueT>
+using EnableIfNotTriviallyDestructibleMap =
+    typename std::enable_if_t<!std::is_trivially_destructible_v<KeyT> ||
+                              !std::is_trivially_destructible_v<ValueT>>;
+
+// OOM safe DenseMap bucket
+// part II: specialization for trivially destructible types
+template <typename KeyT, typename ValueT>
+struct DenseMapPairImpl<KeyT, ValueT,
+                        EnableIfTriviallyDestructibleMap<KeyT, ValueT>>
+    : public DenseMapPair<KeyT, ValueT> {
+  LLVM_ATTRIBUTE_ALWAYS_INLINE void setKeyConstructed(bool) {}
+  LLVM_ATTRIBUTE_ALWAYS_INLINE bool isKeyConstructed() const { return false; }
+  LLVM_ATTRIBUTE_ALWAYS_INLINE void setValueConstructed(bool) {}
+  LLVM_ATTRIBUTE_ALWAYS_INLINE bool isValueConstructed() const { return false; }
+};
+
+// OOM safe DenseMap bucket
+// part III: specialization for non-trivially destructible types
+template <typename KeyT, typename ValueT>
+struct DenseMapPairImpl<KeyT, ValueT,
+                        EnableIfNotTriviallyDestructibleMap<KeyT, ValueT>>
+    : DenseMapPair<KeyT, ValueT> {
+private:
+  bool m_keyConstructed;
+  bool m_valueConstructed;
+
+public:
+  LLVM_ATTRIBUTE_ALWAYS_INLINE void setKeyConstructed(bool constructed) {
+    m_keyConstructed = constructed;
+  }
+  LLVM_ATTRIBUTE_ALWAYS_INLINE bool isKeyConstructed() const {
+    return m_keyConstructed;
+  }
+  LLVM_ATTRIBUTE_ALWAYS_INLINE void setValueConstructed(bool constructed) {
+    m_valueConstructed = constructed;
+  }
+  LLVM_ATTRIBUTE_ALWAYS_INLINE bool isValueConstructed() const {
+    return m_valueConstructed;
+  }
+};
+
 } // end namespace detail
 
 template <typename KeyT, typename ValueT,
           typename KeyInfoT = DenseMapInfo<KeyT>,
-          typename Bucket = llvm::detail::DenseMapPair<KeyT, ValueT>,
+          typename BucketT = llvm::detail::DenseMapPairImpl<KeyT, ValueT>,
+          typename BucketBaseT = llvm::detail::DenseMapPair<KeyT, ValueT>,
           bool IsConst = false>
 class DenseMapIterator;
 
 template <typename DerivedT, typename KeyT, typename ValueT, typename KeyInfoT,
-          typename BucketT>
+          typename BucketT, typename BucketBaseT>
 class DenseMapBase : public DebugEpochBase {
   template <typename T>
   using const_arg_type_t = typename const_pointer_or_const_ref<T>::type;
@@ -66,11 +121,12 @@ class DenseMapBase : public DebugEpochBase {
   using size_type = unsigned;
   using key_type = KeyT;
   using mapped_type = ValueT;
-  using value_type = BucketT;
+  using value_type = BucketBaseT;
 
-  using iterator = DenseMapIterator<KeyT, ValueT, KeyInfoT, BucketT>;
+  using iterator =
+      DenseMapIterator<KeyT, ValueT, KeyInfoT, BucketT, BucketBaseT>;
   using const_iterator =
-      DenseMapIterator<KeyT, ValueT, KeyInfoT, BucketT, true>;
+      DenseMapIterator<KeyT, ValueT, KeyInfoT, BucketT, BucketBaseT, true>;
 
   inline iterator begin() {
     // When the map is empty, avoid the overhead of advancing/retreating past
@@ -125,19 +181,15 @@ class DenseMapBase : public DebugEpochBase {
       for (BucketT *P = getBuckets(), *E = getBucketsEnd(); P != E; ++P)
         P->getFirst() = EmptyKey;
     } else {
-      const KeyT TombstoneKey = getTombstoneKey();
-      unsigned NumEntries = getNumEntries();
       for (BucketT *P = getBuckets(), *E = getBucketsEnd(); P != E; ++P) {
         if (!KeyInfoT::isEqual(P->getFirst(), EmptyKey)) {
-          if (!KeyInfoT::isEqual(P->getFirst(), TombstoneKey)) {
+          if (P->isValueConstructed()) {
             P->getSecond().~ValueT();
-            --NumEntries;
+            P->setValueConstructed(false);
           }
           P->getFirst() = EmptyKey;
         }
       }
-      assert(NumEntries == 0 && "Node count imbalance!");
-      (void)NumEntries;
     }
     setNumEntries(0);
     setNumTombstones(0);
@@ -340,17 +392,22 @@ class DenseMapBase : public DebugEpochBase {
       return false; // not in map.
 
     TheBucket->getSecond().~ValueT();
+    TheBucket->setValueConstructed(false);
     TheBucket->getFirst() = getTombstoneKey();
     decrementNumEntries();
     incrementNumTombstones();
     return true;
   }
   void erase(iterator I) {
-    BucketT *TheBucket = &*I;
-    TheBucket->getSecond().~ValueT();
-    TheBucket->getFirst() = getTombstoneKey();
-    decrementNumEntries();
-    incrementNumTombstones();
+    BucketT *TheBucket = static_cast<BucketT *>(&*I);
+    // Iterator can point to nullptr in case of memory malfunctions
+    if (TheBucket != nullptr) {
+      TheBucket->getSecond().~ValueT();
+      TheBucket->setValueConstructed(false);
+      TheBucket->getFirst() = getTombstoneKey();
+      decrementNumEntries();
+      incrementNumTombstones();
+    }
   }
 
   LLVM_DEPRECATED("Use [Key] instead", "[Key]")
@@ -406,12 +463,24 @@ class DenseMapBase : public DebugEpochBase {
     if (getNumBuckets() == 0) // Nothing to do.
       return;
 
-    const KeyT EmptyKey = getEmptyKey(), TombstoneKey = getTombstoneKey();
     for (BucketT *P = getBuckets(), *E = getBucketsEnd(); P != E; ++P) {
-      if (!KeyInfoT::isEqual(P->getFirst(), EmptyKey) &&
-          !KeyInfoT::isEqual(P->getFirst(), TombstoneKey))
+      if (P->isValueConstructed()) {
         P->getSecond().~ValueT();
-      P->getFirst().~KeyT();
+        P->setValueConstructed(false);
+      }
+      if (P->isKeyConstructed()) {
+        P->getFirst().~KeyT();
+        P->setKeyConstructed(false);
+      }
+    }
+  }
+
+  void initUnitialized() {
+    BucketT *B = getBuckets();
+    BucketT *E = getBucketsEnd();
+    for (; B != E; ++B) {
+      B->setKeyConstructed(false);
+      B->setValueConstructed(false);
     }
   }
 
@@ -422,8 +491,16 @@ class DenseMapBase : public DebugEpochBase {
     assert((getNumBuckets() & (getNumBuckets() - 1)) == 0 &&
            "# initial buckets must be a power of two!");
     const KeyT EmptyKey = getEmptyKey();
-    for (BucketT *B = getBuckets(), *E = getBucketsEnd(); B != E; ++B)
+#ifndef NDEBUG
+    for (BucketT *B = getBuckets(), *E = getBucketsEnd(); B != E; ++B) {
+      assert(!B->isKeyConstructed());
+      assert(!B->isValueConstructed());
+    }
+#endif
+    for (BucketT *B = getBuckets(), *E = getBucketsEnd(); B != E; ++B) {
       ::new (&B->getFirst()) KeyT(EmptyKey);
+      B->setKeyConstructed(true);
+    }
   }
 
   /// Returns the number of buckets to allocate to ensure that the DenseMap can
@@ -453,18 +530,23 @@ class DenseMapBase : public DebugEpochBase {
         assert(!FoundVal && "Key already in new map?");
         DestBucket->getFirst() = std::move(B->getFirst());
         ::new (&DestBucket->getSecond()) ValueT(std::move(B->getSecond()));
+        DestBucket->setValueConstructed(true);
         incrementNumEntries();
-
-        // Free the value.
+      }
+      if (B->isValueConstructed()) {
         B->getSecond().~ValueT();
+        B->setValueConstructed(false);
+      }
+      if (B->isKeyConstructed()) {
+        B->getFirst().~KeyT();
+        B->setKeyConstructed(false);
       }
-      B->getFirst().~KeyT();
     }
   }
 
   template <typename OtherBaseT>
-  void copyFrom(
-      const DenseMapBase<OtherBaseT, KeyT, ValueT, KeyInfoT, BucketT> &other) {
+  void copyFrom(const DenseMapBase<OtherBaseT, KeyT, ValueT, KeyInfoT, BucketT,
+                                   BucketBaseT> &other) {
     assert(&other != this);
     assert(getNumBuckets() == other.getNumBuckets());
 
@@ -483,9 +565,12 @@ class DenseMapBase : public DebugEpochBase {
       const KeyT TombstoneKey = getTombstoneKey();
       for (size_t I = 0; I < NumBuckets; ++I) {
         ::new (&Buckets[I].getFirst()) KeyT(OtherBuckets[I].getFirst());
+        Buckets[I].setKeyConstructed(true);
         if (!KeyInfoT::isEqual(Buckets[I].getFirst(), EmptyKey) &&
-            !KeyInfoT::isEqual(Buckets[I].getFirst(), TombstoneKey))
+            !KeyInfoT::isEqual(Buckets[I].getFirst(), TombstoneKey)) {
           ::new (&Buckets[I].getSecond()) ValueT(OtherBuckets[I].getSecond());
+          Buckets[I].setValueConstructed(true);
+        }
       }
     }
   }
@@ -575,9 +660,9 @@ class DenseMapBase : public DebugEpochBase {
   BucketT *InsertIntoBucket(BucketT *TheBucket, KeyArg &&Key,
                             ValueArgs &&...Values) {
     TheBucket = InsertIntoBucketImpl(Key, TheBucket);
-
     TheBucket->getFirst() = std::forward<KeyArg>(Key);
     ::new (&TheBucket->getSecond()) ValueT(std::forward<ValueArgs>(Values)...);
+    TheBucket->setValueConstructed(true);
     return TheBucket;
   }
 
@@ -585,9 +670,9 @@ class DenseMapBase : public DebugEpochBase {
   BucketT *InsertIntoBucketWithLookup(BucketT *TheBucket, KeyT &&Key,
                                       ValueT &&Value, LookupKeyT &Lookup) {
     TheBucket = InsertIntoBucketImpl(Lookup, TheBucket);
-
     TheBucket->getFirst() = std::move(Key);
     ::new (&TheBucket->getSecond()) ValueT(std::move(Value));
+    TheBucket->setValueConstructed(true);
     return TheBucket;
   }
 
@@ -712,6 +797,48 @@ class DenseMapBase : public DebugEpochBase {
     }
   }
 
+protected:
+  // helper class to guarantee deallocation of buffer
+  class ReleaseOldBuffer {
+    BucketT *m_buckets;
+    unsigned m_numBuckets;
+
+  public:
+    ReleaseOldBuffer(BucketT *buckets, unsigned numBuckets)
+        : m_buckets(buckets), m_numBuckets(numBuckets) {}
+    ~ReleaseOldBuffer() {
+#ifndef NDEBUG
+      memset((void *)m_buckets, 0x5a, sizeof(BucketT) * m_numBuckets);
+#endif
+      // Free the old table.
+      size_t const alignment{alignof(BucketT)};
+      deallocate_buffer(static_cast<void *>(m_buckets),
+                        sizeof(BucketT) * m_numBuckets, alignment);
+    }
+  };
+
+  // helper class to guarantee destruction of bucket content
+  class ReleaseOldBuckets {
+    BucketT *m_buckets;
+    unsigned m_numBuckets;
+
+  public:
+    ReleaseOldBuckets(BucketT *buckets, unsigned numBuckets)
+        : m_buckets(buckets), m_numBuckets(numBuckets) {}
+    ~ReleaseOldBuckets() {
+      for (BucketT *B = m_buckets, *E = m_buckets + m_numBuckets; B != E; ++B) {
+        if (B->isValueConstructed()) {
+          B->getSecond().~ValueT();
+          B->setValueConstructed(false);
+        }
+        if (B->isKeyConstructed()) {
+          B->getFirst().~KeyT();
+          B->setKeyConstructed(false);
+        }
+      }
+    }
+  };
+
 public:
   /// Return the approximate size (in bytes) of the actual map.
   /// This is just the raw memory used by DenseMap.
@@ -727,10 +854,11 @@ class DenseMapBase : public DebugEpochBase {
 /// Equivalent to N calls to RHS.find and N value comparisons. Amortized
 /// complexity is linear, worst case is O(N^2) (if every hash collides).
 template <typename DerivedT, typename KeyT, typename ValueT, typename KeyInfoT,
-          typename BucketT>
-bool operator==(
-    const DenseMapBase<DerivedT, KeyT, ValueT, KeyInfoT, BucketT> &LHS,
-    const DenseMapBase<DerivedT, KeyT, ValueT, KeyInfoT, BucketT> &RHS) {
+          typename BucketT, typename BucketBaseT>
+bool operator==(const DenseMapBase<DerivedT, KeyT, ValueT, KeyInfoT, BucketT,
+                                   BucketBaseT> &LHS,
+                const DenseMapBase<DerivedT, KeyT, ValueT, KeyInfoT, BucketT,
+                                   BucketBaseT> &RHS) {
   if (LHS.size() != RHS.size())
     return false;
 
@@ -747,23 +875,28 @@ bool operator==(
 ///
 /// Equivalent to !(LHS == RHS). See operator== for performance notes.
 template <typename DerivedT, typename KeyT, typename ValueT, typename KeyInfoT,
-          typename BucketT>
-bool operator!=(
-    const DenseMapBase<DerivedT, KeyT, ValueT, KeyInfoT, BucketT> &LHS,
-    const DenseMapBase<DerivedT, KeyT, ValueT, KeyInfoT, BucketT> &RHS) {
+          typename BucketT, typename BucketBaseT>
+bool operator!=(const DenseMapBase<DerivedT, KeyT, ValueT, KeyInfoT, BucketT,
+                                   BucketBaseT> &LHS,
+                const DenseMapBase<DerivedT, KeyT, ValueT, KeyInfoT, BucketT,
+                                   BucketBaseT> &RHS) {
   return !(LHS == RHS);
 }
 
 template <typename KeyT, typename ValueT,
           typename KeyInfoT = DenseMapInfo<KeyT>,
-          typename BucketT = llvm::detail::DenseMapPair<KeyT, ValueT>>
-class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
-                                     KeyT, ValueT, KeyInfoT, BucketT> {
-  friend class DenseMapBase<DenseMap, KeyT, ValueT, KeyInfoT, BucketT>;
+          typename BucketT = llvm::detail::DenseMapPairImpl<KeyT, ValueT>,
+          typename BucketBaseT = llvm::detail::DenseMapPair<KeyT, ValueT>>
+class DenseMap : public DenseMapBase<
+                     DenseMap<KeyT, ValueT, KeyInfoT, BucketT, BucketBaseT>,
+                     KeyT, ValueT, KeyInfoT, BucketT, BucketBaseT> {
+  friend class DenseMapBase<DenseMap, KeyT, ValueT, KeyInfoT, BucketT,
+                            BucketBaseT>;
 
   // Lift some types from the dependent base class into this class for
   // simplicity of referring to them.
-  using BaseT = DenseMapBase<DenseMap, KeyT, ValueT, KeyInfoT, BucketT>;
+  using BaseT =
+      DenseMapBase<DenseMap, KeyT, ValueT, KeyInfoT, BucketT, BucketBaseT>;
 
   BucketT *Buckets;
   unsigned NumEntries;
@@ -776,28 +909,45 @@ class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
   explicit DenseMap(unsigned InitialReserve = 0) { init(InitialReserve); }
 
   DenseMap(const DenseMap &other) : BaseT() {
+    auto CtorGuard = llvm::make_scope_exit([&] { destroy(); });
     init(0);
     copyFrom(other);
+    CtorGuard.release();
   }
 
   DenseMap(DenseMap &&other) : BaseT() {
+    auto CtorGuard = llvm::make_scope_exit([&] { destroy(); });
     init(0);
     swap(other);
+    CtorGuard.release();
   }
 
   template <typename InputIt> DenseMap(const InputIt &I, const InputIt &E) {
+    auto CtorGuard = llvm::make_scope_exit([&] { destroy(); });
     init(std::distance(I, E));
     this->insert(I, E);
+    CtorGuard.release();
   }
 
   DenseMap(std::initializer_list<typename BaseT::value_type> Vals) {
+    auto CtorGuard = llvm::make_scope_exit([&] { destroy(); });
     init(Vals.size());
     this->insert(Vals.begin(), Vals.end());
+    CtorGuard.release();
   }
 
-  ~DenseMap() {
+  ~DenseMap() { destroy(); }
+
+  void destroy() {
     this->destroyAll();
+    deallocateBuckets();
+  }
+
+  void deallocateBuckets() {
     deallocate_buffer(Buckets, sizeof(BucketT) * NumBuckets, alignof(BucketT));
+    Buckets = nullptr;
+    NumBuckets = 0;
+    NumTombstones = 0;
   }
 
   void swap(DenseMap &RHS) {
@@ -816,16 +966,14 @@ class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
   }
 
   DenseMap &operator=(DenseMap &&other) {
-    this->destroyAll();
-    deallocate_buffer(Buckets, sizeof(BucketT) * NumBuckets, alignof(BucketT));
+    destroy();
     init(0);
     swap(other);
     return *this;
   }
 
   void copyFrom(const DenseMap &other) {
-    this->destroyAll();
-    deallocate_buffer(Buckets, sizeof(BucketT) * NumBuckets, alignof(BucketT));
+    destroy();
     if (allocateBuckets(other.NumBuckets)) {
       this->BaseT::copyFrom(other);
     } else {
@@ -855,16 +1003,16 @@ class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
       this->BaseT::initEmpty();
       return;
     }
-
+    // ensure that old buckets are always released, in normal case as well as in
+    // case of an exception Insert all the old elements.
+    typename BaseT::ReleaseOldBuffer ReleaseOldBuffer(OldBuckets,
+                                                      OldNumBuckets);
+    typename BaseT::ReleaseOldBuckets ReleaseOldBuckets(OldBuckets,
+                                                        OldNumBuckets);
     this->moveFromOldBuckets(OldBuckets, OldBuckets + OldNumBuckets);
-
-    // Free the old table.
-    deallocate_buffer(OldBuckets, sizeof(BucketT) * OldNumBuckets,
-                      alignof(BucketT));
   }
 
   void shrink_and_clear() {
-    unsigned OldNumBuckets = NumBuckets;
     unsigned OldNumEntries = NumEntries;
     this->destroyAll();
 
@@ -877,8 +1025,7 @@ class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
       return;
     }
 
-    deallocate_buffer(Buckets, sizeof(BucketT) * OldNumBuckets,
-                      alignof(BucketT));
+    deallocateBuckets();
     init(NewNumBuckets);
   }
 
@@ -896,30 +1043,36 @@ class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
   unsigned getNumBuckets() const { return NumBuckets; }
 
   bool allocateBuckets(unsigned Num) {
-    NumBuckets = Num;
-    if (NumBuckets == 0) {
+    if (Num == 0) {
       Buckets = nullptr;
+      NumBuckets = 0;
       return false;
     }
 
     Buckets = static_cast<BucketT *>(
-        allocate_buffer(sizeof(BucketT) * NumBuckets, alignof(BucketT)));
+        allocate_buffer(sizeof(BucketT) * Num, alignof(BucketT)));
+    // update NumBuckets only if reallocation is successful
+    NumBuckets = Num;
+    this->BaseT::initUnitialized();
     return true;
   }
 };
 
 template <typename KeyT, typename ValueT, unsigned InlineBuckets = 4,
           typename KeyInfoT = DenseMapInfo<KeyT>,
-          typename BucketT = llvm::detail::DenseMapPair<KeyT, ValueT>>
+          typename BucketT = llvm::detail::DenseMapPairImpl<KeyT, ValueT>,
+          typename BucketBaseT = llvm::detail::DenseMapPair<KeyT, ValueT>>
 class SmallDenseMap
-    : public DenseMapBase<
-          SmallDenseMap<KeyT, ValueT, InlineBuckets, KeyInfoT, BucketT>, KeyT,
-          ValueT, KeyInfoT, BucketT> {
-  friend class DenseMapBase<SmallDenseMap, KeyT, ValueT, KeyInfoT, BucketT>;
+    : public DenseMapBase<SmallDenseMap<KeyT, ValueT, InlineBuckets, KeyInfoT,
+                                        BucketT, BucketBaseT>,
+                          KeyT, ValueT, KeyInfoT, BucketT, BucketBaseT> {
+  friend class DenseMapBase<SmallDenseMap, KeyT, ValueT, KeyInfoT, BucketT,
+                            BucketBaseT>;
 
   // Lift some types from the dependent base class into this class for
   // simplicity of referring to them.
-  using BaseT = DenseMapBase<SmallDenseMap, KeyT, ValueT, KeyInfoT, BucketT>;
+  using BaseT =
+      DenseMapBase<SmallDenseMap, KeyT, ValueT, KeyInfoT, BucketT, BucketBaseT>;
 
   static_assert(isPowerOf2_64(InlineBuckets),
                 "InlineBuckets must be a power of 2.");
@@ -935,35 +1088,46 @@ class SmallDenseMap
 
   /// A "union" of an inline bucket array and the struct representing
   /// a large bucket. This union will be discriminated by the 'Small' bit.
-  AlignedCharArrayUnion<BucketT[InlineBuckets], LargeRep> storage;
+  AlignedCharArrayUnion<BucketT[InlineBuckets]> storage;
+  LargeRep largeRep;
 
 public:
   explicit SmallDenseMap(unsigned NumInitBuckets = 0) {
+    auto CtorGuard = llvm::make_scope_exit([&] { destroy(); });
     if (NumInitBuckets > InlineBuckets)
       NumInitBuckets = llvm::bit_ceil(NumInitBuckets);
-    init(NumInitBuckets);
+    initFromCtor(NumInitBuckets);
+    CtorGuard.release();
   }
 
   SmallDenseMap(const SmallDenseMap &other) : BaseT() {
-    init(0);
+    auto CtorGuard = llvm::make_scope_exit([&] { destroy(); });
+    initFromCtor(0);
     copyFrom(other);
+    CtorGuard.release();
   }
 
   SmallDenseMap(SmallDenseMap &&other) : BaseT() {
-    init(0);
+    auto CtorGuard = llvm::make_scope_exit([&] { destroy(); });
+    initFromCtor(0);
     swap(other);
+    CtorGuard.release();
   }
 
   template <typename InputIt>
   SmallDenseMap(const InputIt &I, const InputIt &E) {
-    init(NextPowerOf2(std::distance(I, E)));
+    auto CtorGuard = llvm::make_scope_exit([&] { destroy(); });
+    initFromCtor(NextPowerOf2(std::distance(I, E)));
     this->insert(I, E);
+    CtorGuard.release();
   }
 
   SmallDenseMap(std::initializer_list<typename BaseT::value_type> Vals)
       : SmallDenseMap(Vals.begin(), Vals.end()) {}
 
-  ~SmallDenseMap() {
+  ~SmallDenseMap() { destroy(); }
+
+  void destroy() {
     this->destroyAll();
     deallocateBuckets();
   }
@@ -997,10 +1161,14 @@ class SmallDenseMap
         std::swap(LHSB->getFirst(), RHSB->getFirst());
         if (hasLHSValue) {
           ::new (&RHSB->getSecond()) ValueT(std::move(LHSB->getSecond()));
+          RHSB->setValueConstructed(true);
           LHSB->getSecond().~ValueT();
+          LHSB->setValueConstructed(false);
         } else if (hasRHSValue) {
           ::new (&LHSB->getSecond()) ValueT(std::move(RHSB->getSecond()));
+          LHSB->setValueConstructed(true);
           RHSB->getSecond().~ValueT();
+          RHSB->setValueConstructed(false);
         }
       }
       return;
@@ -1026,18 +1194,22 @@ class SmallDenseMap
       BucketT *NewB = &LargeSide.getInlineBuckets()[i],
               *OldB = &SmallSide.getInlineBuckets()[i];
       ::new (&NewB->getFirst()) KeyT(std::move(OldB->getFirst()));
+      NewB->setKeyConstructed(true);
       OldB->getFirst().~KeyT();
+      OldB->setKeyConstructed(false);
       if (!KeyInfoT::isEqual(NewB->getFirst(), EmptyKey) &&
           !KeyInfoT::isEqual(NewB->getFirst(), TombstoneKey)) {
         ::new (&NewB->getSecond()) ValueT(std::move(OldB->getSecond()));
+        NewB->setValueConstructed(true);
         OldB->getSecond().~ValueT();
+        OldB->setValueConstructed(false);
       }
     }
 
     // The hard part of moving the small buckets across is done, just move
     // the TmpRep into its new home.
+    ::new (SmallSide.getLargeRep()) LargeRep(std::move(TmpRep));
     SmallSide.Small = false;
-    new (SmallSide.getLargeRep()) LargeRep(std::move(TmpRep));
   }
 
   SmallDenseMap &operator=(const SmallDenseMap &other) {
@@ -1047,33 +1219,34 @@ class SmallDenseMap
   }
 
   SmallDenseMap &operator=(SmallDenseMap &&other) {
-    this->destroyAll();
-    deallocateBuckets();
+    destroy();
     init(0);
     swap(other);
     return *this;
   }
 
   void copyFrom(const SmallDenseMap &other) {
-    this->destroyAll();
-    deallocateBuckets();
-    Small = true;
+    destroy();
     if (other.getNumBuckets() > InlineBuckets) {
-      Small = false;
-      new (getLargeRep()) LargeRep(allocateBuckets(other.getNumBuckets()));
+      allocateBuckets(other.getNumBuckets());
     }
     this->BaseT::copyFrom(other);
   }
 
-  void init(unsigned InitBuckets) {
+  void initFromCtor(unsigned InitBuckets) {
     Small = true;
+    this->BaseT::initUnitialized();
+    init(InitBuckets);
+  }
+
+  void init(unsigned InitBuckets) {
     if (InitBuckets > InlineBuckets) {
-      Small = false;
-      new (getLargeRep()) LargeRep(allocateBuckets(InitBuckets));
+      allocateBuckets(InitBuckets);
     }
     this->BaseT::initEmpty();
   }
 
+public:
   void grow(unsigned AtLeast) {
     if (AtLeast > InlineBuckets)
       AtLeast = std::max<unsigned>(64, NextPowerOf2(AtLeast - 1));
@@ -1088,44 +1261,61 @@ class SmallDenseMap
       // temporary storage. Have the loop move the TmpEnd forward as it goes.
       const KeyT EmptyKey = this->getEmptyKey();
       const KeyT TombstoneKey = this->getTombstoneKey();
+
+      // ensure that old buckets are always released, in normal case as well as
+      // in case of an exception
+      typename BaseT::ReleaseOldBuckets ReleaseOldBuckets(TmpBegin,
+                                                          InlineBuckets);
+
+      for (BucketT *P = TmpBegin, *E = P + InlineBuckets; P != E; ++P) {
+        P->setKeyConstructed(false);
+        P->setValueConstructed(false);
+      }
+
       for (BucketT *P = getBuckets(), *E = P + InlineBuckets; P != E; ++P) {
         if (!KeyInfoT::isEqual(P->getFirst(), EmptyKey) &&
             !KeyInfoT::isEqual(P->getFirst(), TombstoneKey)) {
           assert(size_t(TmpEnd - TmpBegin) < InlineBuckets &&
                  "Too many inline buckets!");
           ::new (&TmpEnd->getFirst()) KeyT(std::move(P->getFirst()));
+          TmpEnd->setKeyConstructed(true);
           ::new (&TmpEnd->getSecond()) ValueT(std::move(P->getSecond()));
+          TmpEnd->setValueConstructed(true);
           ++TmpEnd;
+        }
+        if (P->isValueConstructed()) {
           P->getSecond().~ValueT();
+          P->setValueConstructed(false);
+        }
+        if (P->isKeyConstructed()) {
+          P->getFirst().~KeyT();
+          P->setKeyConstructed(false);
         }
-        P->getFirst().~KeyT();
       }
 
       // AtLeast == InlineBuckets can happen if there are many tombstones,
       // and grow() is used to remove them. Usually we always switch to the
       // large rep here.
       if (AtLeast > InlineBuckets) {
-        Small = false;
-        new (getLargeRep()) LargeRep(allocateBuckets(AtLeast));
+        allocateBuckets(AtLeast);
       }
       this->moveFromOldBuckets(TmpBegin, TmpEnd);
       return;
     }
-
     LargeRep OldRep = std::move(*getLargeRep());
     getLargeRep()->~LargeRep();
-    if (AtLeast <= InlineBuckets) {
-      Small = true;
-    } else {
-      new (getLargeRep()) LargeRep(allocateBuckets(AtLeast));
+    Small = true;
+    // ensure that old buckets are always released, in normal case as well as in
+    // case of an exception
+    typename BaseT::ReleaseOldBuffer ReleaseOldBuffer(OldRep.Buckets,
+                                                      OldRep.NumBuckets);
+    typename BaseT::ReleaseOldBuckets ReleaseOldBuckets(OldRep.Buckets,
+                                                        OldRep.NumBuckets);
+    if (AtLeast > InlineBuckets) {
+      allocateBuckets(AtLeast);
     }
-
     this->moveFromOldBuckets(OldRep.Buckets,
                              OldRep.Buckets + OldRep.NumBuckets);
-
-    // Free the old table.
-    deallocate_buffer(OldRep.Buckets, sizeof(BucketT) * OldRep.NumBuckets,
-                      alignof(BucketT));
   }
 
   void shrink_and_clear() {
@@ -1176,9 +1366,8 @@ class SmallDenseMap
   }
 
   const LargeRep *getLargeRep() const {
-    assert(!Small);
     // Note, same rule about aliasing as with getInlineBuckets.
-    return reinterpret_cast<const LargeRep *>(&storage);
+    return reinterpret_cast<const LargeRep *>(&largeRep);
   }
 
   LargeRep *getLargeRep() {
@@ -1200,6 +1389,7 @@ class SmallDenseMap
   }
 
   void deallocateBuckets() {
+    NumEntries = 0;
     if (Small)
       return;
 
@@ -1207,38 +1397,45 @@ class SmallDenseMap
                       sizeof(BucketT) * getLargeRep()->NumBuckets,
                       alignof(BucketT));
     getLargeRep()->~LargeRep();
+    Small = true;
   }
 
-  LargeRep allocateBuckets(unsigned Num) {
+  void allocateBuckets(unsigned Num) {
     assert(Num > InlineBuckets && "Must allocate more buckets than are inline");
-    LargeRep Rep = {static_cast<BucketT *>(allocate_buffer(
-                        sizeof(BucketT) * Num, alignof(BucketT))),
-                    Num};
-    return Rep;
+    largeRep = {static_cast<BucketT *>(
+                    allocate_buffer(sizeof(BucketT) * Num, alignof(BucketT))),
+                Num};
+    Small = false;
+    this->BaseT::initUnitialized();
   }
 };
 
-template <typename KeyT, typename ValueT, typename KeyInfoT, typename Bucket,
-          bool IsConst>
+template <typename KeyT, typename ValueT, typename KeyInfoT, typename BucketT,
+          typename BucketBaseT, bool IsConst>
 class DenseMapIterator : DebugEpochBase::HandleBase {
-  friend class DenseMapIterator<KeyT, ValueT, KeyInfoT, Bucket, true>;
-  friend class DenseMapIterator<KeyT, ValueT, KeyInfoT, Bucket, false>;
+  friend class DenseMapIterator<KeyT, ValueT, KeyInfoT, BucketT, BucketBaseT,
+                                true>;
+  friend class DenseMapIterator<KeyT, ValueT, KeyInfoT, BucketT, BucketBaseT,
+                                false>;
 
 public:
   using difference_type = ptrdiff_t;
-  using value_type = std::conditional_t<IsConst, const Bucket, Bucket>;
+  using inc_type = std::conditional_t<IsConst, const BucketT, BucketT>;
+  using inc_pointer = inc_type *;
+  using value_type =
+      std::conditional_t<IsConst, const BucketBaseT, BucketBaseT>;
   using pointer = value_type *;
   using reference = value_type &;
   using iterator_category = std::forward_iterator_tag;
 
 private:
-  pointer Ptr = nullptr;
-  pointer End = nullptr;
+  inc_pointer Ptr = nullptr;
+  inc_pointer End = nullptr;
 
 public:
   DenseMapIterator() = default;
 
-  DenseMapIterator(pointer Pos, pointer E, const DebugEpochBase &Epoch,
+  DenseMapIterator(inc_pointer Pos, inc_pointer E, const DebugEpochBase &Epoch,
                    bool NoAdvance = false)
       : DebugEpochBase::HandleBase(&Epoch), Ptr(Pos), End(E) {
     assert(isHandleInSync() && "invalid construction!");
@@ -1257,8 +1454,8 @@ class DenseMapIterator : DebugEpochBase::HandleBase {
   // constructor.
   template <bool IsConstSrc,
             typename = std::enable_if_t<!IsConstSrc && IsConst>>
-  DenseMapIterator(
-      const DenseMapIterator<KeyT, ValueT, KeyInfoT, Bucket, IsConstSrc> &I)
+  DenseMapIterator(const DenseMapIterator<KeyT, ValueT, KeyInfoT, BucketT,
+                                          BucketBaseT, IsConstSrc> &I)
       : DebugEpochBase::HandleBase(I), Ptr(I.Ptr), End(I.End) {}
 
   reference operator*() const {
diff --git a/llvm/include/llvm/ADT/DenseSet.h b/llvm/include/llvm/ADT/DenseSet.h
index a307bd8c9198b9..b9a1b0d5087646 100644
--- a/llvm/include/llvm/ADT/DenseSet.h
+++ b/llvm/include/llvm/ADT/DenseSet.h
@@ -27,11 +27,25 @@ namespace llvm {
 
 namespace detail {
 
+// flag to indicate if key is constructed
+// this is necessary since assigning empty key or tombstone key throws for some
+// key types
+
+// Helpers to decide which specialization of DenseSetPair to use
+template <typename KeyT>
+using EnableIfTriviallyDestructibleSet =
+    typename std::enable_if_t<std::is_trivially_destructible_v<KeyT>>;
+template <typename KeyT>
+using EnableIfNotTriviallyDestructibleSet =
+    typename std::enable_if_t<!std::is_trivially_destructible_v<KeyT>>;
+
 struct DenseSetEmpty {};
 
 // Use the empty base class trick so we can create a DenseMap where the buckets
 // contain only a single item.
+
 template <typename KeyT> class DenseSetPair : public DenseSetEmpty {
+private:
   KeyT key;
 
 public:
@@ -39,6 +53,43 @@ template <typename KeyT> class DenseSetPair : public DenseSetEmpty {
   const KeyT &getFirst() const { return key; }
   DenseSetEmpty &getSecond() { return *this; }
   const DenseSetEmpty &getSecond() const { return *this; }
+
+public:
+  LLVM_ATTRIBUTE_ALWAYS_INLINE void setKeyConstructed(bool) {}
+  LLVM_ATTRIBUTE_ALWAYS_INLINE bool isKeyConstructed() const { return false; }
+  LLVM_ATTRIBUTE_ALWAYS_INLINE void setValueConstructed(bool) {}
+  LLVM_ATTRIBUTE_ALWAYS_INLINE bool isValueConstructed() const { return false; }
+};
+
+template <typename KeyT, typename Enabled = void> class DenseSetPairImpl;
+
+// First specialization for trivially destructible types
+template <typename KeyT>
+class DenseSetPairImpl<KeyT, EnableIfTriviallyDestructibleSet<KeyT>>
+    : public DenseSetPair<KeyT> {
+public:
+  LLVM_ATTRIBUTE_ALWAYS_INLINE void setKeyConstructed(bool) {}
+  LLVM_ATTRIBUTE_ALWAYS_INLINE bool isKeyConstructed() const { return false; }
+  LLVM_ATTRIBUTE_ALWAYS_INLINE void setValueConstructed(bool) {}
+  LLVM_ATTRIBUTE_ALWAYS_INLINE bool isValueConstructed() const { return false; }
+};
+
+// Second specialization for non-trivially destructible types
+template <typename KeyT>
+class DenseSetPairImpl<KeyT, EnableIfNotTriviallyDestructibleSet<KeyT>>
+    : public DenseSetPair<KeyT> {
+public:
+  bool KeyConstructed;
+
+public:
+  LLVM_ATTRIBUTE_ALWAYS_INLINE void setKeyConstructed(bool Constructed) {
+    KeyConstructed = Constructed;
+  }
+  LLVM_ATTRIBUTE_ALWAYS_INLINE bool isKeyConstructed() const {
+    return KeyConstructed;
+  }
+  LLVM_ATTRIBUTE_ALWAYS_INLINE void setValueConstructed(bool) {}
+  LLVM_ATTRIBUTE_ALWAYS_INLINE bool isValueConstructed() const { return false; }
 };
 
 /// Base class for DenseSet and DenseSmallSet.
@@ -52,8 +103,6 @@ template <typename KeyT> class DenseSetPair : public DenseSetEmpty {
 /// DenseMapInfo "concept".
 template <typename ValueT, typename MapTy, typename ValueInfoT>
 class DenseSetImpl {
-  static_assert(sizeof(typename MapTy::value_type) == sizeof(ValueT),
-                "DenseMap buckets unexpectedly large!");
   MapTy TheMap;
 
   template <typename T>
@@ -274,13 +323,14 @@ template <typename ValueT, typename ValueInfoT = DenseMapInfo<ValueT>>
 class DenseSet : public detail::DenseSetImpl<
                      ValueT,
                      DenseMap<ValueT, detail::DenseSetEmpty, ValueInfoT,
+                              detail::DenseSetPairImpl<ValueT>,
                               detail::DenseSetPair<ValueT>>,
                      ValueInfoT> {
-  using BaseT =
-      detail::DenseSetImpl<ValueT,
-                           DenseMap<ValueT, detail::DenseSetEmpty, ValueInfoT,
-                                    detail::DenseSetPair<ValueT>>,
-                           ValueInfoT>;
+  using BaseT = detail::DenseSetImpl<
+      ValueT,
+      DenseMap<ValueT, detail::DenseSetEmpty, ValueInfoT,
+               detail::DenseSetPairImpl<ValueT>, detail::DenseSetPair<ValueT>>,
+      ValueInfoT>;
 
 public:
   using BaseT::BaseT;
@@ -294,11 +344,13 @@ class SmallDenseSet
     : public detail::DenseSetImpl<
           ValueT,
           SmallDenseMap<ValueT, detail::DenseSetEmpty, InlineBuckets,
-                        ValueInfoT, detail::DenseSetPair<ValueT>>,
+                        ValueInfoT, detail::DenseSetPairImpl<ValueT>,
+                        detail::DenseSetPair<ValueT>>,
           ValueInfoT> {
   using BaseT = detail::DenseSetImpl<
       ValueT,
       SmallDenseMap<ValueT, detail::DenseSetEmpty, InlineBuckets, ValueInfoT,
+                    detail::DenseSetPairImpl<ValueT>,
                     detail::DenseSetPair<ValueT>>,
       ValueInfoT>;
 
diff --git a/llvm/include/llvm/TextAPI/SymbolSet.h b/llvm/include/llvm/TextAPI/SymbolSet.h
index 6ccabb90772087..e97b63fc5946f5 100644
--- a/llvm/include/llvm/TextAPI/SymbolSet.h
+++ b/llvm/include/llvm/TextAPI/SymbolSet.h
@@ -48,11 +48,12 @@ template <> struct DenseMapInfo<SymbolsMapKey> {
   }
 };
 
-template <typename DerivedT, typename KeyInfoT, typename BucketT>
+template <typename DerivedT, typename KeyInfoT, typename BucketT,
+          typename BucketBaseT>
 bool operator==(const DenseMapBase<DerivedT, SymbolsMapKey, MachO::Symbol *,
-                                   KeyInfoT, BucketT> &LHS,
+                                   KeyInfoT, BucketT, BucketBaseT> &LHS,
                 const DenseMapBase<DerivedT, SymbolsMapKey, MachO::Symbol *,
-                                   KeyInfoT, BucketT> &RHS) {
+                                   KeyInfoT, BucketT, BucketBaseT> &RHS) {
   if (LHS.size() != RHS.size())
     return false;
   for (const auto &KV : LHS) {
@@ -63,11 +64,12 @@ bool operator==(const DenseMapBase<DerivedT, SymbolsMapKey, MachO::Symbol *,
   return true;
 }
 
-template <typename DerivedT, typename KeyInfoT, typename BucketT>
+template <typename DerivedT, typename KeyInfoT, typename BucketT,
+          typename BucketBaseT>
 bool operator!=(const DenseMapBase<DerivedT, SymbolsMapKey, MachO::Symbol *,
-                                   KeyInfoT, BucketT> &LHS,
+                                   KeyInfoT, BucketT, BucketBaseT> &LHS,
                 const DenseMapBase<DerivedT, SymbolsMapKey, MachO::Symbol *,
-                                   KeyInfoT, BucketT> &RHS) {
+                                   KeyInfoT, BucketT, BucketBaseT> &RHS) {
   return !(LHS == RHS);
 }
 
diff --git a/mlir/include/mlir/Support/LLVM.h b/mlir/include/mlir/Support/LLVM.h
index 020c0fba726c84..8e8eb9c8f99ca1 100644
--- a/mlir/include/mlir/Support/LLVM.h
+++ b/mlir/include/mlir/Support/LLVM.h
@@ -50,8 +50,11 @@ class BitVector;
 namespace detail {
 template <typename KeyT, typename ValueT>
 struct DenseMapPair;
+template <typename KeyT, typename ValueT, typename Enable>
+struct DenseMapPairImpl;
 } // namespace detail
-template <typename KeyT, typename ValueT, typename KeyInfoT, typename BucketT>
+template <typename KeyT, typename ValueT, typename KeyInfoT, typename BucketT,
+          typename BucketBaseT>
 class DenseMap;
 template <typename T, typename Enable>
 struct DenseMapInfo;
@@ -122,8 +125,9 @@ template <typename T, typename Enable = void>
 using DenseMapInfo = llvm::DenseMapInfo<T, Enable>;
 template <typename KeyT, typename ValueT,
           typename KeyInfoT = DenseMapInfo<KeyT>,
-          typename BucketT = llvm::detail::DenseMapPair<KeyT, ValueT>>
-using DenseMap = llvm::DenseMap<KeyT, ValueT, KeyInfoT, BucketT>;
+          typename BucketT = llvm::detail::DenseMapPairImpl<KeyT, ValueT, void>,
+          typename BucketBaseT = llvm::detail::DenseMapPair<KeyT, ValueT>>
+using DenseMap = llvm::DenseMap<KeyT, ValueT, KeyInfoT, BucketT, BucketBaseT>;
 template <typename ValueT, typename ValueInfoT = DenseMapInfo<ValueT>>
 using DenseSet = llvm::DenseSet<ValueT, ValueInfoT>;
 template <typename T, typename Vector = llvm::SmallVector<T, 0>,



More information about the Mlir-commits mailing list