3 files changed, 36 insertions, 10 deletions

diff --git a/llvm/include/llvm/ADT/SmallPtrSet.h b/llvm/include/llvm/ADT/SmallPtrSet.h
index fd37cfd0259..c7f5e5bfc0f 100644
--- a/llvm/include/llvm/ADT/SmallPtrSet.h
+++ b/llvm/include/llvm/ADT/SmallPtrSet.h
@@ -140,7 +140,7 @@ protected:
 
   void CopyFrom(const SmallPtrSetImpl &RHS);
 #if LLVM_HAS_RVALUE_REFERENCES
-  void MoveFrom(SmallPtrSetImpl &&RHS);
+  void MoveFrom(unsigned SmallSize, SmallPtrSetImpl &&RHS);
 #endif
 };
 
@@ -300,7 +300,7 @@ public:
 #if LLVM_HAS_RVALUE_REFERENCES
   SmallPtrSet<PtrType, SmallSize>&
   operator=(SmallPtrSet<PtrType, SmallSize> &&RHS) {
-    MoveFrom(std::move(RHS));
+    MoveFrom(SmallSizePowTwo, std::move(RHS));
     return *this;
   }
 #endif
diff --git a/llvm/lib/Support/SmallPtrSet.cpp b/llvm/lib/Support/SmallPtrSet.cpp
index 9b86a793513..e37e23b7bb1 100644
--- a/llvm/lib/Support/SmallPtrSet.cpp
+++ b/llvm/lib/Support/SmallPtrSet.cpp
@@ -206,6 +206,12 @@ SmallPtrSetImpl::SmallPtrSetImpl(const void **SmallStorage, unsigned SmallSize,
   // Otherwise, we steal the large memory allocation and no copy is needed.
   CurArray = that.CurArray;
   that.CurArray = that.SmallArray;
+
+  // Make the "that" object small and empty.
+  that.CurArraySize = SmallSize;
+  assert(that.CurArray == that.SmallArray);
+  that.NumElements = 0;
+  that.NumTombstones = 0;
 }
 #endif
 
@@ -246,7 +252,7 @@ void SmallPtrSetImpl::CopyFrom(const SmallPtrSetImpl &RHS) {
 }
 
 #if LLVM_HAS_RVALUE_REFERENCES
-void SmallPtrSetImpl::MoveFrom(SmallPtrSetImpl &&RHS) {
+void SmallPtrSetImpl::MoveFrom(unsigned SmallSize, SmallPtrSetImpl &&RHS) {
   if (!isSmall())
     free(CurArray);
 
@@ -263,6 +269,12 @@ void SmallPtrSetImpl::MoveFrom(SmallPtrSetImpl &&RHS) {
   CurArraySize = RHS.CurArraySize;
   NumElements = RHS.NumElements;
   NumTombstones = RHS.NumTombstones;
+
+  // Make the RHS small and empty.
+  RHS.CurArraySize = SmallSize;
+  assert(RHS.CurArray == RHS.SmallArray);
+  RHS.NumElements = 0;
+  RHS.NumTombstones = 0;
 }
 #endif
 
diff --git a/llvm/unittests/ADT/SmallPtrSetTest.cpp b/llvm/unittests/ADT/SmallPtrSetTest.cpp
index d885f69b211..9b71cf07f23 100644
--- a/llvm/unittests/ADT/SmallPtrSetTest.cpp
+++ b/llvm/unittests/ADT/SmallPtrSetTest.cpp
@@ -121,28 +121,42 @@ TEST(SmallPtrSetTest, CopyAndMoveTest) {
 
   s1 = s2;
   EXPECT_EQ(4U, s1.size());
+  EXPECT_EQ(4U, s2.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));
+#if LLVM_HAS_RVALUE_REFERENCES
+  SmallPtrSet<int *, 4> s3(std::move(s1));
   EXPECT_EQ(4U, s3.size());
+  EXPECT_TRUE(s1.empty());
   for (int i = 0; i < 8; ++i)
     if (i < 4)
       EXPECT_TRUE(s3.count(&buf[i]));
     else
       EXPECT_FALSE(s3.count(&buf[i]));
 
-  // Move assign to the moved-from object.
+  // Move assign into the moved-from object. Also test move of a non-small
+  // container.
+  s3.insert(&buf[4]);
+  s3.insert(&buf[5]);
+  s3.insert(&buf[6]);
+  s3.insert(&buf[7]);
   s1 = llvm_move(s3);
-  EXPECT_EQ(4U, s1.size());
+  EXPECT_EQ(8U, s1.size());
+  EXPECT_TRUE(s3.empty());
   for (int i = 0; i < 8; ++i)
-    if (i < 4)
-      EXPECT_TRUE(s1.count(&buf[i]));
-    else
-      EXPECT_FALSE(s1.count(&buf[i]));
+    EXPECT_TRUE(s1.count(&buf[i]));
+
+  // Copy assign into a moved-from object.
+  s3 = s1;
+  EXPECT_EQ(8U, s3.size());
+  EXPECT_EQ(8U, s1.size());
+  for (int i = 0; i < 8; ++i)
+    EXPECT_TRUE(s3.count(&buf[i]));
+#endif
 }
 
 TEST(SmallPtrSetTest, SwapTest) {