diff --git a/llvm/include/llvm/ADT/SmallSet.h b/llvm/include/llvm/ADT/SmallSet.h index 8d7511bf0bc8d9..56259ea7cf9d0f 100644 --- a/llvm/include/llvm/ADT/SmallSet.h +++ b/llvm/include/llvm/ADT/SmallSet.h @@ -161,26 +161,10 @@ class SmallSet { /// Returns a pair. The first value of it is an iterator to the inserted /// element or the existing element in the set. The second value is true /// if the element is inserted (it was not in the set before). - std::pair insert(const T &V) { - if (!isSmall()) { - auto [I, Inserted] = Set.insert(V); - return std::make_pair(const_iterator(I), Inserted); - } - - auto I = std::find(Vector.begin(), Vector.end(), V); - if (I != Vector.end()) // Don't reinsert if it already exists. - return std::make_pair(const_iterator(I), false); - if (Vector.size() < N) { - Vector.push_back(V); - return std::make_pair(const_iterator(std::prev(Vector.end())), true); - } + std::pair insert(const T &V) { return insertImpl(V); } - // Otherwise, grow from vector to set. - while (!Vector.empty()) { - Set.insert(Vector.back()); - Vector.pop_back(); - } - return std::make_pair(const_iterator(Set.insert(V).first), true); + std::pair insert(T &&V) { + return insertImpl(std::move(V)); } template @@ -226,6 +210,30 @@ class SmallSet { private: bool isSmall() const { return Set.empty(); } + + template + std::pair insertImpl(ArgType &&V) { + static_assert(std::is_convertible_v, + "ArgType must be convertible to T!"); + if (!isSmall()) { + auto [I, Inserted] = Set.insert(std::forward(V)); + return {const_iterator(I), Inserted}; + } + + auto I = std::find(Vector.begin(), Vector.end(), V); + if (I != Vector.end()) // Don't reinsert if it already exists. + return {const_iterator(I), false}; + if (Vector.size() < N) { + Vector.push_back(std::forward(V)); + return {const_iterator(std::prev(Vector.end())), true}; + } + + // Otherwise, grow from vector to set. + Set.insert(std::make_move_iterator(Vector.begin()), + std::make_move_iterator(Vector.end())); + Vector.clear(); + return {const_iterator(Set.insert(std::forward(V)).first), true}; + } }; /// If this set is of pointer values, transparently switch over to using diff --git a/llvm/unittests/ADT/SmallSetTest.cpp b/llvm/unittests/ADT/SmallSetTest.cpp index b50b368ae66361..0fb20b19df9254 100644 --- a/llvm/unittests/ADT/SmallSetTest.cpp +++ b/llvm/unittests/ADT/SmallSetTest.cpp @@ -41,6 +41,40 @@ TEST(SmallSetTest, Insert) { EXPECT_EQ(0u, s1.count(4)); } +TEST(SmallSetTest, InsertPerfectFwd) { + struct Value { + int Key; + bool Moved; + + Value(int Key) : Key(Key), Moved(false) {} + Value(const Value &) = default; + Value(Value &&Other) : Key(Other.Key), Moved(false) { Other.Moved = true; } + bool operator==(const Value &Other) const { return Key == Other.Key; } + bool operator<(const Value &Other) const { return Key < Other.Key; } + }; + + { + SmallSet S; + Value V1(1), V2(2); + + S.insert(V1); + EXPECT_EQ(V1.Moved, false); + + S.insert(std::move(V2)); + EXPECT_EQ(V2.Moved, true); + } + { + SmallSet S; + Value V1(1), V2(2); + + S.insert(V1); + EXPECT_EQ(V1.Moved, false); + + S.insert(std::move(V2)); + EXPECT_EQ(V2.Moved, true); + } +} + TEST(SmallSetTest, Grow) { SmallSet s1;