implement operator= for smallptrset

llvm-svn: 38460

2 files changed, 50 insertions, 0 deletions

diff --git a/llvm/include/llvm/ADT/SmallPtrSet.h b/llvm/include/llvm/ADT/SmallPtrSet.h
index 40e9bb4139c..80c078f35bd 100644
--- a/llvm/include/llvm/ADT/SmallPtrSet.h
+++ b/llvm/include/llvm/ADT/SmallPtrSet.h
@@ -127,6 +127,10 @@ private:
   
   /// Grow - Allocate a larger backing store for the buckets and move it over.
   void Grow();
+  
+  void operator=(const SmallPtrSetImpl &RHS);  // DO NOT IMPLEMENT.
+protected:
+  void CopyFrom(const SmallPtrSetImpl &RHS);
 };
 
 /// SmallPtrSetIteratorImpl - This is the common base class shared between all
@@ -233,6 +237,16 @@ public:
   inline iterator end() const {
     return iterator(CurArray+CurArraySize);
   }
+  
+  // Allow assignment from any smallptrset with the same element type even if it
+  // doesn't have the same smallsize.
+  template<unsigned RHSSize>
+  const SmallPtrSet<PtrType, SmallSize>
+  operator=(const SmallPtrSet<PtrType, RHSSize> &RHS) {
+    CopyFrom(RHS);
+    return *this;
+  }
+
 };
 
 }
diff --git a/llvm/lib/Support/SmallPtrSet.cpp b/llvm/lib/Support/SmallPtrSet.cpp
index da1a11029e7..81c4bc752fb 100644
--- a/llvm/lib/Support/SmallPtrSet.cpp
+++ b/llvm/lib/Support/SmallPtrSet.cpp
@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Support/MathExtras.h"
 using namespace llvm;
 
 bool SmallPtrSetImpl::insert(void *Ptr) {
@@ -172,3 +173,38 @@ SmallPtrSetImpl::SmallPtrSetImpl(const SmallPtrSetImpl& that) {
     }
   }
 }
+
+/// CopyFrom - implement operator= from a smallptrset that has the same pointer
+/// type, but may have a different small size.
+void SmallPtrSetImpl::CopyFrom(const SmallPtrSetImpl &RHS) {
+  // Allocate space if needed or clear the current elements out of the array.
+  if (CurArraySize < RHS.size()*2) {
+    if (!isSmall())
+      delete [] CurArray;
+    
+    // Get a power of two larger than twice the RHS size.
+    CurArraySize = 1 << Log2_32(RHS.size()*4);
+    
+    // Install the new array.  Clear all the buckets to empty.
+    CurArray = new void*[CurArraySize+1];
+    memset(CurArray, -1, CurArraySize*sizeof(void*));
+    
+    // The end pointer, always valid, is set to a valid element to help the
+    // iterator.
+    CurArray[CurArraySize] = 0;
+    
+  } else if (!empty()) {
+    clear();
+  }
+  
+  // Now that we know we have enough space, and that the current array is empty,
+  // copy over all the elements from the RHS.
+  
+  for (void **BucketPtr = RHS.CurArray, **E = RHS.CurArray+RHS.CurArraySize;
+       BucketPtr != E; ++BucketPtr) {
+    // Copy over the element if it is valid.
+    void *Elt = *BucketPtr;
+    if (Elt != getTombstoneMarker() && Elt != getEmptyMarker())
+      *const_cast<void**>(FindBucketFor(Elt)) = Elt;
+  }
+}