From 8719faafdb3f2eb4971201a9286e24e227cc9d8f Mon Sep 17 00:00:00 2001
From: Zijia Zhu <piggynl@outlook.com>
Date: Mon, 15 Aug 2022 13:43:59 +0800
Subject: [PATCH] [ADT] Make SmallSet::insert(const T &) return const_iterator

This patch makes `SmallSet::insert(const T &)` return
`std::pair<const_iterator, bool>` instead of
`std::pair<NoneType, bool>`. This will exactly match std::set's behavior
and make deduplicating items with SmallSet easier.

Reviewed By: dblaikie, lattner

Differential Revision: https://reviews.llvm.org/D131549
---
 llvm/include/llvm/ADT/SmallSet.h    | 25 ++++++++++++-------------
 llvm/unittests/ADT/SmallSetTest.cpp | 27 +++++++++++++++++++++------
 2 files changed, 33 insertions(+), 19 deletions(-)

diff --git a/llvm/include/llvm/ADT/SmallSet.h b/llvm/include/llvm/ADT/SmallSet.h
index 99f7c55ab8fd..d64620397d3d 100644
--- a/llvm/include/llvm/ADT/SmallSet.h
+++ b/llvm/include/llvm/ADT/SmallSet.h
@@ -141,6 +141,7 @@ class SmallSet {
   std::set<T, C> Set;
 
   using VIterator = typename SmallVector<T, N>::const_iterator;
+  using SIterator = typename std::set<T, C>::const_iterator;
   using mutable_iterator = typename SmallVector<T, N>::iterator;
 
   // In small mode SmallPtrSet uses linear search for the elements, so it is
@@ -171,22 +172,21 @@ public:
   }
 
   /// insert - Insert an element into the set if it isn't already there.
-  /// Returns true if the element is inserted (it was not in the set before).
-  /// The first value of the returned pair is unused and provided for
-  /// partial compatibility with the standard library self-associative container
-  /// concept.
-  // FIXME: Add iterators that abstract over the small and large form, and then
-  // return those here.
-  std::pair<NoneType, bool> insert(const T &V) {
-    if (!isSmall())
-      return std::make_pair(None, Set.insert(V).second);
+  /// 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);
+    }
 
     VIterator I = vfind(V);
     if (I != Vector.end())    // Don't reinsert if it already exists.
-      return std::make_pair(None, false);
+      return std::make_pair(const_iterator(I), false);
     if (Vector.size() < N) {
       Vector.push_back(V);
-      return std::make_pair(None, true);
+      return std::make_pair(const_iterator(std::prev(Vector.end())), true);
     }
 
     // Otherwise, grow from vector to set.
@@ -194,8 +194,7 @@ public:
       Set.insert(Vector.back());
       Vector.pop_back();
     }
-    Set.insert(V);
-    return std::make_pair(None, true);
+    return std::make_pair(const_iterator(Set.insert(V).first), true);
   }
 
   template <typename IterT>
diff --git a/llvm/unittests/ADT/SmallSetTest.cpp b/llvm/unittests/ADT/SmallSetTest.cpp
index 0b7400942e3c..b50b368ae663 100644
--- a/llvm/unittests/ADT/SmallSetTest.cpp
+++ b/llvm/unittests/ADT/SmallSetTest.cpp
@@ -21,11 +21,17 @@ TEST(SmallSetTest, Insert) {
 
   SmallSet<int, 4> s1;
 
-  for (int i = 0; i < 4; i++)
-    s1.insert(i);
+  for (int i = 0; i < 4; i++) {
+    auto InsertResult = s1.insert(i);
+    EXPECT_EQ(*InsertResult.first, i);
+    EXPECT_EQ(InsertResult.second, true);
+  }
 
-  for (int i = 0; i < 4; i++)
-    s1.insert(i);
+  for (int i = 0; i < 4; i++) {
+    auto InsertResult = s1.insert(i);
+    EXPECT_EQ(*InsertResult.first, i);
+    EXPECT_EQ(InsertResult.second, false);
+  }
 
   EXPECT_EQ(4u, s1.size());
 
@@ -38,8 +44,17 @@ TEST(SmallSetTest, Insert) {
 TEST(SmallSetTest, Grow) {
   SmallSet<int, 4> s1;
 
-  for (int i = 0; i < 8; i++)
-    s1.insert(i);
+  for (int i = 0; i < 8; i++) {
+    auto InsertResult = s1.insert(i);
+    EXPECT_EQ(*InsertResult.first, i);
+    EXPECT_EQ(InsertResult.second, true);
+  }
+
+  for (int i = 0; i < 8; i++) {
+    auto InsertResult = s1.insert(i);
+    EXPECT_EQ(*InsertResult.first, i);
+    EXPECT_EQ(InsertResult.second, false);
+  }
 
   EXPECT_EQ(8u, s1.size());