1 files changed, 76 insertions, 20 deletions

diff --git a/include/llvm/ADT/SmallPtrSet.h b/include/llvm/ADT/SmallPtrSet.h
index 3d98e8fac43b6..eaed6aa05dcbc 100644
--- a/include/llvm/ADT/SmallPtrSet.h
+++ b/include/llvm/ADT/SmallPtrSet.h
@@ -21,6 +21,7 @@
 #include <cassert>
 #include <cstddef>
 #include <cstring>
+#include <cstdlib>
 #include <iterator>
 #include <utility>
 
@@ -58,36 +59,45 @@ protected:
   /// CurArraySize - The allocated size of CurArray, always a power of two.
   unsigned CurArraySize;
 
-  // If small, this is # elts allocated consecutively
-  unsigned NumElements;
+  /// Number of elements in CurArray that contain a value or are a tombstone.
+  /// If small, all these elements are at the beginning of CurArray and the rest
+  /// is uninitialized.
+  unsigned NumNonEmpty;
+  /// Number of tombstones in CurArray.
   unsigned NumTombstones;
 
   // Helpers to copy and move construct a SmallPtrSet.
-  SmallPtrSetImplBase(const void **SmallStorage, const SmallPtrSetImplBase &that);
+  SmallPtrSetImplBase(const void **SmallStorage,
+                      const SmallPtrSetImplBase &that);
   SmallPtrSetImplBase(const void **SmallStorage, unsigned SmallSize,
-                  SmallPtrSetImplBase &&that);
-  explicit SmallPtrSetImplBase(const void **SmallStorage, unsigned SmallSize) :
-    SmallArray(SmallStorage), CurArray(SmallStorage), CurArraySize(SmallSize) {
+                      SmallPtrSetImplBase &&that);
+  explicit SmallPtrSetImplBase(const void **SmallStorage, unsigned SmallSize)
+      : SmallArray(SmallStorage), CurArray(SmallStorage),
+        CurArraySize(SmallSize), NumNonEmpty(0), NumTombstones(0) {
     assert(SmallSize && (SmallSize & (SmallSize-1)) == 0 &&
            "Initial size must be a power of two!");
-    clear();
   }
-  ~SmallPtrSetImplBase();
+  ~SmallPtrSetImplBase() {
+    if (!isSmall())
+      free(CurArray);
+  }
 
 public:
   typedef unsigned size_type;
   bool LLVM_ATTRIBUTE_UNUSED_RESULT empty() const { return size() == 0; }
-  size_type size() const { return NumElements; }
+  size_type size() const { return NumNonEmpty - NumTombstones; }
 
   void clear() {
     // If the capacity of the array is huge, and the # elements used is small,
     // shrink the array.
-    if (!isSmall() && NumElements*4 < CurArraySize && CurArraySize > 32)
-      return shrink_and_clear();
+    if (!isSmall()) {
+      if (size() * 4 < CurArraySize && CurArraySize > 32)
+        return shrink_and_clear();
+      // Fill the array with empty markers.
+      memset(CurArray, -1, CurArraySize * sizeof(void *));
+    }
 
-    // Fill the array with empty markers.
-    memset(CurArray, -1, CurArraySize*sizeof(void*));
-    NumElements = 0;
+    NumNonEmpty = 0;
     NumTombstones = 0;
   }
 
@@ -99,10 +109,42 @@ protected:
     return reinterpret_cast<void*>(-1);
   }
 
+  const void **EndPointer() const {
+    return isSmall() ? CurArray + NumNonEmpty : CurArray + CurArraySize;
+  }
+
   /// insert_imp - This returns true if the pointer was new to the set, false if
   /// it was already in the set.  This is hidden from the client so that the
   /// derived class can check that the right type of pointer is passed in.
-  std::pair<const void *const *, bool> insert_imp(const void *Ptr);
+  std::pair<const void *const *, bool> insert_imp(const void *Ptr) {
+    if (isSmall()) {
+      // Check to see if it is already in the set.
+      const void **LastTombstone = nullptr;
+      for (const void **APtr = SmallArray, **E = SmallArray + NumNonEmpty;
+           APtr != E; ++APtr) {
+        const void *Value = *APtr;
+        if (Value == Ptr)
+          return std::make_pair(APtr, false);
+        if (Value == getTombstoneMarker())
+          LastTombstone = APtr;
+      }
+
+      // Did we find any tombstone marker?
+      if (LastTombstone != nullptr) {
+        *LastTombstone = Ptr;
+        --NumTombstones;
+        return std::make_pair(LastTombstone, true);
+      }
+
+      // Nope, there isn't.  If we stay small, just 'pushback' now.
+      if (NumNonEmpty < CurArraySize) {
+        SmallArray[NumNonEmpty++] = Ptr;
+        return std::make_pair(SmallArray + (NumNonEmpty - 1), true);
+      }
+      // Otherwise, hit the big set case, which will call grow.
+    }
+    return insert_imp_big(Ptr);
+  }
 
   /// erase_imp - If the set contains the specified pointer, remove it and
   /// return true, otherwise return false.  This is hidden from the client so
@@ -114,7 +156,7 @@ protected:
     if (isSmall()) {
       // Linear search for the item.
       for (const void *const *APtr = SmallArray,
-                      *const *E = SmallArray+NumElements; APtr != E; ++APtr)
+                      *const *E = SmallArray + NumNonEmpty; APtr != E; ++APtr)
         if (*APtr == Ptr)
           return true;
       return false;
@@ -127,6 +169,8 @@ protected:
 private:
   bool isSmall() const { return CurArray == SmallArray; }
 
+  std::pair<const void *const *, bool> insert_imp_big(const void *Ptr);
+
   const void * const *FindBucketFor(const void *Ptr) const;
   void shrink_and_clear();
 
@@ -142,6 +186,12 @@ protected:
 
   void CopyFrom(const SmallPtrSetImplBase &RHS);
   void MoveFrom(unsigned SmallSize, SmallPtrSetImplBase &&RHS);
+
+private:
+  /// Code shared by MoveFrom() and move constructor.
+  void MoveHelper(unsigned SmallSize, SmallPtrSetImplBase &&RHS);
+  /// Code shared by CopyFrom() and copy constructor.
+  void CopyHelper(const SmallPtrSetImplBase &RHS);
 };
 
 /// SmallPtrSetIteratorImpl - This is the common base class shared between all
@@ -154,7 +204,7 @@ protected:
 public:
   explicit SmallPtrSetIteratorImpl(const void *const *BP, const void*const *E)
     : Bucket(BP), End(E) {
-      AdvanceIfNotValid();
+    AdvanceIfNotValid();
   }
 
   bool operator==(const SmallPtrSetIteratorImpl &RHS) const {
@@ -266,7 +316,7 @@ public:
   /// the element equal to Ptr.
   std::pair<iterator, bool> insert(PtrType Ptr) {
     auto p = insert_imp(PtrTraits::getAsVoidPointer(Ptr));
-    return std::make_pair(iterator(p.first, CurArray + CurArraySize), p.second);
+    return std::make_pair(iterator(p.first, EndPointer()), p.second);
   }
 
   /// erase - If the set contains the specified pointer, remove it and return
@@ -287,10 +337,11 @@ public:
   }
 
   inline iterator begin() const {
-    return iterator(CurArray, CurArray+CurArraySize);
+    return iterator(CurArray, EndPointer());
   }
   inline iterator end() const {
-    return iterator(CurArray+CurArraySize, CurArray+CurArraySize);
+    const void *const *End = EndPointer();
+    return iterator(End, End);
   }
 };
 
@@ -300,6 +351,11 @@ public:
 /// SmallPtrSetImplBase for details of the algorithm.
 template<class PtrType, unsigned SmallSize>
 class SmallPtrSet : public SmallPtrSetImpl<PtrType> {
+  // In small mode SmallPtrSet uses linear search for the elements, so it is
+  // not a good idea to choose this value too high. You may consider using a
+  // DenseSet<> instead if you expect many elements in the set.
+  static_assert(SmallSize <= 32, "SmallSize should be small");
+
   typedef SmallPtrSetImpl<PtrType> BaseT;
 
   // Make sure that SmallSize is a power of two, round up if not.