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