Index: llvm/trunk/include/llvm/ADT/PriorityWorklist.h =================================================================== --- llvm/trunk/include/llvm/ADT/PriorityWorklist.h +++ llvm/trunk/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 +#include +#include +#include + +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 MapT = DenseMap> +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 + bool erase_if(UnaryPredicate P) { + typename VectorT::iterator E = std::remove_if( + V.begin(), V.end(), TestAndEraseFromMap(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 + 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 +class SmallPriorityWorklist + : public PriorityWorklist, + SmallDenseMap> { +public: + SmallPriorityWorklist() {} +}; + +} + +#endif Index: llvm/trunk/unittests/ADT/CMakeLists.txt =================================================================== --- llvm/trunk/unittests/ADT/CMakeLists.txt +++ llvm/trunk/unittests/ADT/CMakeLists.txt @@ -30,6 +30,7 @@ PointerSumTypeTest.cpp PointerUnionTest.cpp PostOrderIteratorTest.cpp + PriorityWorklistTest.cpp RangeAdapterTest.cpp SCCIteratorTest.cpp SequenceTest.cpp Index: llvm/trunk/unittests/ADT/PriorityWorklistTest.cpp =================================================================== --- llvm/trunk/unittests/ADT/PriorityWorklistTest.cpp +++ llvm/trunk/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 class PriorityWorklistTest : public ::testing::Test {}; +typedef ::testing::Types, SmallPriorityWorklist> + 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()); +} + +}