Adding reserve and capacity methods to FoldingSet

http://reviews.llvm.org/D20930

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@271669 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/ADT/FoldingSet.h

@@ -180,11 +180,21 @@ public:
   /// empty - Returns true if there are no nodes in the folding set.
   bool empty() const { return NumNodes == 0; }
 
+  void reserve(unsigned EltCount);
+  unsigned capacity() {
+    // We allow a load factor of up to 2.0,
+    // so that means our capacity is NumBuckets * 2
+    return NumBuckets * 2;
+  }
+
 private:
   /// GrowHashTable - Double the size of the hash table and rehash everything.
-  ///
   void GrowHashTable();
 
+  /// GrowBucketCount - resize the hash table and rehash everything.
+  /// NewBucketCount must be a power of two, and must be greater than the old
+  /// bucket count.
+  void GrowBucketCount(unsigned NewBucketCount);
 protected:
   /// GetNodeProfile - Instantiations of the FoldingSet template implement
   /// this function to gather data bits for the given node.

lib/Support/FoldingSet.cpp

@@ -266,12 +266,12 @@ void FoldingSetImpl::clear() {
   NumNodes = 0;
 }
 
-/// GrowHashTable - Double the size of the hash table and rehash everything.
-///
-void FoldingSetImpl::GrowHashTable() {
+void FoldingSetImpl::GrowBucketCount(unsigned NewBucketCount) {
+  assert((NewBucketCount > NumBuckets) && "Can't shrink a folding set with GrowBucketCount");
+  assert(isPowerOf2_32(NewBucketCount) && "Bad bucket count!");
   void **OldBuckets = Buckets;
   unsigned OldNumBuckets = NumBuckets;
-  NumBuckets <<= 1;
+  NumBuckets = NewBucketCount;
   
   // Clear out new buckets.
   Buckets = AllocateBuckets(NumBuckets);
@@ -298,6 +298,21 @@ void FoldingSetImpl::GrowHashTable() {
   free(OldBuckets);
 }
 
+/// GrowHashTable - Double the size of the hash table and rehash everything.
+///
+void FoldingSetImpl::GrowHashTable() {
+  GrowBucketCount(NumBuckets * 2);
+}
+
+void FoldingSetImpl::reserve(unsigned EltCount) {
+  // This will give us somewhere between EltCount / 2 and
+  // EltCount buckets.  This puts us in the load factor
+  // range of 1.0 - 2.0.
+  if(EltCount < capacity())
+    return;
+  GrowBucketCount(PowerOf2Floor(EltCount));
+}
+
 /// FindNodeOrInsertPos - Look up the node specified by ID.  If it exists,
 /// return it.  If not, return the insertion token that will make insertion
 /// faster.
@@ -330,7 +345,7 @@ FoldingSetImpl::Node
 void FoldingSetImpl::InsertNode(Node *N, void *InsertPos) {
   assert(!N->getNextInBucket());
   // Do we need to grow the hashtable?
-  if (NumNodes+1 > NumBuckets*2) {
+  if (NumNodes+1 > capacity()) {
     GrowHashTable();
     FoldingSetNodeID TempID;
     InsertPos = GetBucketFor(ComputeNodeHash(N, TempID), Buckets, NumBuckets);

unittests/ADT/FoldingSet.cpp

@@ -35,5 +35,138 @@ TEST(FoldingSetTest, UnalignedStringTest) {
   EXPECT_EQ(a.ComputeHash(), b.ComputeHash());
 }
 
+struct TrivialPair : public FoldingSetNode {
+  unsigned Key = 0;
+  unsigned Value = 0;
+  TrivialPair(unsigned K, unsigned V) : FoldingSetNode(), Key(K), Value(V) {}
+
+  void Profile(FoldingSetNodeID &ID) const {
+    ID.AddInteger(Key);
+    ID.AddInteger(Value);
+  }
+};
+
+TEST(FoldingSetTest, IDComparison) {
+  FoldingSet<TrivialPair> Trivial;
+
+  TrivialPair T(99, 42);
+  Trivial.InsertNode(&T);
+
+  void *InsertPos = nullptr;
+  FoldingSetNodeID ID;
+  T.Profile(ID);
+  TrivialPair *N = Trivial.FindNodeOrInsertPos(ID, InsertPos);
+  EXPECT_EQ(&T, N);
+  EXPECT_EQ(nullptr, InsertPos);
+}
+
+TEST(FoldingSetTest, MissedIDComparison) {
+  FoldingSet<TrivialPair> Trivial;
+
+  TrivialPair S(100, 42);
+  TrivialPair T(99, 42);
+  Trivial.InsertNode(&T);
+
+  void *InsertPos = nullptr;
+  FoldingSetNodeID ID;
+  S.Profile(ID);
+  TrivialPair *N = Trivial.FindNodeOrInsertPos(ID, InsertPos);
+  EXPECT_EQ(nullptr, N);
+  EXPECT_NE(nullptr, InsertPos);
+}
+
+TEST(FoldingSetTest, RemoveNodeThatIsPresent) {
+  FoldingSet<TrivialPair> Trivial;
+
+  TrivialPair T(99, 42);
+  Trivial.InsertNode(&T);
+  EXPECT_EQ(Trivial.size(), 1U);
+
+  bool WasThere = Trivial.RemoveNode(&T);
+  EXPECT_TRUE(WasThere);
+  EXPECT_EQ(0U, Trivial.size());
+}
+
+TEST(FoldingSetTest, RemoveNodeThatIsAbsent) {
+  FoldingSet<TrivialPair> Trivial;
+
+  TrivialPair T(99, 42);
+  bool WasThere = Trivial.RemoveNode(&T);
+  EXPECT_FALSE(WasThere);
+  EXPECT_EQ(0U, Trivial.size());
+}
+
+TEST(FoldingSetTest, GetOrInsertInserting) {
+  FoldingSet<TrivialPair> Trivial;
+
+  TrivialPair T(99, 42);
+  TrivialPair *N = Trivial.GetOrInsertNode(&T);
+  EXPECT_EQ(&T, N);
+}
+
+TEST(FoldingSetTest, GetOrInsertGetting) {
+  FoldingSet<TrivialPair> Trivial;
+
+  TrivialPair T(99, 42);
+  TrivialPair T2(99, 42);
+  Trivial.InsertNode(&T);
+  TrivialPair *N = Trivial.GetOrInsertNode(&T2);
+  EXPECT_EQ(&T, N);
+}
+
+TEST(FoldingSetTest, InsertAtPos) {
+  FoldingSet<TrivialPair> Trivial;
+
+  void *InsertPos = nullptr;
+  TrivialPair Finder(99, 42);
+  FoldingSetNodeID ID;
+  Finder.Profile(ID);
+  Trivial.FindNodeOrInsertPos(ID, InsertPos);
+
+  TrivialPair T(99, 42);
+  Trivial.InsertNode(&T, InsertPos);
+  EXPECT_EQ(1U, Trivial.size());
+}
+
+TEST(FoldingSetTest, EmptyIsTrue) {
+  FoldingSet<TrivialPair> Trivial;
+  EXPECT_TRUE(Trivial.empty());
+}
+
+TEST(FoldingSetTest, EmptyIsFalse) {
+  FoldingSet<TrivialPair> Trivial;
+  TrivialPair T(99, 42);
+  Trivial.InsertNode(&T);
+  EXPECT_FALSE(Trivial.empty());
+}
+
+TEST(FoldingSetTest, ClearOnEmpty) {
+  FoldingSet<TrivialPair> Trivial;
+  Trivial.clear();
+  EXPECT_TRUE(Trivial.empty());
+}
+
+TEST(FoldingSetTest, ClearOnNonEmpty) {
+  FoldingSet<TrivialPair> Trivial;
+  TrivialPair T(99, 42);
+  Trivial.InsertNode(&T);
+  Trivial.clear();
+  EXPECT_TRUE(Trivial.empty());
+}
+
+TEST(FoldingSetTest, CapacityLargerThanReserve) {
+  FoldingSet<TrivialPair> Trivial;
+  auto OldCapacity = Trivial.capacity();
+  Trivial.reserve(OldCapacity + 1);
+  EXPECT_GE(Trivial.capacity(), OldCapacity + 1);
+}
+
+TEST(FoldingSetTest, SmallReserveChangesNothing) {
+  FoldingSet<TrivialPair> Trivial;
+  auto OldCapacity = Trivial.capacity();
+  Trivial.reserve(OldCapacity - 1);
+  EXPECT_EQ(Trivial.capacity(), OldCapacity);
+}
+
 }