2 files changed, 37 insertions, 3 deletions
diff --git a/include/llvm/ADT/SmallPtrSet.h b/include/llvm/ADT/SmallPtrSet.h
index d2a68ff90f..b9f7b1fc26 100644
--- a/include/llvm/ADT/SmallPtrSet.h
+++ b/include/llvm/ADT/SmallPtrSet.h
@@ -67,8 +67,6 @@ public:
       delete[] CurArray;
   }
   
-  bool isSmall() const { return CurArray == &SmallArray[0]; }
-
   static void *getTombstoneMarker() { return reinterpret_cast<void*>(-2); }
   static void *getEmptyMarker() {
     // Note that -1 is chosen to make clear() efficiently implementable with
@@ -86,6 +84,10 @@ public:
   /// was already in the set.
   bool insert(void *Ptr);
   
+  /// erase - If the set contains the specified pointer, remove it and return
+  /// true, otherwise return false.
+  bool erase(void *Ptr);
+  
   bool count(void *Ptr) const {
     if (isSmall()) {
       // Linear search for the item.
@@ -101,6 +103,8 @@ public:
   }
   
 private:
+  bool isSmall() const { return CurArray == &SmallArray[0]; }
+
   unsigned Hash(void *Ptr) const {
     return ((uintptr_t)Ptr >> 4) & (CurArraySize-1);
   }
@@ -188,7 +192,10 @@ struct NextPowerOfTwo {
 };
 
 
-/// SmallPtrSet - This class implements 
+/// SmallPtrSet - This class implements a set which is optimizer for holding
+/// SmallSize or less elements.  This internally rounds up SmallSize to the next
+/// power of two if it is not already a power of two.  See the comments above
+/// SmallPtrSetImpl for details of the algorithm.
 template<class PtrType, unsigned SmallSize>
 class SmallPtrSet : public SmallPtrSetImpl {
   // Make sure that SmallSize is a power of two, round up if not.
diff --git a/lib/Support/SmallPtrSet.cpp b/lib/Support/SmallPtrSet.cpp
index 48552a56fc..1eea7272e0 100644
--- a/lib/Support/SmallPtrSet.cpp
+++ b/lib/Support/SmallPtrSet.cpp
@@ -45,6 +45,33 @@ bool SmallPtrSetImpl::insert(void *Ptr) {
   return true;
 }
 
+bool SmallPtrSetImpl::erase(void *Ptr) {
+  if (isSmall()) {
+    // Check to see if it is in the set.
+    for (void **APtr = SmallArray, **E = SmallArray+NumElements;
+         APtr != E; ++APtr)
+      if (*APtr == Ptr) {
+        // If it is in the set, move everything over, replacing this element.
+        memmove(APtr, APtr+1, sizeof(void*)*(E-APtr-1));
+        // Clear the end element.
+        E[-1] = getEmptyMarker();
+        --NumElements;
+        return false;
+      }
+    
+    return false;
+  }
+  
+  // Okay, we know we have space.  Find a hash bucket.
+  void **Bucket = const_cast<void**>(FindBucketFor(Ptr));
+  if (*Bucket != Ptr) return false;  // Not in the set?
+
+  // Set this as a tombstone.
+  *Bucket = getTombstoneMarker();
+  --NumElements;
+  return true;
+}
+
 void * const *SmallPtrSetImpl::FindBucketFor(void *Ptr) const {
   unsigned Bucket = Hash(Ptr);
   unsigned ArraySize = CurArraySize;