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  typename std::enable_if<!std::is_convertible<SequenceT, T>::value>::type
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