[ADT] Add SmallDenseSet.

[jlebar]

Summary

This matches SmallDenseMap.

llvm/include/llvm/ADT/DenseSet.h

 //===- llvm/ADT/DenseSet.h - Dense probed hash table ------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
-// This file defines the DenseSet class.
+// This file defines the DenseSet and SmallDenseSet classes.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ADT_DENSESET_H
 #define LLVM_ADT_DENSESET_H
 
 #include "llvm/ADT/DenseMap.h"
 
 namespace llvm {
 
 namespace detail {
 struct DenseSetEmpty {};
 
 // Use the empty base class trick so we can create a DenseMap where the buckets
 // contain only a single item.
 template <typename KeyT> class DenseSetPair : public DenseSetEmpty {
   KeyT key;
 
 public:
   KeyT &getFirst() { return key; }
   const KeyT &getFirst() const { return key; }
   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;
 
 public:
   typedef ValueT key_type;
   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(); }
   size_t getMemorySize() const { return TheMap.getMemorySize(); }
 
   /// Grow the DenseSet so that it has at least Size buckets. Will not shrink
   /// the Size of the set.
   void resize(size_t Size) { TheMap.resize(Size); }
public:
   size_type count(const ValueT &V) const {
     return TheMap.count(V);
   }
 
   bool erase(const ValueT &V) {
     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;
     typedef ValueT value_type;
     typedef value_type *pointer;
     typedef value_type &reference;
     typedef std::forward_iterator_tag iterator_category;
 
public:
   // Range insertion of values.
   template<typename InputIt>
   void insert(InputIt I, InputIt E) {
     for (; I != E; ++I)
       insert(*I);
   }
 };
 
+} // 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

llvm/unittests/ADT/DenseSetTest.cpp

static T GetTestSet() {
     Set.insert(1);
     Set.insert(2);
     return Set;
   }
 };
 
 // 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);
 
 TYPED_TEST(DenseSetTest, FindAsTest) {
   auto &set = this->Set;
   // Size tests
   EXPECT_EQ(3u, set.size());