diff --git a/include/llvm/ADT/SmallPtrSet.h b/include/llvm/ADT/SmallPtrSet.h index 9992858d67b..70693d5b9aa 100644 --- a/include/llvm/ADT/SmallPtrSet.h +++ b/include/llvm/ADT/SmallPtrSet.h @@ -137,6 +137,10 @@ private: void operator=(const SmallPtrSetImpl &RHS); // DO NOT IMPLEMENT. protected: + /// swap - Swaps the elements of two sets. + /// Note: This method assumes that both sets have the same small size. + void swap(SmallPtrSetImpl &RHS); + void CopyFrom(const SmallPtrSetImpl &RHS); }; @@ -287,8 +291,20 @@ public: return *this; } + /// swap - Swaps the elements of two sets. + void swap(SmallPtrSet &RHS) { + SmallPtrSetImpl::swap(RHS); + } }; } +namespace std { + /// Implement std::swap in terms of SmallPtrSet swap. + template + inline void swap(llvm::SmallPtrSet &LHS, llvm::SmallPtrSet &RHS) { + LHS.swap(RHS); + } +} + #endif diff --git a/lib/Support/SmallPtrSet.cpp b/lib/Support/SmallPtrSet.cpp index 997ce0b74cd..dd516d37f05 100644 --- a/lib/Support/SmallPtrSet.cpp +++ b/lib/Support/SmallPtrSet.cpp @@ -14,6 +14,7 @@ #include "llvm/ADT/SmallPtrSet.h" #include "llvm/Support/MathExtras.h" +#include #include using namespace llvm; @@ -223,6 +224,55 @@ void SmallPtrSetImpl::CopyFrom(const SmallPtrSetImpl &RHS) { NumTombstones = RHS.NumTombstones; } +void SmallPtrSetImpl::swap(SmallPtrSetImpl &RHS) { + if (this == &RHS) return; + + // We can only avoid copying elements if neither set is small. + if (!this->isSmall() && !RHS.isSmall()) { + std::swap(this->CurArray, RHS.CurArray); + std::swap(this->CurArraySize, RHS.CurArraySize); + std::swap(this->NumElements, RHS.NumElements); + std::swap(this->NumTombstones, RHS.NumTombstones); + return; + } + + // FIXME: From here on we assume that both sets have the same small size. + + // If only RHS is small, copy the small elements into LHS and move the pointer + // from LHS to RHS. + if (!this->isSmall() && RHS.isSmall()) { + std::copy(RHS.SmallArray, RHS.SmallArray+RHS.NumElements, this->SmallArray); + std::swap(this->NumElements, RHS.NumElements); + std::swap(this->CurArraySize, RHS.CurArraySize); + RHS.CurArray = this->CurArray; + RHS.NumTombstones = this->NumTombstones; + this->CurArray = this->SmallArray; + this->NumTombstones = 0; + return; + } + + // If only LHS is small, copy the small elements into RHS and move the pointer + // from RHS to LHS. + if (this->isSmall() && !RHS.isSmall()) { + std::copy(this->SmallArray, this->SmallArray+this->NumElements, + RHS.SmallArray); + std::swap(RHS.NumElements, this->NumElements); + std::swap(RHS.CurArraySize, this->CurArraySize); + this->CurArray = RHS.CurArray; + this->NumTombstones = RHS.NumTombstones; + RHS.CurArray = RHS.SmallArray; + RHS.NumTombstones = 0; + return; + } + + // Both a small, just swap the small elements. + assert(this->isSmall() && RHS.isSmall()); + assert(this->CurArraySize == RHS.CurArraySize); + unsigned MaxElems = std::max(this->NumElements, RHS.NumElements); + std::swap_ranges(this->SmallArray, this->SmallArray+MaxElems, RHS.SmallArray); + std::swap(this->NumElements, RHS.NumElements); +} + SmallPtrSetImpl::~SmallPtrSetImpl() { if (!isSmall()) free(CurArray); diff --git a/unittests/ADT/SmallPtrSetTest.cpp b/unittests/ADT/SmallPtrSetTest.cpp new file mode 100644 index 00000000000..9114875e003 --- /dev/null +++ b/unittests/ADT/SmallPtrSetTest.cpp @@ -0,0 +1,72 @@ +//===- llvm/unittest/ADT/SmallPtrSetTest.cpp ------------------------------===// +// +// The LLVM Compiler Infrastructure +// +// This file is distributed under the University of Illinois Open Source +// License. See LICENSE.TXT for details. +// +//===----------------------------------------------------------------------===// +// +// SmallPtrSet unit tests. +// +//===----------------------------------------------------------------------===// + +#include "gtest/gtest.h" +#include "llvm/ADT/SmallPtrSet.h" + +using namespace llvm; + +// SmallPtrSet swapping test. +TEST(SmallPtrSetTest, SwapTest) { + int buf[10]; + + SmallPtrSet a; + SmallPtrSet b; + + a.insert(&buf[0]); + a.insert(&buf[1]); + b.insert(&buf[2]); + + std::swap(a, b); + + EXPECT_EQ(1U, a.size()); + EXPECT_EQ(2U, b.size()); + EXPECT_TRUE(a.count(&buf[2])); + EXPECT_TRUE(b.count(&buf[0])); + EXPECT_TRUE(b.count(&buf[1])); + + b.insert(&buf[3]); + std::swap(a, b); + + EXPECT_EQ(3U, a.size()); + EXPECT_EQ(1U, b.size()); + EXPECT_TRUE(a.count(&buf[0])); + EXPECT_TRUE(a.count(&buf[1])); + EXPECT_TRUE(a.count(&buf[3])); + EXPECT_TRUE(b.count(&buf[2])); + + std::swap(a, b); + + EXPECT_EQ(1U, a.size()); + EXPECT_EQ(3U, b.size()); + EXPECT_TRUE(a.count(&buf[2])); + EXPECT_TRUE(b.count(&buf[0])); + EXPECT_TRUE(b.count(&buf[1])); + EXPECT_TRUE(b.count(&buf[3])); + + a.insert(&buf[4]); + a.insert(&buf[5]); + a.insert(&buf[6]); + + std::swap(b, a); + + EXPECT_EQ(3U, a.size()); + EXPECT_EQ(4U, b.size()); + EXPECT_TRUE(b.count(&buf[2])); + EXPECT_TRUE(b.count(&buf[4])); + EXPECT_TRUE(b.count(&buf[5])); + EXPECT_TRUE(b.count(&buf[6])); + EXPECT_TRUE(a.count(&buf[0])); + EXPECT_TRUE(a.count(&buf[1])); + EXPECT_TRUE(a.count(&buf[3])); +}