[llvm] r334887 - [SmallSet] Add SmallSetIterator.

[Florian Hahn via llvm-commits]

Author: fhahn
Date: Sat Jun 16 05:36:19 2018
New Revision: 334887

URL: http://llvm.org/viewvc/llvm-project?rev=334887&view=rev
Log:
[SmallSet] Add SmallSetIterator.

This patch adds a simple const_iterator implementation for SmallSet by
delegating to either a SmallVector::const_iterator or
std::set::const_iterator, depending on which storage is used by the
SmallSet.

Reviewers: dblaikie, craig.topper

Reviewed By: dblaikie

Differential Revision: https://reviews.llvm.org/D47942

Modified:
    llvm/trunk/include/llvm/ADT/SmallSet.h
    llvm/trunk/unittests/ADT/SmallSetTest.cpp

Modified: llvm/trunk/include/llvm/ADT/SmallSet.h
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/include/llvm/ADT/SmallSet.h?rev=334887&r1=334886&r2=334887&view=diff
==============================================================================
--- llvm/trunk/include/llvm/ADT/SmallSet.h (original)
+++ llvm/trunk/include/llvm/ADT/SmallSet.h Sat Jun 16 05:36:19 2018
@@ -17,14 +17,70 @@
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/iterator.h"
 #include "llvm/Support/Compiler.h"
 #include <cstddef>
 #include <functional>
 #include <set>
+#include <type_traits>
 #include <utility>
 
 namespace llvm {
 
+/// SmallSetIterator - This class implements a const_iterator for SmallSet by
+/// delegating to the underlying SmallVector or Set iterators.
+template <typename T, unsigned N, typename C>
+class SmallSetIterator
+    : public iterator_facade_base<SmallSetIterator<T, N, C>,
+                                  std::forward_iterator_tag, T> {
+private:
+  using SetIterTy = typename std::set<T, C>::const_iterator;
+  using VecIterTy = typename SmallVector<T, N>::const_iterator;
+  using SelfTy = SmallSetIterator<T, N, C>;
+
+  /// Iterators to the parts of the SmallSet containing the data. They are set
+  /// depending on isSmall.
+  union {
+    SetIterTy SetIter;
+    VecIterTy VecIter;
+  };
+
+  bool isSmall;
+
+public:
+  SmallSetIterator(SetIterTy SetIter) : SetIter(SetIter), isSmall(false) {
+    // Use static_assert here, as the SmallSetIterator type is incomplete in the
+    // class scope.
+    static_assert(std::is_trivially_destructible<SelfTy>::value &&
+                      std::is_trivially_copy_constructible<SelfTy>::value &&
+                      std::is_trivially_copy_assignable<SelfTy>::value &&
+                      std::is_move_constructible<SelfTy>::value &&
+                      std::is_trivially_move_assignable<SelfTy>::value,
+                  "SelfTy needs to be trivial copy and move constructors and "
+                  "assignments");
+  }
+
+  SmallSetIterator(VecIterTy VecIter) : VecIter(VecIter), isSmall(true) {}
+
+  bool operator==(const SmallSetIterator &RHS) const {
+    if (isSmall != RHS.isSmall)
+      return false;
+    if (isSmall)
+      return VecIter == RHS.VecIter;
+    return SetIter == RHS.SetIter;
+  }
+
+  SmallSetIterator &operator++() { // Preincrement
+    if (isSmall)
+      VecIter++;
+    else
+      SetIter++;
+    return *this;
+  }
+
+  const T &operator*() const { return isSmall ? *VecIter : *SetIter; }
+};
+
 /// SmallSet - This maintains a set of unique values, optimizing for the case
 /// when the set is small (less than N).  In this case, the set can be
 /// maintained with no mallocs.  If the set gets large, we expand to using an
@@ -50,6 +106,7 @@ class SmallSet {
 
 public:
   using size_type = size_t;
+  using const_iterator = SmallSetIterator<T, N, C>;
 
   SmallSet() = default;
 
@@ -121,6 +178,18 @@ public:
     Set.clear();
   }
 
+  const_iterator begin() const {
+    if (isSmall())
+      return {Vector.begin()};
+    return {Set.begin()};
+  }
+
+  const_iterator end() const {
+    if (isSmall())
+      return {Vector.end()};
+    return {Set.end()};
+  }
+
 private:
   bool isSmall() const { return Set.empty(); }
 

Modified: llvm/trunk/unittests/ADT/SmallSetTest.cpp
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/unittests/ADT/SmallSetTest.cpp?rev=334887&r1=334886&r2=334887&view=diff
==============================================================================
--- llvm/trunk/unittests/ADT/SmallSetTest.cpp (original)
+++ llvm/trunk/unittests/ADT/SmallSetTest.cpp Sat Jun 16 05:36:19 2018
@@ -13,6 +13,7 @@
 
 #include "llvm/ADT/SmallSet.h"
 #include "gtest/gtest.h"
+#include <string>
 
 using namespace llvm;
 
@@ -68,3 +69,57 @@ TEST(SmallSetTest, Erase) {
 
   EXPECT_EQ(0u, s1.count(8));
 }
+
+TEST(SmallSetTest, IteratorInt) {
+  SmallSet<int, 4> s1;
+
+  // Test the 'small' case.
+  for (int i = 0; i < 3; i++)
+    s1.insert(i);
+
+  std::vector<int> V(s1.begin(), s1.end());
+  // Make sure the elements are in the expected order.
+  std::sort(V.begin(), V.end());
+  for (int i = 0; i < 3; i++)
+    EXPECT_EQ(i, V[i]);
+
+  // Test the 'big' case by adding a few more elements to switch to std::set
+  // internally.
+  for (int i = 3; i < 6; i++)
+    s1.insert(i);
+
+  V.assign(s1.begin(), s1.end());
+  // Make sure the elements are in the expected order.
+  std::sort(V.begin(), V.end());
+  for (int i = 0; i < 6; i++)
+    EXPECT_EQ(i, V[i]);
+}
+
+TEST(SmallSetTest, IteratorString) {
+  // Test SmallSetIterator for SmallSet with a type with non-trivial
+  // ctors/dtors.
+  SmallSet<std::string, 2> s1;
+
+  s1.insert("str 1");
+  s1.insert("str 2");
+  s1.insert("str 1");
+
+  std::vector<std::string> V(s1.begin(), s1.end());
+  std::sort(V.begin(), V.end());
+  EXPECT_EQ(2u, s1.size());
+  EXPECT_EQ("str 1", V[0]);
+  EXPECT_EQ("str 2", V[1]);
+
+  s1.insert("str 4");
+  s1.insert("str 0");
+  s1.insert("str 4");
+
+  V.assign(s1.begin(), s1.end());
+  // Make sure the elements are in the expected order.
+  std::sort(V.begin(), V.end());
+  EXPECT_EQ(4u, s1.size());
+  EXPECT_EQ("str 0", V[0]);
+  EXPECT_EQ("str 1", V[1]);
+  EXPECT_EQ("str 2", V[2]);
+  EXPECT_EQ("str 4", V[3]);
+}