include/llvm/ADT/PriorityWorklist.h

   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
  29 namespace 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.
  54 template <typename T, typename VectorT = std::vector<T>,
  55           typename MapT = DenseMap<T, ptrdiff_t>>
  56 class PriorityWorklist {
  57 public:
  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
 217 private:
 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.
 255 template <typename T, unsigned N>
 256 class SmallPriorityWorklist
 257     : public PriorityWorklist<T, SmallVector<T, N>,
 258                               SmallDenseMap<T, ptrdiff_t>> {
 259 public:
 260   SmallPriorityWorklist() = default;
 261 };
 262
 263 } // end namespace llvm
 264
 265 #endif // LLVM_ADT_PRIORITYWORKLIST_H