[llvm] [ADT] Move shrink_and_clear to DenseMapBase (NFC) (PR #165103)

[Kazu Hirata via llvm-commits]

https://github.com/kazutakahirata updated https://github.com/llvm/llvm-project/pull/165103

>From 3fa678625aec39db32d9aeced56116e1b88a4adf Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Fri, 24 Oct 2025 21:05:53 -0700
Subject: [PATCH 1/2] [ADT] Move shrink_and_clear to DenseMapBase (NFC)

Without this patch, DenseMap and SmallDenseMap have distinct
implementations of shrink_and_clear.  These implementations mix a
common high-level algorithm with class-specific logic.

This patch moves the common algorithm into
DenseMapBase::shrink_and_clear.  A new private helper,
planShrinkAndClear, now handles the class-specific logic for deciding
whether to shrink the buffer.  The base class method now serves as the
single public entry point.
---
 llvm/include/llvm/ADT/DenseMap.h | 76 ++++++++++++++------------------
 1 file changed, 34 insertions(+), 42 deletions(-)

diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
index baa91f3a5f533..b2e8224a83fa8 100644
--- a/llvm/include/llvm/ADT/DenseMap.h
+++ b/llvm/include/llvm/ADT/DenseMap.h
@@ -154,6 +154,17 @@ class DenseMapBase : public DebugEpochBase {
     setNumTombstones(0);
   }
 
+  void shrink_and_clear() {
+    auto [Reallocate, NewNumBuckets] = derived().planShrinkAndClear();
+    destroyAll();
+    if (!Reallocate) {
+      initEmpty();
+      return;
+    }
+    derived().deallocateBuckets();
+    derived().init(NewNumBuckets);
+  }
+
   /// Return true if the specified key is in the map, false otherwise.
   [[nodiscard]] bool contains(const_arg_type_t<KeyT> Val) const {
     return doFind(Val) != nullptr;
@@ -532,8 +543,6 @@ class DenseMapBase : public DebugEpochBase {
 
   void grow(unsigned AtLeast) { derived().grow(AtLeast); }
 
-  void shrink_and_clear() { derived().shrink_and_clear(); }
-
   template <typename LookupKeyT>
   BucketT *findBucketForInsertion(const LookupKeyT &Lookup,
                                   BucketT *TheBucket) {
@@ -773,25 +782,6 @@ class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
     return *this;
   }
 
-  void shrink_and_clear() {
-    unsigned OldNumBuckets = NumBuckets;
-    unsigned OldNumEntries = NumEntries;
-    this->destroyAll();
-
-    // Reduce the number of buckets.
-    unsigned NewNumBuckets = 0;
-    if (OldNumEntries)
-      NewNumBuckets = std::max(64, 1 << (Log2_32_Ceil(OldNumEntries) + 1));
-    if (NewNumBuckets == NumBuckets) {
-      this->BaseT::initEmpty();
-      return;
-    }
-
-    deallocate_buffer(Buckets, sizeof(BucketT) * OldNumBuckets,
-                      alignof(BucketT));
-    init(NewNumBuckets);
-  }
-
 private:
   unsigned getNumEntries() const { return NumEntries; }
 
@@ -861,6 +851,15 @@ class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
     deallocate_buffer(OldBuckets, sizeof(BucketT) * OldNumBuckets,
                       alignof(BucketT));
   }
+
+  std::pair<bool, unsigned> planShrinkAndClear() const {
+    unsigned NewNumBuckets = 0;
+    if (NumEntries)
+      NewNumBuckets = std::max(64u, 1u << (Log2_32_Ceil(NumEntries) + 1));
+    if (NewNumBuckets == NumBuckets)
+      return {false, 0};          // Reuse
+    return {true, NewNumBuckets}; // Reallocate
+  }
 };
 
 template <typename KeyT, typename ValueT, unsigned InlineBuckets = 4,
@@ -1013,27 +1012,6 @@ class SmallDenseMap
     return *this;
   }
 
-  void shrink_and_clear() {
-    unsigned OldSize = this->size();
-    this->destroyAll();
-
-    // Reduce the number of buckets.
-    unsigned NewNumBuckets = 0;
-    if (OldSize) {
-      NewNumBuckets = 1 << (Log2_32_Ceil(OldSize) + 1);
-      if (NewNumBuckets > InlineBuckets && NewNumBuckets < 64u)
-        NewNumBuckets = 64;
-    }
-    if ((Small && NewNumBuckets <= InlineBuckets) ||
-        (!Small && NewNumBuckets == getLargeRep()->NumBuckets)) {
-      this->BaseT::initEmpty();
-      return;
-    }
-
-    deallocateBuckets();
-    init(NewNumBuckets);
-  }
-
 private:
   unsigned getNumEntries() const { return NumEntries; }
 
@@ -1168,6 +1146,20 @@ class SmallDenseMap
     deallocate_buffer(OldRep.Buckets, sizeof(BucketT) * OldRep.NumBuckets,
                       alignof(BucketT));
   }
+
+  std::pair<bool, unsigned> planShrinkAndClear() const {
+    unsigned NewNumBuckets = 0;
+    if (this->size()) {
+      NewNumBuckets = 1u << (Log2_32_Ceil(this->size()) + 1);
+      if (NewNumBuckets > InlineBuckets)
+        NewNumBuckets = std::max(64u, NewNumBuckets);
+    }
+    bool Reuse = Small ? NewNumBuckets <= InlineBuckets
+                       : NewNumBuckets == getLargeRep()->NumBuckets;
+    if (Reuse)
+      return {false, 0};          // Reuse
+    return {true, NewNumBuckets}; // Reallocate
+  }
 };
 
 template <typename KeyT, typename ValueT, typename KeyInfoT, typename Bucket,

>From 18989e108914fb9331618f43e715af0c14fdbedd Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Sat, 25 Oct 2025 08:40:22 -0700
Subject: [PATCH 2/2] Address comments.

---
 llvm/include/llvm/ADT/DenseMap.h | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
index b2e8224a83fa8..608de35e45281 100644
--- a/llvm/include/llvm/ADT/DenseMap.h
+++ b/llvm/include/llvm/ADT/DenseMap.h
@@ -852,13 +852,16 @@ class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
                       alignof(BucketT));
   }
 
+  // Plan how to shrink the bucket table.  Return:
+  // - {false, 0} to reuse the existing bucket table
+  // - {true, N} to reallocate a bucket table with N entries
   std::pair<bool, unsigned> planShrinkAndClear() const {
     unsigned NewNumBuckets = 0;
     if (NumEntries)
       NewNumBuckets = std::max(64u, 1u << (Log2_32_Ceil(NumEntries) + 1));
     if (NewNumBuckets == NumBuckets)
-      return {false, 0};          // Reuse
-    return {true, NewNumBuckets}; // Reallocate
+      return {false, 0};          // Reuse.
+    return {true, NewNumBuckets}; // Reallocate.
   }
 };
 
@@ -1147,9 +1150,12 @@ class SmallDenseMap
                       alignof(BucketT));
   }
 
+  // Plan how to shrink the bucket table.  Return:
+  // - {false, 0} to reuse the existing bucket table
+  // - {true, N} to reallocate a bucket table with N entries
   std::pair<bool, unsigned> planShrinkAndClear() const {
     unsigned NewNumBuckets = 0;
-    if (this->size()) {
+    if (!this->empty()) {
       NewNumBuckets = 1u << (Log2_32_Ceil(this->size()) + 1);
       if (NewNumBuckets > InlineBuckets)
         NewNumBuckets = std::max(64u, NewNumBuckets);
@@ -1157,8 +1163,8 @@ class SmallDenseMap
     bool Reuse = Small ? NewNumBuckets <= InlineBuckets
                        : NewNumBuckets == getLargeRep()->NumBuckets;
     if (Reuse)
-      return {false, 0};          // Reuse
-    return {true, NewNumBuckets}; // Reallocate
+      return {false, 0};          // Reuse.
+    return {true, NewNumBuckets}; // Reallocate.
   }
 };