From fe593fe15f780517a703c4c108fc162028f180bb Mon Sep 17 00:00:00 2001
From: Nikita Popov <nikita.ppv@gmail.com>
Date: Wed, 11 Dec 2019 21:17:29 +0100
Subject: [PATCH] [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    | 16 ++++++++--------
 llvm/unittests/ADT/DenseMapTest.cpp | 18 ++++++++++++++++++
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
index 948a6e6bfb38..148d319c8603 100644
--- a/llvm/include/llvm/ADT/DenseMap.h
+++ b/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;
     }
diff --git a/llvm/unittests/ADT/DenseMapTest.cpp b/llvm/unittests/ADT/DenseMapTest.cpp
index 7e85ce7011fd..1ff12a6366cb 100644
--- 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));