diff --git a/llvm/include/llvm/ADT/SmallPtrSet.h b/llvm/include/llvm/ADT/SmallPtrSet.h --- a/llvm/include/llvm/ADT/SmallPtrSet.h +++ b/llvm/include/llvm/ADT/SmallPtrSet.h @@ -409,6 +409,32 @@ } }; +/// Equality comparison for SmallPtrSet. +/// +/// Iterates over elements of LHS confirming that each value from LHS is also in +/// RHS, and that no additional values are in RHS. +template +bool operator==(const SmallPtrSetImpl &LHS, + const SmallPtrSetImpl &RHS) { + if (LHS.size() != RHS.size()) + return false; + + for (const auto *KV : LHS) + if (!RHS.count(KV)) + return false; + + return true; +} + +/// Inequality comparison for SmallPtrSet. +/// +/// Equivalent to !(LHS == RHS). +template +bool operator!=(const SmallPtrSetImpl &LHS, + const SmallPtrSetImpl &RHS) { + return !(LHS == RHS); +} + /// SmallPtrSet - This class implements a set which is optimized 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 diff --git a/llvm/unittests/ADT/SmallPtrSetTest.cpp b/llvm/unittests/ADT/SmallPtrSetTest.cpp --- a/llvm/unittests/ADT/SmallPtrSetTest.cpp +++ b/llvm/unittests/ADT/SmallPtrSetTest.cpp @@ -329,3 +329,41 @@ EXPECT_EQ(IntSet.count(Pair), 1u); EXPECT_NE(IntSet.find(Pair), IntSet.end()); } + +// Test equality comparison. +TEST(SmallPtrSetTest, EqualityComparison) { + int buf[3]; + for (int i = 0; i < 3; ++i) + buf[i] = 0; + + SmallPtrSet a; + a.insert(&buf[0]); + a.insert(&buf[1]); + + SmallPtrSet b; + b.insert(&buf[1]); + b.insert(&buf[0]); + + SmallPtrSet c; + c.insert(&buf[1]); + c.insert(&buf[2]); + + SmallPtrSet d; + d.insert(&buf[0]); + + SmallPtrSet e; + e.insert(&buf[0]); + e.insert(&buf[1]); + e.insert(&buf[2]); + + EXPECT_EQ(a, b); + EXPECT_EQ(b, a); + EXPECT_NE(b, c); + EXPECT_NE(c, a); + EXPECT_NE(d, a); + EXPECT_NE(a, d); + EXPECT_NE(a, e); + EXPECT_NE(e, a); + EXPECT_NE(c, e); + EXPECT_NE(e, d); +}