1303231Sdim//===- PriorityWorklist.h - Worklist with insertion priority ----*- C++ -*-===//
2303231Sdim//
3353358Sdim// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4353358Sdim// See https://llvm.org/LICENSE.txt for license information.
5353358Sdim// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6303231Sdim//
7303231Sdim//===----------------------------------------------------------------------===//
8303231Sdim///
9303231Sdim/// \file
10303231Sdim///
11303231Sdim/// This file provides a priority worklist. See the class comments for details.
12303231Sdim///
13303231Sdim//===----------------------------------------------------------------------===//
14303231Sdim
15303231Sdim#ifndef LLVM_ADT_PRIORITYWORKLIST_H
16303231Sdim#define LLVM_ADT_PRIORITYWORKLIST_H
17303231Sdim
18303231Sdim#include "llvm/ADT/DenseMap.h"
19314564Sdim#include "llvm/ADT/STLExtras.h"
20303231Sdim#include "llvm/ADT/SmallVector.h"
21314564Sdim#include "llvm/Support/Compiler.h"
22303231Sdim#include <algorithm>
23303231Sdim#include <cassert>
24314564Sdim#include <cstddef>
25321369Sdim#include <iterator>
26321369Sdim#include <type_traits>
27303231Sdim#include <vector>
28303231Sdim
29303231Sdimnamespace llvm {
30303231Sdim
31303231Sdim/// A FILO worklist that prioritizes on re-insertion without duplication.
32303231Sdim///
33303231Sdim/// This is very similar to a \c SetVector with the primary difference that
34303231Sdim/// while re-insertion does not create a duplicate, it does adjust the
35303231Sdim/// visitation order to respect the last insertion point. This can be useful
36303231Sdim/// when the visit order needs to be prioritized based on insertion point
37303231Sdim/// without actually having duplicate visits.
38303231Sdim///
39303231Sdim/// Note that this doesn't prevent re-insertion of elements which have been
40303231Sdim/// visited -- if you need to break cycles, a set will still be necessary.
41303231Sdim///
42303231Sdim/// The type \c T must be default constructable to a null value that will be
43303231Sdim/// ignored. It is an error to insert such a value, and popping elements will
44303231Sdim/// never produce such a value. It is expected to be used with common nullable
45303231Sdim/// types like pointers or optionals.
46303231Sdim///
47303231Sdim/// Internally this uses a vector to store the worklist and a map to identify
48303231Sdim/// existing elements in the worklist. Both of these may be customized, but the
49303231Sdim/// map must support the basic DenseMap API for mapping from a T to an integer
50303231Sdim/// index into the vector.
51303231Sdim///
52303231Sdim/// A partial specialization is provided to automatically select a SmallVector
53303231Sdim/// and a SmallDenseMap if custom data structures are not provided.
54303231Sdimtemplate <typename T, typename VectorT = std::vector<T>,
55303231Sdim          typename MapT = DenseMap<T, ptrdiff_t>>
56303231Sdimclass PriorityWorklist {
57303231Sdimpublic:
58321369Sdim  using value_type = T;
59321369Sdim  using key_type = T;
60321369Sdim  using reference = T&;
61321369Sdim  using const_reference = const T&;
62321369Sdim  using size_type = typename MapT::size_type;
63303231Sdim
64303231Sdim  /// Construct an empty PriorityWorklist
65314564Sdim  PriorityWorklist() = default;
66303231Sdim
67303231Sdim  /// Determine if the PriorityWorklist is empty or not.
68303231Sdim  bool empty() const {
69303231Sdim    return V.empty();
70303231Sdim  }
71303231Sdim
72303231Sdim  /// Returns the number of elements in the worklist.
73303231Sdim  size_type size() const {
74303231Sdim    return M.size();
75303231Sdim  }
76303231Sdim
77303231Sdim  /// Count the number of elements of a given key in the PriorityWorklist.
78303231Sdim  /// \returns 0 if the element is not in the PriorityWorklist, 1 if it is.
79303231Sdim  size_type count(const key_type &key) const {
80303231Sdim    return M.count(key);
81303231Sdim  }
82303231Sdim
83303231Sdim  /// Return the last element of the PriorityWorklist.
84303231Sdim  const T &back() const {
85303231Sdim    assert(!empty() && "Cannot call back() on empty PriorityWorklist!");
86303231Sdim    return V.back();
87303231Sdim  }
88303231Sdim
89303231Sdim  /// Insert a new element into the PriorityWorklist.
90303231Sdim  /// \returns true if the element was inserted into the PriorityWorklist.
91303231Sdim  bool insert(const T &X) {
92303231Sdim    assert(X != T() && "Cannot insert a null (default constructed) value!");
93303231Sdim    auto InsertResult = M.insert({X, V.size()});
94303231Sdim    if (InsertResult.second) {
95303231Sdim      // Fresh value, just append it to the vector.
96303231Sdim      V.push_back(X);
97303231Sdim      return true;
98303231Sdim    }
99303231Sdim
100303231Sdim    auto &Index = InsertResult.first->second;
101303231Sdim    assert(V[Index] == X && "Value not actually at index in map!");
102303231Sdim    if (Index != (ptrdiff_t)(V.size() - 1)) {
103303231Sdim      // If the element isn't at the back, null it out and append a fresh one.
104303231Sdim      V[Index] = T();
105303231Sdim      Index = (ptrdiff_t)V.size();
106303231Sdim      V.push_back(X);
107303231Sdim    }
108303231Sdim    return false;
109303231Sdim  }
110303231Sdim
111314564Sdim  /// Insert a sequence of new elements into the PriorityWorklist.
112314564Sdim  template <typename SequenceT>
113314564Sdim  typename std::enable_if<!std::is_convertible<SequenceT, T>::value>::type
114314564Sdim  insert(SequenceT &&Input) {
115314564Sdim    if (std::begin(Input) == std::end(Input))
116314564Sdim      // Nothing to do for an empty input sequence.
117314564Sdim      return;
118314564Sdim
119314564Sdim    // First pull the input sequence into the vector as a bulk append
120314564Sdim    // operation.
121314564Sdim    ptrdiff_t StartIndex = V.size();
122314564Sdim    V.insert(V.end(), std::begin(Input), std::end(Input));
123314564Sdim    // Now walk backwards fixing up the index map and deleting any duplicates.
124314564Sdim    for (ptrdiff_t i = V.size() - 1; i >= StartIndex; --i) {
125314564Sdim      auto InsertResult = M.insert({V[i], i});
126314564Sdim      if (InsertResult.second)
127314564Sdim        continue;
128314564Sdim
129314564Sdim      // If the existing index is before this insert's start, nuke that one and
130314564Sdim      // move it up.
131314564Sdim      ptrdiff_t &Index = InsertResult.first->second;
132314564Sdim      if (Index < StartIndex) {
133314564Sdim        V[Index] = T();
134314564Sdim        Index = i;
135314564Sdim        continue;
136314564Sdim      }
137314564Sdim
138314564Sdim      // Otherwise the existing one comes first so just clear out the value in
139314564Sdim      // this slot.
140314564Sdim      V[i] = T();
141314564Sdim    }
142314564Sdim  }
143314564Sdim
144303231Sdim  /// Remove the last element of the PriorityWorklist.
145303231Sdim  void pop_back() {
146303231Sdim    assert(!empty() && "Cannot remove an element when empty!");
147303231Sdim    assert(back() != T() && "Cannot have a null element at the back!");
148303231Sdim    M.erase(back());
149303231Sdim    do {
150303231Sdim      V.pop_back();
151303231Sdim    } while (!V.empty() && V.back() == T());
152303231Sdim  }
153303231Sdim
154314564Sdim  LLVM_NODISCARD T pop_back_val() {
155303231Sdim    T Ret = back();
156303231Sdim    pop_back();
157303231Sdim    return Ret;
158303231Sdim  }
159303231Sdim
160303231Sdim  /// Erase an item from the worklist.
161303231Sdim  ///
162303231Sdim  /// Note that this is constant time due to the nature of the worklist implementation.
163303231Sdim  bool erase(const T& X) {
164303231Sdim    auto I = M.find(X);
165303231Sdim    if (I == M.end())
166303231Sdim      return false;
167303231Sdim
168303231Sdim    assert(V[I->second] == X && "Value not actually at index in map!");
169303231Sdim    if (I->second == (ptrdiff_t)(V.size() - 1)) {
170303231Sdim      do {
171303231Sdim        V.pop_back();
172303231Sdim      } while (!V.empty() && V.back() == T());
173303231Sdim    } else {
174303231Sdim      V[I->second] = T();
175303231Sdim    }
176303231Sdim    M.erase(I);
177303231Sdim    return true;
178303231Sdim  }
179303231Sdim
180303231Sdim  /// Erase items from the set vector based on a predicate function.
181303231Sdim  ///
182303231Sdim  /// This is intended to be equivalent to the following code, if we could
183303231Sdim  /// write it:
184303231Sdim  ///
185303231Sdim  /// \code
186314564Sdim  ///   V.erase(remove_if(V, P), V.end());
187303231Sdim  /// \endcode
188303231Sdim  ///
189303231Sdim  /// However, PriorityWorklist doesn't expose non-const iterators, making any
190303231Sdim  /// algorithm like remove_if impossible to use.
191303231Sdim  ///
192303231Sdim  /// \returns true if any element is removed.
193303231Sdim  template <typename UnaryPredicate>
194303231Sdim  bool erase_if(UnaryPredicate P) {
195314564Sdim    typename VectorT::iterator E =
196314564Sdim        remove_if(V, TestAndEraseFromMap<UnaryPredicate>(P, M));
197303231Sdim    if (E == V.end())
198303231Sdim      return false;
199303231Sdim    for (auto I = V.begin(); I != E; ++I)
200303231Sdim      if (*I != T())
201303231Sdim        M[*I] = I - V.begin();
202303231Sdim    V.erase(E, V.end());
203303231Sdim    return true;
204303231Sdim  }
205303231Sdim
206314564Sdim  /// Reverse the items in the PriorityWorklist.
207314564Sdim  ///
208314564Sdim  /// This does an in-place reversal. Other kinds of reverse aren't easy to
209314564Sdim  /// support in the face of the worklist semantics.
210314564Sdim
211303231Sdim  /// Completely clear the PriorityWorklist
212303231Sdim  void clear() {
213303231Sdim    M.clear();
214303231Sdim    V.clear();
215303231Sdim  }
216303231Sdim
217303231Sdimprivate:
218303231Sdim  /// A wrapper predicate designed for use with std::remove_if.
219303231Sdim  ///
220303231Sdim  /// This predicate wraps a predicate suitable for use with std::remove_if to
221303231Sdim  /// call M.erase(x) on each element which is slated for removal. This just
222303231Sdim  /// allows the predicate to be move only which we can't do with lambdas
223303231Sdim  /// today.
224303231Sdim  template <typename UnaryPredicateT>
225303231Sdim  class TestAndEraseFromMap {
226303231Sdim    UnaryPredicateT P;
227303231Sdim    MapT &M;
228303231Sdim
229303231Sdim  public:
230303231Sdim    TestAndEraseFromMap(UnaryPredicateT P, MapT &M)
231303231Sdim        : P(std::move(P)), M(M) {}
232303231Sdim
233303231Sdim    bool operator()(const T &Arg) {
234303231Sdim      if (Arg == T())
235303231Sdim        // Skip null values in the PriorityWorklist.
236303231Sdim        return false;
237303231Sdim
238303231Sdim      if (P(Arg)) {
239303231Sdim        M.erase(Arg);
240303231Sdim        return true;
241303231Sdim      }
242303231Sdim      return false;
243303231Sdim    }
244303231Sdim  };
245303231Sdim
246303231Sdim  /// The map from value to index in the vector.
247303231Sdim  MapT M;
248303231Sdim
249303231Sdim  /// The vector of elements in insertion order.
250303231Sdim  VectorT V;
251303231Sdim};
252303231Sdim
253303231Sdim/// A version of \c PriorityWorklist that selects small size optimized data
254303231Sdim/// structures for the vector and map.
255303231Sdimtemplate <typename T, unsigned N>
256303231Sdimclass SmallPriorityWorklist
257303231Sdim    : public PriorityWorklist<T, SmallVector<T, N>,
258303231Sdim                              SmallDenseMap<T, ptrdiff_t>> {
259303231Sdimpublic:
260314564Sdim  SmallPriorityWorklist() = default;
261303231Sdim};
262303231Sdim
263314564Sdim} // end namespace llvm
264303231Sdim
265314564Sdim#endif // LLVM_ADT_PRIORITYWORKLIST_H
266