[ADT] Fix SmallDenseMap assertion with large InlineBuckets

Fixes issue encountered in D56362, where I tried to use a
SmallSetVector<Instruction*, 128> with an excessively large number
of inline elements. This triggers an "Must allocate more buckets
than are inline" assertion inside allocateBuckets() under certain
usage patterns.

The issue is as follows: The grow() method is used either to grow
the map, or to rehash it and remove tombstones. The latter is done
if the fraction of empty (non-used, non-tombstone) elements is
below 1/8. In this case grow() is invoked with the current number
of buckets.

This is currently incorrectly handled for dense maps using the small
rep. The current implementation will switch them over to the large
rep, which violates the invariant that the large rep is only used
if there are more than InlineBuckets buckets.

This patch fixes the issue by staying in the small rep and only
moving the buckets. An alternative, if we do want to switch to the
large rep in this case, would be to relax the assertion in
allocateBuckets().

Differential Revision: https://reviews.llvm.org/D56455

llvm/include/llvm/ADT/DenseMap.h

@@ -1006,13 +1006,10 @@ public:
   }
 
   void grow(unsigned AtLeast) {
-    if (AtLeast >= InlineBuckets)
+    if (AtLeast > InlineBuckets)
       AtLeast = std::max<unsigned>(64, NextPowerOf2(AtLeast-1));
 
     if (Small) {
-      if (AtLeast < InlineBuckets)
-        return; // Nothing to do.
-
       // First move the inline buckets into a temporary storage.
       AlignedCharArrayUnion<BucketT[InlineBuckets]> TmpStorage;
       BucketT *TmpBegin = reinterpret_cast<BucketT *>(TmpStorage.buffer);
@@ -1035,10 +1032,13 @@ public:
         P->getFirst().~KeyT();
       }
 
-      // Now make this map use the large rep, and move all the entries back
-      // into it.
-      Small = false;
-      new (getLargeRep()) LargeRep(allocateBuckets(AtLeast));
+      // 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));
+      }
       this->moveFromOldBuckets(TmpBegin, TmpEnd);
       return;
     }

llvm/unittests/ADT/DenseMapTest.cpp

@@ -580,6 +580,24 @@ TEST(DenseMapCustomTest, SmallDenseMapGrowTest) {
   EXPECT_TRUE(map.find(32) == map.end());
 }
 
+TEST(DenseMapCustomTest, LargeSmallDenseMapCompaction) {
+  SmallDenseMap<unsigned, unsigned, 128, ContiguousDenseMapInfo> map;
+  // Fill to < 3/4 load.
+  for (unsigned i = 0; i < 95; ++i)
+    map[i] = i;
+  // And erase, leaving behind tombstones.
+  for (unsigned i = 0; i < 95; ++i)
+    map.erase(i);
+  // Fill further, so that less than 1/8 are empty, but still below 3/4 load.
+  for (unsigned i = 95; i < 128; ++i)
+    map[i] = i;
+
+  EXPECT_EQ(33u, map.size());
+  // Similar to the previous test, check for a non-existing element, as an
+  // indirect check that tombstones have been removed.
+  EXPECT_TRUE(map.find(0) == map.end());
+}
+
 TEST(DenseMapCustomTest, TryEmplaceTest) {
   DenseMap<int, std::unique_ptr<int>> Map;
   std::unique_ptr<int> P(new int(2));