3 files changed, 115 insertions, 5 deletions

diff --git a/include/llvm/ADT/SmallPtrSet.h b/include/llvm/ADT/SmallPtrSet.h
index bd0d8838ef..fd37cfd025 100644
--- a/include/llvm/ADT/SmallPtrSet.h
+++ b/include/llvm/ADT/SmallPtrSet.h
@@ -60,8 +60,12 @@ protected:
   unsigned NumElements;
   unsigned NumTombstones;
 
-  // Helper to copy construct a SmallPtrSet.
-  SmallPtrSetImpl(const void **SmallStorage, const SmallPtrSetImpl& that);
+  // Helpers to copy and move construct a SmallPtrSet.
+  SmallPtrSetImpl(const void **SmallStorage, const SmallPtrSetImpl &that);
+#if LLVM_HAS_RVALUE_REFERENCES
+  SmallPtrSetImpl(const void **SmallStorage, unsigned SmallSize,
+                  SmallPtrSetImpl &&that);
+#endif
   explicit SmallPtrSetImpl(const void **SmallStorage, unsigned SmallSize) :
     SmallArray(SmallStorage), CurArray(SmallStorage), CurArraySize(SmallSize) {
     assert(SmallSize && (SmallSize & (SmallSize-1)) == 0 &&
@@ -135,6 +139,9 @@ protected:
   void swap(SmallPtrSetImpl &RHS);
 
   void CopyFrom(const SmallPtrSetImpl &RHS);
+#if LLVM_HAS_RVALUE_REFERENCES
+  void MoveFrom(SmallPtrSetImpl &&RHS);
+#endif
 };
 
 /// SmallPtrSetIteratorImpl - This is the common base class shared between all
@@ -242,6 +249,10 @@ class SmallPtrSet : public SmallPtrSetImpl {
 public:
   SmallPtrSet() : SmallPtrSetImpl(SmallStorage, SmallSizePowTwo) {}
   SmallPtrSet(const SmallPtrSet &that) : SmallPtrSetImpl(SmallStorage, that) {}
+#if LLVM_HAS_RVALUE_REFERENCES
+  SmallPtrSet(SmallPtrSet &&that)
+      : SmallPtrSetImpl(SmallStorage, SmallSizePowTwo, std::move(that)) {}
+#endif
 
   template<typename It>
   SmallPtrSet(It I, It E) : SmallPtrSetImpl(SmallStorage, SmallSizePowTwo) {
@@ -280,14 +291,20 @@ public:
     return iterator(CurArray+CurArraySize, CurArray+CurArraySize);
   }
 
-  // Allow assignment from any smallptrset with the same element type even if it
-  // doesn't have the same smallsize.
-  const SmallPtrSet<PtrType, SmallSize>&
+  SmallPtrSet<PtrType, SmallSize> &
   operator=(const SmallPtrSet<PtrType, SmallSize> &RHS) {
     CopyFrom(RHS);
     return *this;
   }
 
+#if LLVM_HAS_RVALUE_REFERENCES
+  SmallPtrSet<PtrType, SmallSize>&
+  operator=(SmallPtrSet<PtrType, SmallSize> &&RHS) {
+    MoveFrom(std::move(RHS));
+    return *this;
+  }
+#endif
+
   /// swap - Swaps the elements of two sets.
   void swap(SmallPtrSet<PtrType, SmallSize> &RHS) {
     SmallPtrSetImpl::swap(RHS);
diff --git a/lib/Support/SmallPtrSet.cpp b/lib/Support/SmallPtrSet.cpp
index dd417b453e..9b86a79351 100644
--- a/lib/Support/SmallPtrSet.cpp
+++ b/lib/Support/SmallPtrSet.cpp
@@ -186,6 +186,29 @@ SmallPtrSetImpl::SmallPtrSetImpl(const void **SmallStorage,
   NumTombstones = that.NumTombstones;
 }
 
+#if LLVM_HAS_RVALUE_REFERENCES
+SmallPtrSetImpl::SmallPtrSetImpl(const void **SmallStorage, unsigned SmallSize,
+                                 SmallPtrSetImpl &&that) {
+  SmallArray = SmallStorage;
+
+  // Copy over the basic members.
+  CurArraySize = that.CurArraySize;
+  NumElements = that.NumElements;
+  NumTombstones = that.NumTombstones;
+
+  // When small, just copy into our small buffer.
+  if (that.isSmall()) {
+    CurArray = SmallArray;
+    memcpy(CurArray, that.CurArray, sizeof(void *) * CurArraySize);
+    return;
+  }
+
+  // Otherwise, we steal the large memory allocation and no copy is needed.
+  CurArray = that.CurArray;
+  that.CurArray = that.SmallArray;
+}
+#endif
+
 /// 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) {
@@ -222,6 +245,27 @@ void SmallPtrSetImpl::CopyFrom(const SmallPtrSetImpl &RHS) {
   NumTombstones = RHS.NumTombstones;
 }
 
+#if LLVM_HAS_RVALUE_REFERENCES
+void SmallPtrSetImpl::MoveFrom(SmallPtrSetImpl &&RHS) {
+  if (!isSmall())
+    free(CurArray);
+
+  if (RHS.isSmall()) {
+    // Copy a small RHS rather than moving.
+    CurArray = SmallArray;
+    memcpy(CurArray, RHS.CurArray, sizeof(void*)*RHS.CurArraySize);
+  } else {
+    CurArray = RHS.CurArray;
+    RHS.CurArray = RHS.SmallArray;
+  }
+
+  // Copy the rest of the trivial members.
+  CurArraySize = RHS.CurArraySize;
+  NumElements = RHS.NumElements;
+  NumTombstones = RHS.NumTombstones;
+}
+#endif
+
 void SmallPtrSetImpl::swap(SmallPtrSetImpl &RHS) {
   if (this == &RHS) return;
 
diff --git a/unittests/ADT/SmallPtrSetTest.cpp b/unittests/ADT/SmallPtrSetTest.cpp
index f85d7c941e..1b564ac05a 100644
--- a/unittests/ADT/SmallPtrSetTest.cpp
+++ b/unittests/ADT/SmallPtrSetTest.cpp
@@ -71,6 +71,55 @@ TEST(SmallPtrSetTest, GrowthTest) {
       EXPECT_EQ(1,buf[i]);
 }
 
+TEST(SmallPtrSetTest, CopyAndMoveTest) {
+  int buf[8];
+  for (int i = 0; i < 8; ++i)
+    buf[i] = 0;
+
+  SmallPtrSet<int *, 4> s1;
+  s1.insert(&buf[0]);
+  s1.insert(&buf[1]);
+  s1.insert(&buf[2]);
+  s1.insert(&buf[3]);
+  EXPECT_EQ(4U, s1.size());
+  for (int i = 0; i < 8; ++i)
+    if (i < 4)
+      EXPECT_TRUE(s1.count(&buf[i]));
+    else
+      EXPECT_FALSE(s1.count(&buf[i]));
+
+  SmallPtrSet<int *, 4> s2(s1);
+  EXPECT_EQ(4U, s2.size());
+  for (int i = 0; i < 8; ++i)
+    if (i < 4)
+      EXPECT_TRUE(s2.count(&buf[i]));
+    else
+      EXPECT_FALSE(s2.count(&buf[i]));
+
+  s1 = s2;
+  EXPECT_EQ(4U, s1.size());
+  for (int i = 0; i < 8; ++i)
+    if (i < 4)
+      EXPECT_TRUE(s1.count(&buf[i]));
+    else
+      EXPECT_FALSE(s1.count(&buf[i]));
+
+  SmallPtrSet<int *, 4> s3(llvm_move(s1));
+  EXPECT_EQ(4U, s3.size());
+  for (int i = 0; i < 8; ++i)
+    if (i < 4)
+      EXPECT_TRUE(s3.count(&buf[i]));
+    else
+      EXPECT_FALSE(s3.count(&buf[i]));
+
+  s1 = llvm_move(s3);
+  EXPECT_EQ(4U, s1.size());
+  for (int i = 0; i < 8; ++i)
+    if (i < 4)
+      EXPECT_TRUE(s1.count(&buf[i]));
+    else
+      EXPECT_FALSE(s1.count(&buf[i]));
+}
 
 TEST(SmallPtrSetTest, SwapTest) {
   int buf[10];