1//===- PriorityWorklist.h - Worklist with insertion priority ----*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8///
9/// \file
10///
11/// This file provides a priority worklist. See the class comments for details.
12///
13//===----------------------------------------------------------------------===//
14
15#ifndef LLVM_ADT_PRIORITYWORKLIST_H
16#define LLVM_ADT_PRIORITYWORKLIST_H
17
18#include "llvm/ADT/DenseMap.h"
19#include "llvm/ADT/STLExtras.h"
20#include "llvm/ADT/SmallVector.h"
21#include "llvm/Support/Compiler.h"
22#include <algorithm>
23#include <cassert>
24#include <cstddef>
25#include <iterator>
26#include <type_traits>
27#include <vector>
28
29namespace llvm {
30
31/// A FILO worklist that prioritizes on re-insertion without duplication.
32///
33/// This is very similar to a \c SetVector with the primary difference that
34/// while re-insertion does not create a duplicate, it does adjust the
35/// visitation order to respect the last insertion point. This can be useful
36/// when the visit order needs to be prioritized based on insertion point
37/// without actually having duplicate visits.
38///
39/// Note that this doesn't prevent re-insertion of elements which have been
40/// visited -- if you need to break cycles, a set will still be necessary.
41///
42/// The type \c T must be default constructable to a null value that will be
43/// ignored. It is an error to insert such a value, and popping elements will
44/// never produce such a value. It is expected to be used with common nullable
45/// types like pointers or optionals.
46///
47/// Internally this uses a vector to store the worklist and a map to identify
48/// existing elements in the worklist. Both of these may be customized, but the
49/// map must support the basic DenseMap API for mapping from a T to an integer
50/// index into the vector.
51///
52/// A partial specialization is provided to automatically select a SmallVector
53/// and a SmallDenseMap if custom data structures are not provided.
54template <typename T, typename VectorT = std::vector<T>,
55 typename MapT = DenseMap<T, ptrdiff_t>>
56class PriorityWorklist {
57public:
58 using value_type = T;
59 using key_type = T;
60 using reference = T&;
61 using const_reference = const T&;
62 using size_type = typename MapT::size_type;
63
64 /// Construct an empty PriorityWorklist
65 PriorityWorklist() = default;
66
67 /// Determine if the PriorityWorklist is empty or not.
68 bool empty() const {
69 return V.empty();
70 }
71
72 /// Returns the number of elements in the worklist.
73 size_type size() const {
74 return M.size();
75 }
76
77 /// Count the number of elements of a given key in the PriorityWorklist.
78 /// \returns 0 if the element is not in the PriorityWorklist, 1 if it is.
79 size_type count(const key_type &key) const {
80 return M.count(key);
81 }
82
83 /// Return the last element of the PriorityWorklist.
84 const T &back() const {
85 assert(!empty() && "Cannot call back() on empty PriorityWorklist!");
86 return V.back();
87 }
88
89 /// Insert a new element into the PriorityWorklist.
90 /// \returns true if the element was inserted into the PriorityWorklist.
91 bool insert(const T &X) {
92 assert(X != T() && "Cannot insert a null (default constructed) value!");
93 auto InsertResult = M.insert({X, V.size()});
94 if (InsertResult.second) {
95 // Fresh value, just append it to the vector.
96 V.push_back(X);
97 return true;
98 }
99
100 auto &Index = InsertResult.first->second;
101 assert(V[Index] == X && "Value not actually at index in map!");
102 if (Index != (ptrdiff_t)(V.size() - 1)) {
103 // If the element isn't at the back, null it out and append a fresh one.
104 V[Index] = T();
105 Index = (ptrdiff_t)V.size();
106 V.push_back(X);
107 }
108 return false;
109 }
110
111 /// Insert a sequence of new elements into the PriorityWorklist.
112 template <typename SequenceT>
113 std::enable_if_t<!std::is_convertible<SequenceT, T>::value>
114 insert(SequenceT &&Input) {
115 if (std::begin(Input) == std::end(Input))
116 // Nothing to do for an empty input sequence.
117 return;
118
119 // First pull the input sequence into the vector as a bulk append
120 // operation.
121 ptrdiff_t StartIndex = V.size();
122 V.insert(V.end(), std::begin(Input), std::end(Input));
123 // Now walk backwards fixing up the index map and deleting any duplicates.
124 for (ptrdiff_t i = V.size() - 1; i >= StartIndex; --i) {
125 auto InsertResult = M.insert({V[i], i});
126 if (InsertResult.second)
127 continue;
128
129 // If the existing index is before this insert's start, nuke that one and
130 // move it up.
131 ptrdiff_t &Index = InsertResult.first->second;
132 if (Index < StartIndex) {
133 V[Index] = T();
134 Index = i;
135 continue;
136 }
137
138 // Otherwise the existing one comes first so just clear out the value in
139 // this slot.
140 V[i] = T();
141 }
142 }
143
144 /// Remove the last element of the PriorityWorklist.
145 void pop_back() {
146 assert(!empty() && "Cannot remove an element when empty!");
147 assert(back() != T() && "Cannot have a null element at the back!");
148 M.erase(back());
149 do {
150 V.pop_back();
151 } while (!V.empty() && V.back() == T());
152 }
153
154 LLVM_NODISCARD T pop_back_val() {
155 T Ret = back();
156 pop_back();
157 return Ret;
158 }
159
160 /// Erase an item from the worklist.
161 ///
162 /// Note that this is constant time due to the nature of the worklist implementation.
163 bool erase(const T& X) {
164 auto I = M.find(X);
165 if (I == M.end())
166 return false;
167
168 assert(V[I->second] == X && "Value not actually at index in map!");
169 if (I->second == (ptrdiff_t)(V.size() - 1)) {
170 do {
171 V.pop_back();
172 } while (!V.empty() && V.back() == T());
173 } else {
174 V[I->second] = T();
175 }
176 M.erase(I);
177 return true;
178 }
179
180 /// Erase items from the set vector based on a predicate function.
181 ///
182 /// This is intended to be equivalent to the following code, if we could
183 /// write it:
184 ///
185 /// \code
186 /// V.erase(remove_if(V, P), V.end());
187 /// \endcode
188 ///
189 /// However, PriorityWorklist doesn't expose non-const iterators, making any
190 /// algorithm like remove_if impossible to use.
191 ///
192 /// \returns true if any element is removed.
193 template <typename UnaryPredicate>
194 bool erase_if(UnaryPredicate P) {
195 typename VectorT::iterator E =
196 remove_if(V, TestAndEraseFromMap<UnaryPredicate>(P, M));
197 if (E == V.end())
198 return false;
199 for (auto I = V.begin(); I != E; ++I)
200 if (*I != T())
201 M[*I] = I - V.begin();
202 V.erase(E, V.end());
203 return true;
204 }
205
206 /// Reverse the items in the PriorityWorklist.
207 ///
208 /// This does an in-place reversal. Other kinds of reverse aren't easy to
209 /// support in the face of the worklist semantics.
210
211 /// Completely clear the PriorityWorklist
212 void clear() {
213 M.clear();
214 V.clear();
215 }
216
217private:
218 /// A wrapper predicate designed for use with std::remove_if.
219 ///
220 /// This predicate wraps a predicate suitable for use with std::remove_if to
221 /// call M.erase(x) on each element which is slated for removal. This just
222 /// allows the predicate to be move only which we can't do with lambdas
223 /// today.
224 template <typename UnaryPredicateT>
225 class TestAndEraseFromMap {
226 UnaryPredicateT P;
227 MapT &M;
228
229 public:
230 TestAndEraseFromMap(UnaryPredicateT P, MapT &M)
231 : P(std::move(P)), M(M) {}
232
233 bool operator()(const T &Arg) {
234 if (Arg == T())
235 // Skip null values in the PriorityWorklist.
236 return false;
237
238 if (P(Arg)) {
239 M.erase(Arg);
240 return true;
241 }
242 return false;
243 }
244 };
245
246 /// The map from value to index in the vector.
247 MapT M;
248
249 /// The vector of elements in insertion order.
250 VectorT V;
251};
252
253/// A version of \c PriorityWorklist that selects small size optimized data
254/// structures for the vector and map.
255template <typename T, unsigned N>
256class SmallPriorityWorklist
257 : public PriorityWorklist<T, SmallVector<T, N>,
258 SmallDenseMap<T, ptrdiff_t>> {
259public:
260 SmallPriorityWorklist() = default;
261};
262
263} // end namespace llvm
264
265#endif // LLVM_ADT_PRIORITYWORKLIST_H
266