[ADT] Reduce code duplication in SmallDenseMap (NFC)

[kazutakahirata]

This patch reduces code duplication by having allocateBuckets take a
larger role. Specifically, allocateBuckets now checks to see if we
need to allocate heap memory and initializes Small appropriately.

With this patch, allocateBuckets mirrors deallocateBuckets cleanly.
Both methods handle the Small mode without asserting and are
responsible for constructing and destructing LargeRep.

[llvmbot]

@llvm/pr-subscribers-llvm-adt

Author: Kazu Hirata (kazutakahirata)

Changes

This patch reduces code duplication by having allocateBuckets take a
larger role. Specifically, allocateBuckets now checks to see if we
need to allocate heap memory and initializes Small appropriately.

With this patch, allocateBuckets mirrors deallocateBuckets cleanly.
Both methods handle the Small mode without asserting and are
responsible for constructing and destructing LargeRep.

---

Full diff: https://github.com/llvm/llvm-project/pull/160813.diff

1 Files Affected:

• (modified) llvm/include/llvm/ADT/DenseMap.h (+13-25)

diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
index e13a2cb09a412..5f716751751c4 100644
--- a/llvm/include/llvm/ADT/DenseMap.h
+++ b/llvm/include/llvm/ADT/DenseMap.h
@@ -1009,21 +1009,13 @@ class SmallDenseMap
   void copyFrom(const SmallDenseMap &other) {
     this->destroyAll();
     deallocateBuckets();
-    Small = true;
-    if (other.getNumBuckets() > InlineBuckets) {
-      Small = false;
-      new (getLargeRep()) LargeRep(allocateBuckets(other.getNumBuckets()));
-    }
+    allocateBuckets(other.getNumBuckets());
     this->BaseT::copyFrom(other);
   }
 
   void init(unsigned InitNumEntries) {
     auto InitBuckets = BaseT::getMinBucketToReserveForEntries(InitNumEntries);
-    Small = true;
-    if (InitBuckets > InlineBuckets) {
-      Small = false;
-      new (getLargeRep()) LargeRep(allocateBuckets(InitBuckets));
-    }
+    allocateBuckets(InitBuckets);
     this->BaseT::initEmpty();
   }
 
@@ -1057,21 +1049,14 @@ class SmallDenseMap
       // 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(llvm::make_range(TmpBegin, TmpEnd));
       return;
     }
 
     LargeRep OldRep = std::move(*getLargeRep());
     getLargeRep()->~LargeRep();
-    if (AtLeast <= InlineBuckets) {
-      Small = true;
-    } else {
-      new (getLargeRep()) LargeRep(allocateBuckets(AtLeast));
-    }
+    allocateBuckets(AtLeast);
 
     this->moveFromOldBuckets(OldRep.buckets());
 
@@ -1166,12 +1151,15 @@ class SmallDenseMap
     getLargeRep()->~LargeRep();
   }
 
-  LargeRep 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;
+  void allocateBuckets(unsigned Num) {
+    if (Num <= InlineBuckets) {
+      Small = true;
+    } else {
+      Small = false;
+      BucketT *NewBuckets = static_cast<BucketT *>(
+          allocate_buffer(sizeof(BucketT) * Num, alignof(BucketT)));
+      new (getLargeRep()) LargeRep{NewBuckets, Num};
+    }
   }
 };