[ADT] Add a new data structure for managing a priority worklist where

re-insertion of entries into the worklist moves them to the end.

This is fairly similar to a SetVector, but helps in the case where in
addition to not inserting duplicates you want to adjust the sequence of
a pop-off-the-back worklist.

I'm not at all attached to the name of this data structure if others
have better suggestions, but this is one that David Majnemer brought up
in IRC discussions that seems plausible.

I've trimmed the interface down somewhat from SetVector's interface
because several things make less sense here IMO: iteration primarily.
I'd prefer to add these back as we have users that need them. My use
case doesn't even need all of what is provided here. =]

I've also included a basic unittest to make sure this functions
reasonably.

Differential Revision: http://reviews.llvm.org/D21866

llvm-svn: 274198

llvm/include/llvm/ADT/PriorityWorklist.h

@@ -0,0 +1,224 @@
+//===- PriorityWorklist.h - Worklist with insertion priority ----*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+///
+/// This file provides a priority worklist. See the class comments for details.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_ADT_PRIORITYWORKLIST_H
+#define LLVM_ADT_PRIORITYWORKLIST_H
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include <algorithm>
+#include <cassert>
+#include <utility>
+#include <vector>
+
+namespace llvm {
+
+/// A FILO worklist that prioritizes on re-insertion without duplication.
+///
+/// This is very similar to a \c SetVector with the primary difference that
+/// while re-insertion does not create a duplicate, it does adjust the
+/// visitation order to respect the last insertion point. This can be useful
+/// when the visit order needs to be prioritized based on insertion point
+/// without actually having duplicate visits.
+///
+/// Note that this doesn't prevent re-insertion of elements which have been
+/// visited -- if you need to break cycles, a set will still be necessary.
+///
+/// The type \c T must be default constructable to a null value that will be
+/// ignored. It is an error to insert such a value, and popping elements will
+/// never produce such a value. It is expected to be used with common nullable
+/// types like pointers or optionals.
+///
+/// Internally this uses a vector to store the worklist and a map to identify
+/// existing elements in the worklist. Both of these may be customized, but the
+/// map must support the basic DenseMap API for mapping from a T to an integer
+/// index into the vector.
+///
+/// A partial specialization is provided to automatically select a SmallVector
+/// and a SmallDenseMap if custom data structures are not provided.
+template <typename T, typename VectorT = std::vector<T>,
+          typename MapT = DenseMap<T, ptrdiff_t>>
+class PriorityWorklist {
+public:
+  typedef T value_type;
+  typedef T key_type;
+  typedef T& reference;
+  typedef const T& const_reference;
+  typedef typename MapT::size_type size_type;
+
+  /// Construct an empty PriorityWorklist
+  PriorityWorklist() {}
+
+  /// Determine if the PriorityWorklist is empty or not.
+  bool empty() const {
+    return V.empty();
+  }
+
+  /// Returns the number of elements in the worklist.
+  size_type size() const {
+    return M.size();
+  }
+
+  /// Count the number of elements of a given key in the PriorityWorklist.
+  /// \returns 0 if the element is not in the PriorityWorklist, 1 if it is.
+  size_type count(const key_type &key) const {
+    return M.count(key);
+  }
+
+  /// Return the last element of the PriorityWorklist.
+  const T &back() const {
+    assert(!empty() && "Cannot call back() on empty PriorityWorklist!");
+    return V.back();
+  }
+
+  /// Insert a new element into the PriorityWorklist.
+  /// \returns true if the element was inserted into the PriorityWorklist.
+  bool insert(const T &X) {
+    assert(X != T() && "Cannot insert a null (default constructed) value!");
+    auto InsertResult = M.insert({X, V.size()});
+    if (InsertResult.second) {
+      // Fresh value, just append it to the vector.
+      V.push_back(X);
+      return true;
+    }
+
+    auto &Index = InsertResult.first->second;
+    assert(V[Index] == X && "Value not actually at index in map!");
+    if (Index != (ptrdiff_t)(V.size() - 1)) {
+      // If the element isn't at the back, null it out and append a fresh one.
+      V[Index] = T();
+      Index = (ptrdiff_t)V.size();
+      V.push_back(X);
+    }
+    return false;
+  }
+
+  /// Remove the last element of the PriorityWorklist.
+  void pop_back() {
+    assert(!empty() && "Cannot remove an element when empty!");
+    assert(back() != T() && "Cannot have a null element at the back!");
+    M.erase(back());
+    do {
+      V.pop_back();
+    } while (!V.empty() && V.back() == T());
+  }
+
+  T LLVM_ATTRIBUTE_UNUSED_RESULT pop_back_val() {
+    T Ret = back();
+    pop_back();
+    return Ret;
+  }
+
+  /// Erase an item from the worklist.
+  ///
+  /// Note that this is constant time due to the nature of the worklist implementation.
+  bool erase(const T& X) {
+    auto I = M.find(X);
+    if (I == M.end())
+      return false;
+
+    assert(V[I->second] == X && "Value not actually at index in map!");
+    if (I->second == (ptrdiff_t)(V.size() - 1)) {
+      do {
+        V.pop_back();
+      } while (!V.empty() && V.back() == T());
+    } else {
+      V[I->second] = T();
+    }
+    M.erase(I);
+    return true;
+  }
+
+  /// Erase items from the set vector based on a predicate function.
+  ///
+  /// This is intended to be equivalent to the following code, if we could
+  /// write it:
+  ///
+  /// \code
+  ///   V.erase(std::remove_if(V.begin(), V.end(), P), V.end());
+  /// \endcode
+  ///
+  /// However, PriorityWorklist doesn't expose non-const iterators, making any
+  /// algorithm like remove_if impossible to use.
+  ///
+  /// \returns true if any element is removed.
+  template <typename UnaryPredicate>
+  bool erase_if(UnaryPredicate P) {
+    typename VectorT::iterator E = std::remove_if(
+        V.begin(), V.end(), TestAndEraseFromMap<UnaryPredicate>(P, M));
+    if (E == V.end())
+      return false;
+    for (auto I = V.begin(); I != E; ++I)
+      if (*I != T())
+        M[*I] = I - V.begin();
+    V.erase(E, V.end());
+    return true;
+  }
+
+  /// Completely clear the PriorityWorklist
+  void clear() {
+    M.clear();
+    V.clear();
+  }
+
+private:
+  /// A wrapper predicate designed for use with std::remove_if.
+  ///
+  /// This predicate wraps a predicate suitable for use with std::remove_if to
+  /// call M.erase(x) on each element which is slated for removal. This just
+  /// allows the predicate to be move only which we can't do with lambdas
+  /// today.
+  template <typename UnaryPredicateT>
+  class TestAndEraseFromMap {
+    UnaryPredicateT P;
+    MapT &M;
+
+  public:
+    TestAndEraseFromMap(UnaryPredicateT P, MapT &M)
+        : P(std::move(P)), M(M) {}
+
+    bool operator()(const T &Arg) {
+      if (Arg == T())
+        // Skip null values in the PriorityWorklist.
+        return false;
+
+      if (P(Arg)) {
+        M.erase(Arg);
+        return true;
+      }
+      return false;
+    }
+  };
+
+  /// The map from value to index in the vector.
+  MapT M;
+
+  /// The vector of elements in insertion order.
+  VectorT V;
+};
+
+/// A version of \c PriorityWorklist that selects small size optimized data
+/// structures for the vector and map.
+template <typename T, unsigned N>
+class SmallPriorityWorklist
+    : public PriorityWorklist<T, SmallVector<T, N>,
+                              SmallDenseMap<T, ptrdiff_t>> {
+public:
+  SmallPriorityWorklist() {}
+};
+
+}
+
+#endif

llvm/unittests/ADT/CMakeLists.txt

@@ -30,6 +30,7 @@ set(ADTSources
   PointerSumTypeTest.cpp
   PointerUnionTest.cpp
   PostOrderIteratorTest.cpp
+  PriorityWorklistTest.cpp
   RangeAdapterTest.cpp
   SCCIteratorTest.cpp
   SequenceTest.cpp

llvm/unittests/ADT/PriorityWorklistTest.cpp

@@ -0,0 +1,106 @@
+//===- llvm/unittest/ADT/PriorityWorklist.cpp -----------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// PriorityWorklist unit tests.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/PriorityWorklist.h"
+#include "gtest/gtest.h"
+
+namespace {
+
+using namespace llvm;
+
+template <typename T> class PriorityWorklistTest : public ::testing::Test {};
+typedef ::testing::Types<PriorityWorklist<int>, SmallPriorityWorklist<int, 2>>
+    TestTypes;
+TYPED_TEST_CASE(PriorityWorklistTest, TestTypes);
+
+TYPED_TEST(PriorityWorklistTest, Basic) {
+  TypeParam W;
+  EXPECT_TRUE(W.empty());
+  EXPECT_EQ(0u, W.size());
+  EXPECT_FALSE(W.count(42));
+
+  EXPECT_TRUE(W.insert(21));
+  EXPECT_TRUE(W.insert(42));
+  EXPECT_TRUE(W.insert(17));
+
+  EXPECT_FALSE(W.empty());
+  EXPECT_EQ(3u, W.size());
+  EXPECT_TRUE(W.count(42));
+
+  EXPECT_FALSE(W.erase(75));
+  EXPECT_EQ(3u, W.size());
+  EXPECT_EQ(17, W.back());
+
+  EXPECT_TRUE(W.erase(17));
+  EXPECT_FALSE(W.count(17));
+  EXPECT_EQ(2u, W.size());
+  EXPECT_EQ(42, W.back());
+
+  W.clear();
+  EXPECT_TRUE(W.empty());
+  EXPECT_EQ(0u, W.size());
+
+  EXPECT_TRUE(W.insert(21));
+  EXPECT_TRUE(W.insert(42));
+  EXPECT_TRUE(W.insert(12));
+  EXPECT_TRUE(W.insert(17));
+  EXPECT_TRUE(W.count(12));
+  EXPECT_TRUE(W.count(17));
+  EXPECT_EQ(4u, W.size());
+  EXPECT_EQ(17, W.back());
+  EXPECT_TRUE(W.erase(12));
+  EXPECT_FALSE(W.count(12));
+  EXPECT_TRUE(W.count(17));
+  EXPECT_EQ(3u, W.size());
+  EXPECT_EQ(17, W.back());
+
+  EXPECT_FALSE(W.insert(42));
+  EXPECT_EQ(3u, W.size());
+  EXPECT_EQ(42, W.pop_back_val());
+  EXPECT_EQ(17, W.pop_back_val());
+  EXPECT_EQ(21, W.pop_back_val());
+  EXPECT_TRUE(W.empty());
+}
+
+TYPED_TEST(PriorityWorklistTest, EraseIf) {
+  TypeParam W;
+  W.insert(23);
+  W.insert(10);
+  W.insert(47);
+  W.insert(42);
+  W.insert(23);
+  W.insert(13);
+  W.insert(26);
+  W.insert(42);
+  EXPECT_EQ(6u, W.size());
+
+  EXPECT_FALSE(W.erase_if([](int i) { return i > 100; }));
+  EXPECT_EQ(6u, W.size());
+  EXPECT_EQ(42, W.back());
+
+  EXPECT_TRUE(W.erase_if([](int i) {
+    assert(i != 0 && "Saw a null value!");
+    return (i & 1) == 0;
+  }));
+  EXPECT_EQ(3u, W.size());
+  EXPECT_FALSE(W.count(42));
+  EXPECT_FALSE(W.count(26));
+  EXPECT_FALSE(W.count(10));
+  EXPECT_FALSE(W.insert(47));
+  EXPECT_FALSE(W.insert(23));
+  EXPECT_EQ(23, W.pop_back_val());
+  EXPECT_EQ(47, W.pop_back_val());
+  EXPECT_EQ(13, W.pop_back_val());
+}
+
+}