[ADT] Fix SmallDenseMap assertion with large InlineBuckets

nikic · nikic · commit fe593fe15f78 · 2019-12-11T21:41:14.000+01:00
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
diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
@@ -1006,13 +1006,10 @@ class SmallDenseMap
   }
 
   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 @@ class SmallDenseMap
         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;
     }
diff --git a/llvm/unittests/ADT/DenseMapTest.cpp b/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));
