From 4bccb581bfe326e6b3aea4a2340c15cdc13e2008 Mon Sep 17 00:00:00 2001
From: Justin Lebar <jlebar@google.com>
Date: Mon, 17 Oct 2016 22:24:28 +0000
Subject: [PATCH] [ADT] Add SmallDenseSet.

Summary: This matches SmallDenseMap.

Reviewers: timshen

Subscribers: llvm-commits

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

llvm-svn: 284432
---
 llvm/include/llvm/ADT/DenseSet.h    | 63 ++++++++++++++++++++++++++++++-------
 llvm/unittests/ADT/DenseSetTest.cpp |  6 +++-
 2 files changed, 56 insertions(+), 13 deletions(-)

diff --git a/llvm/include/llvm/ADT/DenseSet.h b/llvm/include/llvm/ADT/DenseSet.h
index fe44b43..21f7dfa 100644
--- a/llvm/include/llvm/ADT/DenseSet.h
+++ b/llvm/include/llvm/ADT/DenseSet.h
@@ -7,7 +7,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file defines the DenseSet class.
+// This file defines the DenseSet and SmallDenseSet classes.
 //
 //===----------------------------------------------------------------------===//
 
@@ -32,13 +32,18 @@ public:
   DenseSetEmpty &getSecond() { return *this; }
   const DenseSetEmpty &getSecond() const { return *this; }
 };
-}
 
-/// DenseSet - This implements a dense probed hash-table based set.
-template<typename ValueT, typename ValueInfoT = DenseMapInfo<ValueT> >
-class DenseSet {
-  typedef DenseMap<ValueT, detail::DenseSetEmpty, ValueInfoT,
-                   detail::DenseSetPair<ValueT>> MapTy;
+/// Base class for DenseSet and DenseSmallSet.
+///
+/// MapTy should be either
+///
+///   DenseMap<ValueT, detail::DenseSetEmpty, ValueInfoT,
+///            detail::DenseSetPair<ValueT>>
+///
+/// or the equivalent SmallDenseMap type.  ValueInfoT must implement the
+/// DenseMapInfo "concept".
+template <typename ValueT, typename MapTy, typename ValueInfoT>
+class DenseSetImpl {
   static_assert(sizeof(typename MapTy::value_type) == sizeof(ValueT),
                 "DenseMap buckets unexpectedly large!");
   MapTy TheMap;
@@ -48,7 +53,7 @@ public:
   typedef ValueT value_type;
   typedef unsigned size_type;
 
-  explicit DenseSet(unsigned NumInitBuckets = 0) : TheMap(NumInitBuckets) {}
+  explicit DenseSetImpl(unsigned InitialReserve = 0) : TheMap(InitialReserve) {}
 
   bool empty() const { return TheMap.empty(); }
   size_type size() const { return TheMap.size(); }
@@ -75,15 +80,13 @@ public:
     return TheMap.erase(V);
   }
 
-  void swap(DenseSet& RHS) {
-    TheMap.swap(RHS.TheMap);
-  }
+  void swap(DenseSetImpl &RHS) { TheMap.swap(RHS.TheMap); }
 
   // Iterators.
 
   class Iterator {
     typename MapTy::iterator I;
-    friend class DenseSet;
+    friend class DenseSetImpl;
 
   public:
     typedef typename MapTy::iterator::difference_type difference_type;
@@ -186,6 +189,42 @@ public:
   }
 };
 
+} // namespace detail
+
+/// Implements a dense probed hash-table based set.
+template <typename ValueT, typename ValueInfoT = DenseMapInfo<ValueT>>
+class DenseSet : public detail::DenseSetImpl<
+                     ValueT, DenseMap<ValueT, detail::DenseSetEmpty, ValueInfoT,
+                                      detail::DenseSetPair<ValueT>>,
+                     ValueInfoT> {
+  using BaseT =
+      detail::DenseSetImpl<ValueT,
+                           DenseMap<ValueT, detail::DenseSetEmpty, ValueInfoT,
+                                    detail::DenseSetPair<ValueT>>,
+                           ValueInfoT>;
+
+public:
+  using BaseT::BaseT;
+};
+
+/// Implements a dense probed hash-table based set with some number of buckets
+/// stored inline.
+template <typename ValueT, unsigned InlineBuckets = 4,
+          typename ValueInfoT = DenseMapInfo<ValueT>>
+class SmallDenseSet
+    : public detail::DenseSetImpl<
+          ValueT, SmallDenseMap<ValueT, detail::DenseSetEmpty, InlineBuckets,
+                                ValueInfoT, detail::DenseSetPair<ValueT>>,
+          ValueInfoT> {
+  using BaseT = detail::DenseSetImpl<
+      ValueT, SmallDenseMap<ValueT, detail::DenseSetEmpty, InlineBuckets,
+                            ValueInfoT, detail::DenseSetPair<ValueT>>,
+      ValueInfoT>;
+
+public:
+  using BaseT::BaseT;
+};
+
 } // end namespace llvm
 
 #endif
diff --git a/llvm/unittests/ADT/DenseSetTest.cpp b/llvm/unittests/ADT/DenseSetTest.cpp
index 1e1978e..8397c60 100644
--- a/llvm/unittests/ADT/DenseSetTest.cpp
+++ b/llvm/unittests/ADT/DenseSetTest.cpp
@@ -56,7 +56,11 @@ private:
 
 // Register these types for testing.
 typedef ::testing::Types<DenseSet<unsigned, TestDenseSetInfo>,
-                         const DenseSet<unsigned, TestDenseSetInfo>>
+                         const DenseSet<unsigned, TestDenseSetInfo>,
+                         SmallDenseSet<unsigned, 1, TestDenseSetInfo>,
+                         SmallDenseSet<unsigned, 4, TestDenseSetInfo>,
+                         const SmallDenseSet<unsigned, 4, TestDenseSetInfo>,
+                         SmallDenseSet<unsigned, 64, TestDenseSetInfo>>
     DenseSetTestTypes;
 TYPED_TEST_CASE(DenseSetTest, DenseSetTestTypes);
 
-- 
2.7.4