//===- PriorityWorklist.h - Worklist with insertion priority ----*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
///
/// \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/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Compiler.h"
#include <cassert>
#include <cstddef>
#include <iterator>
#include <type_traits>
#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:
using value_type = T;
using key_type = T;
using reference = T&;
using const_reference = const T&;
using size_type = typename MapT::size_type;
/// Construct an empty PriorityWorklist
PriorityWorklist() = default;
/// 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;
}
/// Insert a sequence of new elements into the PriorityWorklist.
template <typename SequenceT>
std::enable_if_t<!std::is_convertible<SequenceT, T>::value>
insert(SequenceT &&Input) {
if (std::begin(Input) == std::end(Input))
// Nothing to do for an empty input sequence.
return;
// First pull the input sequence into the vector as a bulk append
// operation.
ptrdiff_t StartIndex = V.size();
V.insert(V.end(), std::begin(Input), std::end(Input));
// Now walk backwards fixing up the index map and deleting any duplicates.
for (ptrdiff_t i = V.size() - 1; i >= StartIndex; --i) {
auto InsertResult = M.insert({V[i], i});
if (InsertResult.second)
continue;
// If the existing index is before this insert's start, nuke that one and
// move it up.
ptrdiff_t &Index = InsertResult.first->second;
if (Index < StartIndex) {
V[Index] = T();
Index = i;
continue;
}
// Otherwise the existing one comes first so just clear out the value in
// this slot.
V[i] = T();
}
}
/// 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());
}
[[nodiscard]] T 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(remove_if(V, 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 =
remove_if(V, 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;
}
/// Reverse the items in the PriorityWorklist.
///
/// This does an in-place reversal. Other kinds of reverse aren't easy to
/// support in the face of the worklist semantics.
/// 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() = default;
};
} // end namespace llvm
#endif // LLVM_ADT_PRIORITYWORKLIST_H