--- /dev/null
+//===- 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
--- /dev/null
+//===- 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());
+}
+
+}