[ADT] Add a new data structure for managing a priority worklist where re-insertion of entries into the worklist moves them to the end.

[chandlerc]

Summary

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.

[sanjoy]

Drop by comments.

include/llvm/ADT/PriorityWorklist.h

28

Don't we have autobrief these days?

40

Can we re-use DenseMapInfo<T> for this (say the tombstone value)?

88

Why T in one place and value_type in another?

95

We should have a way to fail an assert if someone tries to insert the sentinel "null" value.

177

This is only to avoid templates and lamda, right? If so, would be nice to state that explicitly.

[MatzeB]

I'm wondering if it is worth creating a whole new templated type here, do you expect more than one user?

include/llvm/ADT/PriorityWorklist.h

110–114

Should this rather be V.back() == T()? Or alternatively you could document that T needs to be implicitely convertible to bool with everything but the NULL value being false...

(Similar in the assert and another instance later).

[chandlerc]

In D21866#470712, @MatzeB wrote:

I'm wondering if it is worth creating a whole new templated type here, do you expect more than one user?

If the new template were complex, I wouldn't have. But this is actually easier than me getting this right in the one user I have in mind, and that user already needs two of them (one for SCC* and the other for RefSCC*).

Working on the rest of the comments now....

[chandlerc]

include/llvm/ADT/PriorityWorklist.h

28

I just forgot. Fixed.

40

We could I guess, but I expect these to essentially always be pointers, and the dense map empty and tombstones for those are going to be less efficient than testing for zero... I mean, probably it doesn't matter, but T() more accurately matches the naive code I would have written by hand and am packaging up here and so that's what I used.

If you or others feel strongly, I would use the empty key as these are in fact empty slots.

88

the value_type came from SetVector. I prefer the easier to read T personally, made it consistent.

95

Done.

110–114

Yes, good catch. Fixed here and elsewhere.

177

I guess, this was copied directly from SetVector. I can just delete all of this until someone needs it if you like (my one user doesn't need it).

[chandlerc]

Updates based on code review feedback.

[MatzeB]

LGTM

[sanjoy]

lgtm too

include/llvm/ADT/PriorityWorklist.h

42

ok

179

Given that there are tests, I'm fine with this.

[chandlerc]

Thanks all, submitting. Let me know if there are further suggestions around this.

include/llvm/ADT/PriorityWorklist.h

179

There was actually a bug here... I've improved the test and fixed the bug. =]

include/llvm/ADT/PriorityWorklist.h

+//===- 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 {
+
+/// \brief 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 value_type &X) {
+    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((bool)(back()) && "Cannot have a null element at the back!");
+    M.erase(back());
+    do {
+      V.pop_back();
+    } while (!V.empty() && !V.back());
+  }
+
+  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 value_type& 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());
+    } 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)
+      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.
+  template <typename UnaryPredicateT>
+  class TestAndEraseFromMap {
+    UnaryPredicateT P;
+    MapT &M;
+
+  public:
+    TestAndEraseFromMap(UnaryPredicateT P, MapT &M)
+        : P(std::move(P)), M(M) {}
+
+    template <typename ArgumentT>
+    bool operator()(const ArgumentT &Arg) {
+      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

unittests/ADT/CMakeLists.txt

set(ADTSources
   MapVectorTest.cpp
   OptionalTest.cpp
   PackedVectorTest.cpp
   PointerEmbeddedIntTest.cpp
   PointerIntPairTest.cpp
   PointerSumTypeTest.cpp
   PointerUnionTest.cpp
   PostOrderIteratorTest.cpp
+  PriorityWorklistTest.cpp
   RangeAdapterTest.cpp
   SCCIteratorTest.cpp
   SequenceTest.cpp
   SetVectorTest.cpp
   SmallPtrSetTest.cpp
   SmallStringTest.cpp
   SmallVectorTest.cpp
   SparseBitVectorTest.cpp

unittests/ADT/PriorityWorklistTest.cpp

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