From 2a29d24ba94dac82e838c694595a0a574e505aab Mon Sep 17 00:00:00 2001
From: Victor Campos <victor.campos@arm.com>
Date: Wed, 25 Sep 2024 12:43:41 +0100
Subject: [PATCH] [ADT] Use perfect forwarding in SmallSet::insert() (#108590)

Previously this method took arguments by const-ref. This patch changes
the implementation to take perfectly forwarded arguments in the form of
a universal reference. Now, the insertion method will take advantage of
arguments passed as rvalue, potentially leading to performance
improvements.
---
 llvm/include/llvm/ADT/SmallSet.h    | 46 +++++++++++++++++------------
 llvm/unittests/ADT/SmallSetTest.cpp | 34 +++++++++++++++++++++
 2 files changed, 61 insertions(+), 19 deletions(-)

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<const_iterator, bool> 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<const_iterator, bool> 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<const_iterator, bool> insert(T &&V) {
+    return insertImpl(std::move(V));
   }
 
   template <typename IterT>
@@ -226,6 +210,30 @@ class SmallSet {
 
 private:
   bool isSmall() const { return Set.empty(); }
+
+  template <typename ArgType>
+  std::pair<const_iterator, bool> insertImpl(ArgType &&V) {
+    static_assert(std::is_convertible_v<ArgType, T>,
+                  "ArgType must be convertible to T!");
+    if (!isSmall()) {
+      auto [I, Inserted] = Set.insert(std::forward<ArgType>(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<ArgType>(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<ArgType>(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<Value, 4> 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<Value, 1> 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<int, 4> s1;