Recommit r334887: [SmallSet] Add SmallSetIterator.

Updated to make sure we properly construct/destroy SetIter if it has a
non-trivial ctors/dtors, like in MSVC.

llvm-svn: 337818

diff --git a/llvm/include/llvm/ADT/SmallSet.h b/llvm/include/llvm/ADT/SmallSet.h
index d52d0f0..5d84627 100644 (file)
--- a/llvm/include/llvm/ADT/SmallSet.h
+++ b/llvm/include/llvm/ADT/SmallSet.h
@@ -17,21 +17,120 @@
 #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 "llvm/Support/type_traits.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) {}
+
+  SmallSetIterator(VecIterTy VecIter) : VecIter(VecIter), isSmall(true) {}
+
+  // Spell out destructor, copy/move constructor and assignment operators for
+  // MSVC STL, where set<T>::const_iterator is not trivially copy constructible.
+  ~SmallSetIterator() {
+    if (isSmall)
+      VecIter.~VecIterTy();
+    else
+      SetIter.~SetIterTy();
+  }
+
+  SmallSetIterator(const SmallSetIterator &Other) : isSmall(Other.isSmall) {
+    if (isSmall)
+      VecIter = Other.VecIter;
+    else
+      // Use placement new, to make sure SetIter is properly constructed, even
+      // if it is not trivially copy-able (e.g. in MSVC).
+      new (&SetIter) SetIterTy(Other.SetIter);
+  }
+
+  SmallSetIterator(SmallSetIterator &&Other) : isSmall(Other.isSmall) {
+    if (isSmall)
+      VecIter = std::move(Other.VecIter);
+    else
+      // Use placement new, to make sure SetIter is properly constructed, even
+      // if it is not trivially copy-able (e.g. in MSVC).
+      new (&SetIter) SetIterTy(std::move(Other.SetIter));
+  }
+
+  SmallSetIterator& operator=(const SmallSetIterator& Other) {
+    // Call destructor for SetIter, so it gets properly destroyed if it is
+    // not trivially destructible in case we are setting VecIter.
+    if (!isSmall)
+      SetIter.~SetIterTy();
+
+    isSmall = Other.isSmall;
+    if (isSmall)
+      VecIter = Other.VecIter;
+    else
+      new (&SetIter) SetIterTy(Other.SetIter);
+    return *this;
+  }
+
+  SmallSetIterator& operator=(SmallSetIterator&& Other) {
+    // Call destructor for SetIter, so it gets properly destroyed if it is
+    // not trivially destructible in case we are setting VecIter.
+    if (!isSmall)
+      SetIter.~SetIterTy();
+
+    isSmall = Other.isSmall;
+    if (isSmall)
+      VecIter = std::move(Other.VecIter);
+    else
+      new (&SetIter) SetIterTy(std::move(Other.SetIter));
+    return *this;
+  }
+
+  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
 /// std::set to maintain reasonable lookup times.
-///
-/// Note that this set does not provide a way to iterate over members in the
-/// set.
 template <typename T, unsigned N, typename C = std::less<T>>
 class SmallSet {
   /// Use a SmallVector to hold the elements here (even though it will never
@@ -50,6 +149,7 @@ class SmallSet {
 
 public:
   using size_type = size_t;
+  using const_iterator = SmallSetIterator<T, N, C>;
 
   SmallSet() = default;
 
@@ -121,6 +221,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(); }
 

diff --git a/llvm/unittests/ADT/SmallSetTest.cpp b/llvm/unittests/ADT/SmallSetTest.cpp
index 7608e65..d78a72b 100644 (file)
--- a/llvm/unittests/ADT/SmallSetTest.cpp
+++ b/llvm/unittests/ADT/SmallSetTest.cpp
@@ -13,6 +13,7 @@
 
 #include "llvm/ADT/SmallSet.h"
 #include "gtest/gtest.h"
+#include <string>
 
 using namespace llvm;
 
@@ -68,3 +69,81 @@ 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]);
+}
+
+TEST(SmallSetTest, IteratorIncMoveCopy) {
+  // 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");
+
+  auto Iter = s1.begin();
+  EXPECT_EQ("str 1", *Iter);
+  ++Iter;
+  EXPECT_EQ("str 2", *Iter);
+
+  s1.insert("str 4");
+  s1.insert("str 0");
+  auto Iter2 = s1.begin();
+  Iter = std::move(Iter2);
+  EXPECT_EQ("str 0", *Iter);
+
+  auto Iter3 = s1.end();
+  Iter3 = Iter2;
+  EXPECT_EQ(Iter3, Iter2);
+}