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

Marc Auberer llvmlistbot at llvm.org
Tue Sep 10 04:56:28 PDT 2024


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

>From 435e491feb0d584df9c9d5e812c8f9351c8fa9cd 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 1/4] [ADT] Make DenseMap/DenseSet more resilient agains OOM
 situations

---
 llvm/cmake/modules/LLVMConfig.cmake.in        |   3 +
 llvm/include/llvm/ADT/DenseMap.h              | 696 +++++++++++-------
 llvm/include/llvm/ADT/DenseSet.h              | 121 ++-
 llvm/include/llvm/TextAPI/SymbolSet.h         |  14 +-
 .../unittests/ADT/DenseMapMalfunctionTest.cpp | 145 ++++
 .../unittests/ADT/DenseSetMalfunctionTest.cpp |  18 +
 llvm/unittests/ADT/MalfunctionTestRunner.h    |  89 +++
 llvm/unittests/ADT/TestIntAlloc.h             | 113 +++
 8 files changed, 911 insertions(+), 288 deletions(-)
 create mode 100644 llvm/unittests/ADT/DenseMapMalfunctionTest.cpp
 create mode 100644 llvm/unittests/ADT/DenseSetMalfunctionTest.cpp
 create mode 100644 llvm/unittests/ADT/MalfunctionTestRunner.h
 create mode 100644 llvm/unittests/ADT/TestIntAlloc.h

diff --git a/llvm/cmake/modules/LLVMConfig.cmake.in b/llvm/cmake/modules/LLVMConfig.cmake.in
index 7e1501a89354c8..dcaed41b458c09 100644
--- a/llvm/cmake/modules/LLVMConfig.cmake.in
+++ b/llvm/cmake/modules/LLVMConfig.cmake.in
@@ -47,6 +47,9 @@ set(LLVM_ENABLE_EXPENSIVE_CHECKS @LLVM_ENABLE_EXPENSIVE_CHECKS@)
 set(LLVM_ENABLE_ASSERTIONS @LLVM_ENABLE_ASSERTIONS@)
 
 set(LLVM_ENABLE_EH @LLVM_ENABLE_EH@)
+if(LLVM_ENABLE_EH)
+  add_definitions(-DLLVM_EH_ENABLED)
+endif()
 
 set(LLVM_ENABLE_FFI @LLVM_ENABLE_FFI@)
 if(LLVM_ENABLE_FFI)
diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
index 00290c9dd0a585..ce84c8ff51f34f 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
@@ -109,34 +165,62 @@ class DenseMapBase : public DebugEpochBase {
 
   void clear() {
     incrementEpoch();
-    if (getNumEntries() == 0 && getNumTombstones() == 0) return;
+    if (getNumEntries() == 0 && getNumTombstones() == 0)
+      return;
 
     // If the capacity of the array is huge, and the # elements used is small,
     // shrink the array.
     if (getNumEntries() * 4 < getNumBuckets() && getNumBuckets() > 64) {
+#ifdef LLVM_EH_ENABLED
+      // catch bad_alloc during assign of empty key in shrink_and_clear
+      try {
+        shrink_and_clear();
+      } catch (std::bad_alloc &) {
+        // catch bad_alloc, but simply ignore, we must not throw bad_alloc in
+        // destruction
+      }
+#else
       shrink_and_clear();
+#endif
       return;
     }
 
     const KeyT EmptyKey = getEmptyKey();
     if (std::is_trivially_destructible<ValueT>::value) {
       // Use a simpler loop when values don't need destruction.
-      for (BucketT *P = getBuckets(), *E = getBucketsEnd(); P != E; ++P)
+      for (BucketT *P = getBuckets(), *E = getBucketsEnd(); P != E; ++P) {
+#ifdef LLVM_EH_ENABLED
+        // catch bad_alloc during assign of empty key in shrink_and_clear
+        try {
+          P->getFirst() = EmptyKey;
+        } catch (std::bad_alloc &) {
+          // catch bad_alloc, but simply ignore, we must not throw bad_alloc in
+          // destruction
+        }
+#else
         P->getFirst() = EmptyKey;
+#endif
+      }
     } 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);
+          }
+#ifdef LLVM_EH_ENABLED
+          // catch bad_alloc during assign of empty key in shrink_and_clear
+          try {
+            P->getFirst() = EmptyKey;
+          } catch (std::bad_alloc &) {
+            // catch bad_alloc, but simply ignore, we must not throw bad_alloc
+            // in destruction
           }
+#else
           P->getFirst() = EmptyKey;
+#endif
         }
       }
-      assert(NumEntries == 0 && "Node count imbalance!");
-      (void)NumEntries;
     }
     setNumEntries(0);
     setNumTombstones(0);
@@ -172,15 +256,14 @@ class DenseMapBase : public DebugEpochBase {
   /// The DenseMapInfo is responsible for supplying methods
   /// getHashValue(LookupKeyT) and isEqual(LookupKeyT, KeyT) for each key
   /// type used.
-  template<class LookupKeyT>
-  iterator find_as(const LookupKeyT &Val) {
+  template <class LookupKeyT> iterator find_as(const LookupKeyT &Val) {
     if (BucketT *Bucket = doFind(Val))
       return makeIterator(
           Bucket, shouldReverseIterate<KeyT>() ? getBuckets() : getBucketsEnd(),
           *this, true);
     return end();
   }
-  template<class LookupKeyT>
+  template <class LookupKeyT>
   const_iterator find_as(const LookupKeyT &Val) const {
     if (const BucketT *Bucket = doFind(Val))
       return makeConstIterator(
@@ -223,7 +306,7 @@ class DenseMapBase : public DebugEpochBase {
   // The value is constructed in-place if the key is not in the map, otherwise
   // it is not moved.
   template <typename... Ts>
-  std::pair<iterator, bool> try_emplace(KeyT &&Key, Ts &&... Args) {
+  std::pair<iterator, bool> try_emplace(KeyT &&Key, Ts &&...Args) {
     BucketT *TheBucket;
     if (LookupBucketFor(Key, TheBucket))
       return std::make_pair(makeIterator(TheBucket,
@@ -248,7 +331,7 @@ class DenseMapBase : public DebugEpochBase {
   // The value is constructed in-place if the key is not in the map, otherwise
   // it is not moved.
   template <typename... Ts>
-  std::pair<iterator, bool> try_emplace(const KeyT &Key, Ts &&... Args) {
+  std::pair<iterator, bool> try_emplace(const KeyT &Key, Ts &&...Args) {
     BucketT *TheBucket;
     if (LookupBucketFor(Key, TheBucket))
       return std::make_pair(makeIterator(TheBucket,
@@ -297,8 +380,7 @@ class DenseMapBase : public DebugEpochBase {
   }
 
   /// insert - Range insertion of pairs.
-  template<typename InputIt>
-  void insert(InputIt I, InputIt E) {
+  template <typename InputIt> void insert(InputIt I, InputIt E) {
     for (; I != E; ++I)
       insert(*I);
   }
@@ -341,17 +423,50 @@ 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();
+    }
+  }
+
+  template <typename LookupKeyT> BucketT *doFind(const LookupKeyT &Val) {
+    BucketT *BucketsPtr = getBuckets();
+    const unsigned NumBuckets = getNumBuckets();
+    if (NumBuckets == 0)
+      return nullptr;
+
+    const KeyT EmptyKey = getEmptyKey();
+    unsigned BucketNo = getHashValue(Val) & (NumBuckets - 1);
+    unsigned ProbeAmt = 1;
+    while (true) {
+      BucketT *Bucket = BucketsPtr + BucketNo;
+      if (LLVM_LIKELY(KeyInfoT::isEqual(Val, Bucket->getFirst())))
+        return Bucket;
+      if (LLVM_LIKELY(KeyInfoT::isEqual(Bucket->getFirst(), EmptyKey)))
+        return nullptr;
+
+      // Otherwise, it's a hash collision or a tombstone, continue quadratic
+      // probing.
+      BucketNo += ProbeAmt++;
+      BucketNo &= NumBuckets - 1;
+    }
+  }
+
+  template <typename LookupKeyT>
+  const BucketT *doFind(const LookupKeyT &Val) const {
+    return const_cast<DenseMapBase *>(this)->doFind(Val); // NOLINT
   }
 
   LLVM_DEPRECATED("Use [Key] instead", "[Key]")
@@ -359,7 +474,6 @@ class DenseMapBase : public DebugEpochBase {
     BucketT *TheBucket;
     if (LookupBucketFor(Key, TheBucket))
       return *TheBucket;
-
     return *InsertIntoBucket(TheBucket, Key);
   }
 
@@ -376,7 +490,6 @@ class DenseMapBase : public DebugEpochBase {
     BucketT *TheBucket;
     if (LookupBucketFor(Key, TheBucket))
       return *TheBucket;
-
     return *InsertIntoBucket(TheBucket, std::move(Key));
   }
 
@@ -404,15 +517,25 @@ class DenseMapBase : public DebugEpochBase {
   DenseMapBase() = default;
 
   void destroyAll() {
-    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([[maybe_unused]] unsigned NumBuckets) {
+    BucketT *B = getBuckets();
+    BucketT *E = getBucketsEnd();
+    //assert(E - B == NumBuckets);
+    for (; B != E; ++B) {
+      B->setKeyConstructed(false);
+      B->setValueConstructed(false);
     }
   }
 
@@ -420,11 +543,19 @@ class DenseMapBase : public DebugEpochBase {
     setNumEntries(0);
     setNumTombstones(0);
 
-    assert((getNumBuckets() & (getNumBuckets()-1)) == 0 &&
+    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,19 +584,24 @@ class DenseMapBase : public DebugEpochBase {
         (void)FoundVal; // silence warning.
         assert(!FoundVal && "Key already in new map?");
         DestBucket->getFirst() = std::move(B->getFirst());
-        ::new (&DestBucket->getSecond()) ValueT(std::move(B->getSecond()));
+        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());
 
@@ -480,10 +616,13 @@ class DenseMapBase : public DebugEpochBase {
       for (size_t i = 0; i < getNumBuckets(); ++i) {
         ::new (&getBuckets()[i].getFirst())
             KeyT(other.getBuckets()[i].getFirst());
+        getBuckets()[i].setKeyConstructed(true);
         if (!KeyInfoT::isEqual(getBuckets()[i].getFirst(), getEmptyKey()) &&
-            !KeyInfoT::isEqual(getBuckets()[i].getFirst(), getTombstoneKey()))
+            !KeyInfoT::isEqual(getBuckets()[i].getFirst(), getTombstoneKey())) {
           ::new (&getBuckets()[i].getSecond())
               ValueT(other.getBuckets()[i].getSecond());
+          getBuckets()[i].setValueConstructed(true);
+        }
       }
   }
 
@@ -491,7 +630,7 @@ class DenseMapBase : public DebugEpochBase {
     return KeyInfoT::getHashValue(Val);
   }
 
-  template<typename LookupKeyT>
+  template <typename LookupKeyT>
   static unsigned getHashValue(const LookupKeyT &Val) {
     return KeyInfoT::getHashValue(Val);
   }
@@ -502,14 +641,11 @@ class DenseMapBase : public DebugEpochBase {
     return KeyInfoT::getEmptyKey();
   }
 
-  static const KeyT getTombstoneKey() {
-    return KeyInfoT::getTombstoneKey();
-  }
+  static const KeyT getTombstoneKey() { return KeyInfoT::getTombstoneKey(); }
 
 private:
-  iterator makeIterator(BucketT *P, BucketT *E,
-                        DebugEpochBase &Epoch,
-                        bool NoAdvance=false) {
+  iterator makeIterator(BucketT *P, BucketT *E, DebugEpochBase &Epoch,
+                        bool NoAdvance = false) {
     if (shouldReverseIterate<KeyT>()) {
       BucketT *B = P == getBucketsEnd() ? getBuckets() : P + 1;
       return iterator(B, E, Epoch, NoAdvance);
@@ -519,7 +655,7 @@ class DenseMapBase : public DebugEpochBase {
 
   const_iterator makeConstIterator(const BucketT *P, const BucketT *E,
                                    const DebugEpochBase &Epoch,
-                                   const bool NoAdvance=false) const {
+                                   const bool NoAdvance = false) const {
     if (shouldReverseIterate<KeyT>()) {
       const BucketT *B = P == getBucketsEnd() ? getBuckets() : P + 1;
       return const_iterator(B, E, Epoch, NoAdvance);
@@ -535,13 +671,9 @@ class DenseMapBase : public DebugEpochBase {
     static_cast<DerivedT *>(this)->setNumEntries(Num);
   }
 
-  void incrementNumEntries() {
-    setNumEntries(getNumEntries() + 1);
-  }
+  void incrementNumEntries() { setNumEntries(getNumEntries() + 1); }
 
-  void decrementNumEntries() {
-    setNumEntries(getNumEntries() - 1);
-  }
+  void decrementNumEntries() { setNumEntries(getNumEntries() - 1); }
 
   unsigned getNumTombstones() const {
     return static_cast<const DerivedT *>(this)->getNumTombstones();
@@ -551,49 +683,37 @@ class DenseMapBase : public DebugEpochBase {
     static_cast<DerivedT *>(this)->setNumTombstones(Num);
   }
 
-  void incrementNumTombstones() {
-    setNumTombstones(getNumTombstones() + 1);
-  }
+  void incrementNumTombstones() { setNumTombstones(getNumTombstones() + 1); }
 
-  void decrementNumTombstones() {
-    setNumTombstones(getNumTombstones() - 1);
-  }
+  void decrementNumTombstones() { setNumTombstones(getNumTombstones() - 1); }
 
   const BucketT *getBuckets() const {
     return static_cast<const DerivedT *>(this)->getBuckets();
   }
 
-  BucketT *getBuckets() {
-    return static_cast<DerivedT *>(this)->getBuckets();
-  }
+  BucketT *getBuckets() { return static_cast<DerivedT *>(this)->getBuckets(); }
 
   unsigned getNumBuckets() const {
     return static_cast<const DerivedT *>(this)->getNumBuckets();
   }
 
-  BucketT *getBucketsEnd() {
-    return getBuckets() + getNumBuckets();
-  }
+  BucketT *getBucketsEnd() { return getBuckets() + getNumBuckets(); }
 
   const BucketT *getBucketsEnd() const {
     return getBuckets() + getNumBuckets();
   }
 
-  void grow(unsigned AtLeast) {
-    static_cast<DerivedT *>(this)->grow(AtLeast);
-  }
+  void grow(unsigned AtLeast) { static_cast<DerivedT *>(this)->grow(AtLeast); }
 
-  void shrink_and_clear() {
-    static_cast<DerivedT *>(this)->shrink_and_clear();
-  }
+  void shrink_and_clear() { static_cast<DerivedT *>(this)->shrink_and_clear(); }
 
   template <typename KeyArg, typename... ValueArgs>
   BucketT *InsertIntoBucket(BucketT *TheBucket, KeyArg &&Key,
-                            ValueArgs &&... Values) {
+                            ValueArgs &&...Values) {
     TheBucket = InsertIntoBucketImpl(Key, Key, TheBucket);
-
     TheBucket->getFirst() = std::forward<KeyArg>(Key);
     ::new (&TheBucket->getSecond()) ValueT(std::forward<ValueArgs>(Values)...);
+    TheBucket->setValueConstructed(true);
     return TheBucket;
   }
 
@@ -601,14 +721,14 @@ class DenseMapBase : public DebugEpochBase {
   BucketT *InsertIntoBucketWithLookup(BucketT *TheBucket, KeyT &&Key,
                                       ValueT &&Value, LookupKeyT &Lookup) {
     TheBucket = InsertIntoBucketImpl(Key, Lookup, TheBucket);
-
     TheBucket->getFirst() = std::move(Key);
     ::new (&TheBucket->getSecond()) ValueT(std::move(Value));
+    TheBucket->setValueConstructed(true);
     return TheBucket;
   }
 
   template <typename LookupKeyT>
-  BucketT *InsertIntoBucketImpl(const KeyT &Key, const LookupKeyT &Lookup,
+  BucketT *InsertIntoBucketImpl(const KeyT &, const LookupKeyT &Lookup,
                                 BucketT *TheBucket) {
     incrementEpoch();
 
@@ -627,8 +747,9 @@ class DenseMapBase : public DebugEpochBase {
       this->grow(NumBuckets * 2);
       LookupBucketFor(Lookup, TheBucket);
       NumBuckets = getNumBuckets();
-    } else if (LLVM_UNLIKELY(NumBuckets-(NewNumEntries+getNumTombstones()) <=
-                             NumBuckets/8)) {
+    } else if (LLVM_UNLIKELY(NumBuckets -
+                                 (NewNumEntries + getNumTombstones()) <=
+                             NumBuckets / 8)) {
       this->grow(NumBuckets);
       LookupBucketFor(Lookup, TheBucket);
     }
@@ -646,34 +767,6 @@ class DenseMapBase : public DebugEpochBase {
     return TheBucket;
   }
 
-  template <typename LookupKeyT> BucketT *doFind(const LookupKeyT &Val) {
-    BucketT *BucketsPtr = getBuckets();
-    const unsigned NumBuckets = getNumBuckets();
-    if (NumBuckets == 0)
-      return nullptr;
-
-    const KeyT EmptyKey = getEmptyKey();
-    unsigned BucketNo = getHashValue(Val) & (NumBuckets - 1);
-    unsigned ProbeAmt = 1;
-    while (true) {
-      BucketT *Bucket = BucketsPtr + BucketNo;
-      if (LLVM_LIKELY(KeyInfoT::isEqual(Val, Bucket->getFirst())))
-        return Bucket;
-      if (LLVM_LIKELY(KeyInfoT::isEqual(Bucket->getFirst(), EmptyKey)))
-        return nullptr;
-
-      // Otherwise, it's a hash collision or a tombstone, continue quadratic
-      // probing.
-      BucketNo += ProbeAmt++;
-      BucketNo &= NumBuckets - 1;
-    }
-  }
-
-  template <typename LookupKeyT>
-  const BucketT *doFind(const LookupKeyT &Val) const {
-    return const_cast<DenseMapBase *>(this)->doFind(Val); // NOLINT
-  }
-
   /// LookupBucketFor - Lookup the appropriate bucket for Val, returning it in
   /// FoundBucket.  If the bucket contains the key and a value, this returns
   /// true, otherwise it returns a bucket with an empty marker or tombstone and
@@ -696,7 +789,7 @@ class DenseMapBase : public DebugEpochBase {
            !KeyInfoT::isEqual(Val, TombstoneKey) &&
            "Empty/Tombstone value shouldn't be inserted into map!");
 
-    unsigned BucketNo = getHashValue(Val) & (NumBuckets-1);
+    unsigned BucketNo = getHashValue(Val) & (NumBuckets - 1);
     unsigned ProbeAmt = 1;
     while (true) {
       BucketT *ThisBucket = BucketsPtr + BucketNo;
@@ -719,23 +812,63 @@ class DenseMapBase : public DebugEpochBase {
       // prefer to return it than something that would require more probing.
       if (KeyInfoT::isEqual(ThisBucket->getFirst(), TombstoneKey) &&
           !FoundTombstone)
-        FoundTombstone = ThisBucket;  // Remember the first tombstone found.
+        FoundTombstone = ThisBucket; // Remember the first tombstone found.
 
       // Otherwise, it's a hash collision or a tombstone, continue quadratic
       // probing.
       BucketNo += ProbeAmt++;
-      BucketNo &= (NumBuckets-1);
+      BucketNo &= (NumBuckets - 1);
     }
   }
 
+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.
   /// If entries are pointers to objects, the size of the referenced objects
   /// are not included.
-  size_t getMemorySize() const {
-    return getNumBuckets() * sizeof(BucketT);
-  }
+  size_t getMemorySize() const { return getNumBuckets() * sizeof(BucketT); }
 };
 
 /// Equality comparison for DenseMap.
@@ -745,10 +878,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;
 
@@ -765,23 +899,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;
@@ -794,32 +933,48 @@ 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) {
+  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) {
+  void swap(DenseMap &RHS) {
     this->incrementEpoch();
     RHS.incrementEpoch();
     std::swap(Buckets, RHS.Buckets);
@@ -828,23 +983,21 @@ class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
     std::swap(NumBuckets, RHS.NumBuckets);
   }
 
-  DenseMap& operator=(const DenseMap& other) {
+  DenseMap &operator=(const DenseMap &other) {
     if (&other != this)
       copyFrom(other);
     return *this;
   }
 
-  DenseMap& operator=(DenseMap &&other) {
-    this->destroyAll();
-    deallocate_buffer(Buckets, sizeof(BucketT) * NumBuckets, alignof(BucketT));
+  DenseMap &operator=(DenseMap &&other) {
+    destroy();
     init(0);
     swap(other);
     return *this;
   }
 
-  void copyFrom(const DenseMap& other) {
-    this->destroyAll();
-    deallocate_buffer(Buckets, sizeof(BucketT) * NumBuckets, alignof(BucketT));
+  void copyFrom(const DenseMap &other) {
+    destroy();
     if (allocateBuckets(other.NumBuckets)) {
       this->BaseT::copyFrom(other);
     } else {
@@ -867,22 +1020,24 @@ class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
     unsigned OldNumBuckets = NumBuckets;
     BucketT *OldBuckets = Buckets;
 
-    allocateBuckets(std::max<unsigned>(64, static_cast<unsigned>(NextPowerOf2(AtLeast-1))));
+    allocateBuckets(std::max<unsigned>(
+        64, static_cast<unsigned>(NextPowerOf2(AtLeast - 1))));
     assert(Buckets);
     if (!OldBuckets) {
       this->BaseT::initEmpty();
       return;
     }
-
-    this->moveFromOldBuckets(OldBuckets, OldBuckets+OldNumBuckets);
-
-    // Free the old table.
-    deallocate_buffer(OldBuckets, sizeof(BucketT) * OldNumBuckets,
-                      alignof(BucketT));
+    // 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);
   }
 
   void shrink_and_clear() {
-    unsigned OldNumBuckets = NumBuckets;
+    // unsigned OldNumBuckets = NumBuckets;
     unsigned OldNumEntries = NumEntries;
     this->destroyAll();
 
@@ -895,61 +1050,54 @@ class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
       return;
     }
 
-    deallocate_buffer(Buckets, sizeof(BucketT) * OldNumBuckets,
-                      alignof(BucketT));
+    deallocateBuckets();
     init(NewNumBuckets);
   }
 
 private:
-  unsigned getNumEntries() const {
-    return NumEntries;
-  }
+  unsigned getNumEntries() const { return NumEntries; }
 
-  void setNumEntries(unsigned Num) {
-    NumEntries = Num;
-  }
+  void setNumEntries(unsigned Num) { NumEntries = Num; }
 
-  unsigned getNumTombstones() const {
-    return NumTombstones;
-  }
+  unsigned getNumTombstones() const { return NumTombstones; }
 
-  void setNumTombstones(unsigned Num) {
-    NumTombstones = Num;
-  }
+  void setNumTombstones(unsigned Num) { NumTombstones = Num; }
 
-  BucketT *getBuckets() const {
-    return Buckets;
-  }
+  BucketT *getBuckets() const { return Buckets; }
 
-  unsigned getNumBuckets() const {
-    return NumBuckets;
-  }
+  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(Num);
     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.");
@@ -965,40 +1113,51 @@ 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>
+  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();
   }
 
-  void swap(SmallDenseMap& RHS) {
+  void swap(SmallDenseMap &RHS) {
     unsigned TmpNumEntries = RHS.NumEntries;
     RHS.NumEntries = NumEntries;
     NumEntries = TmpNumEntries;
@@ -1027,10 +1186,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;
@@ -1056,57 +1219,66 @@ 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.
-    SmallSide.Small = false;
     new (SmallSide.getLargeRep()) LargeRep(std::move(TmpRep));
+    SmallSide.Small = false;
   }
 
-  SmallDenseMap& operator=(const SmallDenseMap& other) {
+  SmallDenseMap &operator=(const SmallDenseMap &other) {
     if (&other != this)
       copyFrom(other);
     return *this;
   }
 
-  SmallDenseMap& operator=(SmallDenseMap &&other) {
-    this->destroyAll();
-    deallocateBuckets();
+  SmallDenseMap &operator=(SmallDenseMap &&other) {
+    destroy();
     init(0);
     swap(other);
     return *this;
   }
 
-  void copyFrom(const SmallDenseMap& other) {
-    this->destroyAll();
-    deallocateBuckets();
-    Small = true;
+  void copyFrom(const SmallDenseMap &other) {
+    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(
+        std::max<unsigned>(InitBuckets, InlineBuckets));
+    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));
+      AtLeast = std::max<unsigned>(64, NextPowerOf2(AtLeast - 1));
+
+    const KeyT EmptyKey = this->getEmptyKey();
+    const KeyT TombstoneKey = this->getTombstoneKey();
 
     if (Small) {
       // First move the inline buckets into a temporary storage.
@@ -1116,45 +1288,61 @@ class SmallDenseMap
 
       // Loop over the buckets, moving non-empty, non-tombstones into the
       // 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()));
-          ::new (&TmpEnd->getSecond()) ValueT(std::move(P->getSecond()));
+          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));
+    this->moveFromOldBuckets(OldRep.Buckets,
+                             OldRep.Buckets + OldRep.NumBuckets);
   }
 
   void shrink_and_clear() {
@@ -1179,9 +1367,7 @@ class SmallDenseMap
   }
 
 private:
-  unsigned getNumEntries() const {
-    return NumEntries;
-  }
+  unsigned getNumEntries() const { return NumEntries; }
 
   void setNumEntries(unsigned Num) {
     // NumEntries is hardcoded to be 31 bits wide.
@@ -1189,13 +1375,9 @@ class SmallDenseMap
     NumEntries = Num;
   }
 
-  unsigned getNumTombstones() const {
-    return NumTombstones;
-  }
+  unsigned getNumTombstones() const { return NumTombstones; }
 
-  void setNumTombstones(unsigned Num) {
-    NumTombstones = Num;
-  }
+  void setNumTombstones(unsigned Num) { NumTombstones = Num; }
 
   const BucketT *getInlineBuckets() const {
     assert(Small);
@@ -1207,18 +1389,17 @@ class SmallDenseMap
 
   BucketT *getInlineBuckets() {
     return const_cast<BucketT *>(
-      const_cast<const SmallDenseMap *>(this)->getInlineBuckets());
+        const_cast<const SmallDenseMap *>(this)->getInlineBuckets());
   }
 
   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() {
     return const_cast<LargeRep *>(
-      const_cast<const SmallDenseMap *>(this)->getLargeRep());
+        const_cast<const SmallDenseMap *>(this)->getLargeRep());
   }
 
   const BucketT *getBuckets() const {
@@ -1227,7 +1408,7 @@ class SmallDenseMap
 
   BucketT *getBuckets() {
     return const_cast<BucketT *>(
-      const_cast<const SmallDenseMap *>(this)->getBuckets());
+        const_cast<const SmallDenseMap *>(this)->getBuckets());
   }
 
   unsigned getNumBuckets() const {
@@ -1235,6 +1416,7 @@ class SmallDenseMap
   }
 
   void deallocateBuckets() {
+    NumEntries = 0;
     if (Small)
       return;
 
@@ -1242,43 +1424,51 @@ 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(Num);
   }
 };
 
-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!");
 
-    if (NoAdvance) return;
+    if (NoAdvance)
+      return;
     if (shouldReverseIterate<KeyT>()) {
       RetreatPastEmptyBuckets();
       return;
@@ -1291,8 +1481,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 {
@@ -1324,7 +1514,7 @@ class DenseMapIterator : DebugEpochBase::HandleBase {
     return !(LHS == RHS);
   }
 
-  inline DenseMapIterator& operator++() {  // Preincrement
+  inline DenseMapIterator &operator++() { // Preincrement
     assert(isHandleInSync() && "invalid iterator access!");
     assert(Ptr != End && "incrementing end() iterator");
     if (shouldReverseIterate<KeyT>()) {
@@ -1336,9 +1526,11 @@ class DenseMapIterator : DebugEpochBase::HandleBase {
     AdvancePastEmptyBuckets();
     return *this;
   }
-  DenseMapIterator operator++(int) {  // Postincrement
+  DenseMapIterator operator++(int) { // Postincrement
     assert(isHandleInSync() && "invalid iterator access!");
-    DenseMapIterator tmp = *this; ++*this; return tmp;
+    DenseMapIterator tmp = *this;
+    ++*this;
+    return tmp;
   }
 
 private:
diff --git a/llvm/include/llvm/ADT/DenseSet.h b/llvm/include/llvm/ADT/DenseSet.h
index b89c88626e43bb..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>
@@ -89,18 +138,12 @@ class DenseSetImpl {
   /// before resizing again.
   void reserve(size_t Size) { TheMap.reserve(Size); }
 
-  void clear() {
-    TheMap.clear();
-  }
+  void clear() { TheMap.clear(); }
 
   /// Return 1 if the specified key is in the set, 0 otherwise.
-  size_type count(const_arg_type_t<ValueT> V) const {
-    return TheMap.count(V);
-  }
+  size_type count(const_arg_type_t<ValueT> V) const { return TheMap.count(V); }
 
-  bool erase(const ValueT &V) {
-    return TheMap.erase(V);
-  }
+  bool erase(const ValueT &V) { return TheMap.erase(V); }
 
   void swap(DenseSetImpl &RHS) { TheMap.swap(RHS.TheMap); }
 
@@ -128,8 +171,15 @@ class DenseSetImpl {
     ValueT *operator->() { return &I->getFirst(); }
     const ValueT *operator->() const { return &I->getFirst(); }
 
-    Iterator& operator++() { ++I; return *this; }
-    Iterator operator++(int) { auto T = *this; ++I; return T; }
+    Iterator &operator++() {
+      ++I;
+      return *this;
+    }
+    Iterator operator++(int) {
+      auto T = *this;
+      ++I;
+      return T;
+    }
     friend bool operator==(const Iterator &X, const Iterator &Y) {
       return X.I == Y.I;
     }
@@ -157,8 +207,15 @@ class DenseSetImpl {
     const ValueT &operator*() const { return I->getFirst(); }
     const ValueT *operator->() const { return &I->getFirst(); }
 
-    ConstIterator& operator++() { ++I; return *this; }
-    ConstIterator operator++(int) { auto T = *this; ++I; return T; }
+    ConstIterator &operator++() {
+      ++I;
+      return *this;
+    }
+    ConstIterator operator++(int) {
+      auto T = *this;
+      ++I;
+      return T;
+    }
     friend bool operator==(const ConstIterator &X, const ConstIterator &Y) {
       return X.I == Y.I;
     }
@@ -191,8 +248,7 @@ class DenseSetImpl {
   /// The DenseMapInfo is responsible for supplying methods
   /// getHashValue(LookupKeyT) and isEqual(LookupKeyT, KeyT) for each key type
   /// used.
-  template <class LookupKeyT>
-  iterator find_as(const LookupKeyT &Val) {
+  template <class LookupKeyT> iterator find_as(const LookupKeyT &Val) {
     return Iterator(TheMap.find_as(Val));
   }
   template <class LookupKeyT>
@@ -226,8 +282,7 @@ class DenseSetImpl {
   }
 
   // Range insertion of values.
-  template<typename InputIt>
-  void insert(InputIt I, InputIt E) {
+  template <typename InputIt> void insert(InputIt I, InputIt E) {
     for (; I != E; ++I)
       insert(*I);
   }
@@ -266,14 +321,16 @@ bool operator!=(const DenseSetImpl<ValueT, MapTy, ValueInfoT> &LHS,
 /// Implements a dense probed hash-table based set.
 template <typename ValueT, typename ValueInfoT = DenseMapInfo<ValueT>>
 class DenseSet : public detail::DenseSetImpl<
-                     ValueT, DenseMap<ValueT, detail::DenseSetEmpty, ValueInfoT,
-                                      detail::DenseSetPair<ValueT>>,
+                     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;
@@ -285,12 +342,16 @@ template <typename ValueT, unsigned InlineBuckets = 4,
           typename ValueInfoT = DenseMapInfo<ValueT>>
 class SmallDenseSet
     : public detail::DenseSetImpl<
-          ValueT, SmallDenseMap<ValueT, detail::DenseSetEmpty, InlineBuckets,
-                                ValueInfoT, detail::DenseSetPair<ValueT>>,
+          ValueT,
+          SmallDenseMap<ValueT, detail::DenseSetEmpty, InlineBuckets,
+                        ValueInfoT, detail::DenseSetPairImpl<ValueT>,
+                        detail::DenseSetPair<ValueT>>,
           ValueInfoT> {
   using BaseT = detail::DenseSetImpl<
-      ValueT, SmallDenseMap<ValueT, detail::DenseSetEmpty, InlineBuckets,
-                            ValueInfoT, detail::DenseSetPair<ValueT>>,
+      ValueT,
+      SmallDenseMap<ValueT, detail::DenseSetEmpty, InlineBuckets, ValueInfoT,
+                    detail::DenseSetPairImpl<ValueT>,
+                    detail::DenseSetPair<ValueT>>,
       ValueInfoT>;
 
 public:
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/llvm/unittests/ADT/DenseMapMalfunctionTest.cpp b/llvm/unittests/ADT/DenseMapMalfunctionTest.cpp
new file mode 100644
index 00000000000000..852d95cf272094
--- /dev/null
+++ b/llvm/unittests/ADT/DenseMapMalfunctionTest.cpp
@@ -0,0 +1,145 @@
+//===- llvm/unittest/ADT/DenseMapMalfunctionTest.cpp - DenseMap malfunction unit
+// tests --*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestIntAlloc.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseMapInfo.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace llvm {
+template <> struct DenseMapInfo<IntAlloc> {
+  static inline IntAlloc getEmptyKey() { return IntAlloc(0xFFFF); }
+  static inline IntAlloc getTombstoneKey() { return IntAlloc(0xFFFF - 1); }
+  static unsigned getHashValue(const IntAlloc &Val) {
+    // magic value copied from llvm
+    return Val.getValue() * 37U;
+  }
+
+  static bool isEqual(const IntAlloc &Lhs, const IntAlloc &Rhs) {
+    return Lhs.getValue() == Rhs.getValue();
+  }
+};
+
+template <> struct DenseMapInfo<IntAllocSpecial> {
+  static inline IntAllocSpecial getEmptyKey() {
+    return IntAllocSpecial(0xFFFF);
+  }
+  static inline IntAllocSpecial getTombstoneKey() {
+    return IntAllocSpecial(0xFFFF - 1);
+  }
+  static unsigned getHashValue(const IntAlloc &Val) {
+    // magic value copied from llvm
+    return Val.getValue() * 37U;
+  }
+
+  static bool isEqual(const IntAllocSpecial &Lhs, const IntAllocSpecial &Rhs) {
+    return Lhs.getValue() == Rhs.getValue();
+  }
+};
+
+} // namespace llvm
+
+namespace {
+
+enum class EnumClass { Val };
+
+// Test class
+template <typename T> class DenseMapMalfunctionTest : public ::testing::Test {
+protected:
+  T Map;
+};
+
+// Register these types for testing.
+// clang-format off
+typedef ::testing::Types<DenseMap<uint32_t, uint32_t>,
+                         DenseMap<uint32_t *, uint32_t *>,
+                         DenseMap<EnumClass, uint32_t>,
+                         SmallDenseMap<uint32_t, uint32_t>,
+                         SmallDenseMap<uint32_t *, uint32_t *>,
+                         SmallDenseMap<EnumClass, uint32_t>
+                         > DenseMapMalfunctionTestTypes;
+// clang-format on
+
+TYPED_TEST_SUITE(DenseMapMalfunctionTest, DenseMapMalfunctionTestTypes, );
+
+TYPED_TEST(DenseMapMalfunctionTest, SingleOperationTest1) {
+  // insert
+  this->Map.insert({1, 2});
+  EXPECT_EQ(1, this->Map.size());
+  EXPECT_EQ(2, this->Map.lookup(1));
+
+  // update (noop)
+  this->Map.insert({1, 3});
+  EXPECT_EQ(1, this->Map.size());
+  EXPECT_EQ(2, this->Map.lookup(1));
+
+  // erase
+  this->Map.erase({1});
+  EXPECT_EQ(0, this->Map.size());
+  EXPECT_FALSE(this->Map.contains(1));
+
+  // clear
+  this->Map.clear();
+  EXPECT_EQ(0, this->Map.size());
+  this->Map.insert({1, 3});
+  EXPECT_EQ(1, this->Map.size());
+  EXPECT_EQ(3, this->Map.lookup(1));
+}
+
+TYPED_TEST(DenseMapMalfunctionTest, SingleOperationTest2) {
+  using MapType = typename TestFixture::T;
+
+  // grow
+  this->Map.grow(1024);
+
+  // insert
+  this->Map.insert({1, 2});
+  EXPECT_EQ(1, this->Map.size());
+  EXPECT_EQ(2, this->Map.lookup(1));
+
+  // copy
+  MapType CopiedDenseMap(this->Map);
+  EXPECT_EQ(1, CopiedDenseMap.size());
+  EXPECT_EQ(2, CopiedDenseMap.lookup(1));
+
+  // move
+  MapType MovedDenseMap(std::move(CopiedDenseMap));
+  EXPECT_EQ(1, MovedDenseMap.size());
+  EXPECT_EQ(2, MovedDenseMap.lookup(1));
+
+  // shrink
+  // add some elements before to exceed inplace buffer of SmallDenseMap
+  this->Map.insert({2, 2});
+  this->Map.insert({3, 2});
+  this->Map.insert({4, 2});
+  this->Map.insert({5, 2});
+  this->Map.shrink_and_clear();
+  EXPECT_EQ(0, this->Map.size());
+}
+
+TYPED_TEST(DenseMapMalfunctionTest, MassOperation) {
+  uint32_t Count = 1024;
+  for (uint32_t I = 0; I != Count; I++) {
+    this->Map.insert({I, 2 * I});
+    if (I % 2 == 0) {
+      this->Map.erase(I);
+    }
+  }
+  for (uint32_t I = 0; I != Count; I++) {
+    if (I % 2 == 0) {
+      EXPECT_FALSE(this->Map.contains(I));
+    } else {
+      EXPECT_EQ(2 * I, this->Map.at(I));
+    }
+  }
+}
+
+} // namespace
diff --git a/llvm/unittests/ADT/DenseSetMalfunctionTest.cpp b/llvm/unittests/ADT/DenseSetMalfunctionTest.cpp
new file mode 100644
index 00000000000000..29250e0d448943
--- /dev/null
+++ b/llvm/unittests/ADT/DenseSetMalfunctionTest.cpp
@@ -0,0 +1,18 @@
+//===- llvm/unittest/ADT/DenseSetMalfunctionTest.cpp - DenseSet malfunction unit tests --*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/DenseSet.h"
+#include "TestIntAlloc.h"
+
+using namespace llvm;
+
+namespace {
+
+
+
+} // namespace
\ No newline at end of file
diff --git a/llvm/unittests/ADT/MalfunctionTestRunner.h b/llvm/unittests/ADT/MalfunctionTestRunner.h
new file mode 100644
index 00000000000000..ea8b9d300c06d5
--- /dev/null
+++ b/llvm/unittests/ADT/MalfunctionTestRunner.h
@@ -0,0 +1,89 @@
+//===- llvm/unittest/ADT/MalfunctionTestRunner.h - Test runner for malfunction
+// testing --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Test runner for OOM scenario tests
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_UNITTESTS_ADT_MALFUNCTIONTESTRUNNER_H
+#define LLVM_UNITTESTS_ADT_MALFUNCTIONTESTRUNNER_H
+
+#ifdef LLVM_EH_ENABLED
+
+#include <iomanip>
+#include <ostream>
+
+namespace llvm {
+
+// Interface for a scenario to be tested for malfunctions
+class IMalfunctionTestScenario {
+public:
+  virtual ~IMalfunctionTestScenario() = default;
+  // is called before execute, but without malfunctions, used to prepare stuff
+  virtual void prepare() {}
+  // execute the scenario
+  virtual void execute() = 0;
+};
+
+// Malfunction test runner
+class MalfunctionTestRunner {
+private:
+  IMalfunctionTestScenario *m_scenario;
+  std::ostream &m_log;
+  static constexpr uint32_t MaxSuccessfulRuns = 3;
+
+public:
+  MalfunctionTestRunner(IMalfunctionTestScenario *Scenario, std::ostream &Log)
+      : m_scenario(Scenario), m_log(Log) {}
+
+  void run() {
+    m_scenario->prepare();
+
+    uint32_t SuccessfulRuns = 0;
+    uint32_t Counter = 0u;
+    while (true) {
+      m_log << std::setw(3) << Counter;
+      m_log.flush(); // need to know the current counter in case of a crash
+
+      bool FinishedWithSuccess = false;
+      try {
+        m_scenario->execute();
+        FinishedWithSuccess = true;
+        m_log << ":ok, ";
+      } catch (const std::bad_alloc &) {
+        m_log << ":ba, ";
+      }
+
+      // Insert new line each 10th iteration
+      if (Counter % 10 == 0) {
+        m_log << std::endl;
+      }
+
+      if (FinishedWithSuccess) {
+        SuccessfulRuns++;
+      } else {
+        // ignore successful runs followed by failures
+        SuccessfulRuns = 0;
+      }
+
+      // Move to the next allocation
+      Counter++;
+
+      if (SuccessfulRuns >= MaxSuccessfulRuns)
+        break;
+    }
+
+    m_log << std::endl << std::endl;
+  }
+};
+
+} // namespace llvm
+
+#endif // LLVM_EH_ENABLED
+#endif // LLVM_UNITTESTS_ADT_MALFUNCTIONTESTRUNNER_H
diff --git a/llvm/unittests/ADT/TestIntAlloc.h b/llvm/unittests/ADT/TestIntAlloc.h
new file mode 100644
index 00000000000000..bdf57174ece63d
--- /dev/null
+++ b/llvm/unittests/ADT/TestIntAlloc.h
@@ -0,0 +1,113 @@
+//===- llvm/unittest/ADT/TestIntAlloc.h - Allocating integer for testing --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Wrapper around an integer, that allocates on assignment.
+// This class is useful when OOM-testing data containers like DenseMap.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_UNITTESTS_ADT_TEST_INT_ALLOC_H
+#define LLVM_UNITTESTS_ADT_TEST_INT_ALLOC_H
+
+#include <cassert>
+#include <cstdint>
+
+namespace llvm {
+
+class IntAlloc {
+  friend class IntAllocSpecial;
+
+private:
+  uint32_t m_value = 0;
+  uint8_t *m_ptr = nullptr;
+  bool m_allocSpecialKeys = false;
+
+public:
+  const uint32_t EMPTY = 0xFFFF;
+  const uint32_t TOMBSTONE = 0xFFFE;
+
+  IntAlloc() = default;
+
+  // NOLINTNEXTLINE(google-explicit-constructor)
+  IntAlloc(uint32_t Value) { alloc(Value); }
+
+  IntAlloc(uint32_t Value, bool AllocSpecial)
+      : m_allocSpecialKeys(AllocSpecial) {
+    alloc(Value);
+  }
+
+  virtual ~IntAlloc() { dealloc(); }
+
+  IntAlloc(const IntAlloc &Other)
+      : IntAlloc(Other.m_value, Other.m_allocSpecialKeys) {}
+
+  IntAlloc(IntAlloc &&Other) noexcept {
+    m_value = Other.m_value;
+    Other.m_value = EMPTY;
+    m_ptr = Other.m_ptr;
+    Other.m_ptr = nullptr;
+    m_allocSpecialKeys = Other.m_allocSpecialKeys;
+  }
+
+  void alloc(uint32_t Value) {
+    // limit the number of actual allocations in mass test to reduce test
+    // runtime
+    bool AllocSpecial =
+        m_allocSpecialKeys && (Value == TOMBSTONE || Value == EMPTY);
+    if (Value <= 4 || Value % 8 == 0 || AllocSpecial) {
+      m_ptr = new uint8_t(0);
+    }
+    m_value = Value;
+  }
+
+  void dealloc() {
+    delete m_ptr;
+    m_ptr = nullptr;
+    m_value = EMPTY;
+  }
+
+  IntAlloc &
+  operator=(const IntAlloc &other) { // NOLINT // self assignment does not occur
+    assert(this != &other);
+    assert(m_allocSpecialKeys == other.m_allocSpecialKeys);
+    dealloc();
+    alloc(other.m_value);
+    return *this;
+  }
+
+  IntAlloc &operator=(
+      IntAlloc &&other) noexcept { // NOLINT // self assignment does not occur
+    assert(this != &other);
+    assert(m_allocSpecialKeys == other.m_allocSpecialKeys);
+    dealloc();
+    m_value = other.m_value;
+    other.m_value = 0xFFFF;
+    m_ptr = other.m_ptr;
+    other.m_ptr = nullptr;
+    return *this;
+  }
+
+  uint32_t getValue() const { return m_value; }
+
+  bool operator==(const IntAlloc &Other) const {
+    return m_value == Other.m_value;
+  }
+
+  // NOLINTNEXTLINE(google-explicit-constructor)
+  operator uint32_t() const { return m_value; }
+};
+
+class IntAllocSpecial : public IntAlloc {
+public:
+  IntAllocSpecial() { m_allocSpecialKeys = true; }
+  explicit IntAllocSpecial(uint32_t Value) : IntAlloc(Value, true) {}
+};
+
+} // end namespace llvm
+
+#endif

>From 80ea658079535586c4125b46ec6fc8343e303216 Mon Sep 17 00:00:00 2001
From: Marc Auberer <marc.auberer at sap.com>
Date: Wed, 4 Sep 2024 14:59:14 +0000
Subject: [PATCH 2/4] Cleanup assertion

---
 llvm/include/llvm/ADT/DenseMap.h | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
index ce84c8ff51f34f..36dca3dd090d7c 100644
--- a/llvm/include/llvm/ADT/DenseMap.h
+++ b/llvm/include/llvm/ADT/DenseMap.h
@@ -529,10 +529,9 @@ class DenseMapBase : public DebugEpochBase {
     }
   }
 
-  void initUnitialized([[maybe_unused]] unsigned NumBuckets) {
+  void initUnitialized() {
     BucketT *B = getBuckets();
     BucketT *E = getBucketsEnd();
-    //assert(E - B == NumBuckets);
     for (; B != E; ++B) {
       B->setKeyConstructed(false);
       B->setValueConstructed(false);
@@ -1078,7 +1077,7 @@ class DenseMap : public DenseMapBase<
         allocate_buffer(sizeof(BucketT) * Num, alignof(BucketT)));
     // update NumBuckets only if reallocation is successful
     NumBuckets = Num;
-    this->BaseT::initUnitialized(Num);
+    this->BaseT::initUnitialized();
     return true;
   }
 };
@@ -1260,8 +1259,7 @@ class SmallDenseMap
 
   void initFromCtor(unsigned InitBuckets) {
     Small = true;
-    this->BaseT::initUnitialized(
-        std::max<unsigned>(InitBuckets, InlineBuckets));
+    this->BaseT::initUnitialized();
     init(InitBuckets);
   }
 
@@ -1433,7 +1431,7 @@ class SmallDenseMap
                     allocate_buffer(sizeof(BucketT) * Num, alignof(BucketT))),
                 Num};
     Small = false;
-    this->BaseT::initUnitialized(Num);
+    this->BaseT::initUnitialized();
   }
 };
 

>From 2b2da4867614faf5af690da144d943bfea251d7f Mon Sep 17 00:00:00 2001
From: Marc Auberer <marc.auberer at sap.com>
Date: Thu, 5 Sep 2024 15:49:35 +0000
Subject: [PATCH 3/4] Fix build error

---
 llvm/include/llvm/ADT/DenseMap.h | 38 ++++----------------------------
 mlir/include/mlir/Support/LLVM.h | 10 ++++++---
 2 files changed, 11 insertions(+), 37 deletions(-)

diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
index 36dca3dd090d7c..51b21941ce5945 100644
--- a/llvm/include/llvm/ADT/DenseMap.h
+++ b/llvm/include/llvm/ADT/DenseMap.h
@@ -171,17 +171,7 @@ class DenseMapBase : public DebugEpochBase {
     // If the capacity of the array is huge, and the # elements used is small,
     // shrink the array.
     if (getNumEntries() * 4 < getNumBuckets() && getNumBuckets() > 64) {
-#ifdef LLVM_EH_ENABLED
-      // catch bad_alloc during assign of empty key in shrink_and_clear
-      try {
-        shrink_and_clear();
-      } catch (std::bad_alloc &) {
-        // catch bad_alloc, but simply ignore, we must not throw bad_alloc in
-        // destruction
-      }
-#else
       shrink_and_clear();
-#endif
       return;
     }
 
@@ -189,17 +179,7 @@ class DenseMapBase : public DebugEpochBase {
     if (std::is_trivially_destructible<ValueT>::value) {
       // Use a simpler loop when values don't need destruction.
       for (BucketT *P = getBuckets(), *E = getBucketsEnd(); P != E; ++P) {
-#ifdef LLVM_EH_ENABLED
-        // catch bad_alloc during assign of empty key in shrink_and_clear
-        try {
-          P->getFirst() = EmptyKey;
-        } catch (std::bad_alloc &) {
-          // catch bad_alloc, but simply ignore, we must not throw bad_alloc in
-          // destruction
-        }
-#else
         P->getFirst() = EmptyKey;
-#endif
       }
     } else {
       for (BucketT *P = getBuckets(), *E = getBucketsEnd(); P != E; ++P) {
@@ -208,17 +188,7 @@ class DenseMapBase : public DebugEpochBase {
             P->getSecond().~ValueT();
             P->setValueConstructed(false);
           }
-#ifdef LLVM_EH_ENABLED
-          // catch bad_alloc during assign of empty key in shrink_and_clear
-          try {
-            P->getFirst() = EmptyKey;
-          } catch (std::bad_alloc &) {
-            // catch bad_alloc, but simply ignore, we must not throw bad_alloc
-            // in destruction
-          }
-#else
           P->getFirst() = EmptyKey;
-#endif
         }
       }
     }
@@ -583,7 +553,7 @@ class DenseMapBase : public DebugEpochBase {
         (void)FoundVal; // silence warning.
         assert(!FoundVal && "Key already in new map?");
         DestBucket->getFirst() = std::move(B->getFirst());
-        new (&DestBucket->getSecond()) ValueT(std::move(B->getSecond()));
+        ::new (&DestBucket->getSecond()) ValueT(std::move(B->getSecond()));
         DestBucket->setValueConstructed(true);
         incrementNumEntries();
       }
@@ -1232,7 +1202,7 @@ class SmallDenseMap
 
     // 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));
+    ::new (SmallSide.getLargeRep()) LargeRep(std::move(TmpRep));
     SmallSide.Small = false;
   }
 
@@ -1302,9 +1272,9 @@ class SmallDenseMap
             !KeyInfoT::isEqual(P->getFirst(), TombstoneKey)) {
           assert(size_t(TmpEnd - TmpBegin) < InlineBuckets &&
                  "Too many inline buckets!");
-          new (&TmpEnd->getFirst()) KeyT(std::move(P->getFirst()));
+          ::new (&TmpEnd->getFirst()) KeyT(std::move(P->getFirst()));
           TmpEnd->setKeyConstructed(true);
-          new (&TmpEnd->getSecond()) ValueT(std::move(P->getSecond()));
+          ::new (&TmpEnd->getSecond()) ValueT(std::move(P->getSecond()));
           TmpEnd->setValueConstructed(true);
           ++TmpEnd;
         }
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>,

>From 6fe1e429d4d9ae8311ff83bfdfa92697be98b29d Mon Sep 17 00:00:00 2001
From: Marc Auberer <marc.auberer at sap.com>
Date: Tue, 10 Sep 2024 11:56:13 +0000
Subject: [PATCH 4/4] Remove LLVM_EH_ENABLED definition again

---
 llvm/cmake/modules/LLVMConfig.cmake.in | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/llvm/cmake/modules/LLVMConfig.cmake.in b/llvm/cmake/modules/LLVMConfig.cmake.in
index dcaed41b458c09..7e1501a89354c8 100644
--- a/llvm/cmake/modules/LLVMConfig.cmake.in
+++ b/llvm/cmake/modules/LLVMConfig.cmake.in
@@ -47,9 +47,6 @@ set(LLVM_ENABLE_EXPENSIVE_CHECKS @LLVM_ENABLE_EXPENSIVE_CHECKS@)
 set(LLVM_ENABLE_ASSERTIONS @LLVM_ENABLE_ASSERTIONS@)
 
 set(LLVM_ENABLE_EH @LLVM_ENABLE_EH@)
-if(LLVM_ENABLE_EH)
-  add_definitions(-DLLVM_EH_ENABLED)
-endif()
 
 set(LLVM_ENABLE_FFI @LLVM_ENABLE_FFI@)
 if(LLVM_ENABLE_FFI)



More information about the Mlir-commits mailing list