llvm/ADT/PriorityWorklist.h

0001 //===- PriorityWorklist.h - Worklist with insertion priority ----*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008 ///
0009 /// \file
0010 ///
0011 /// This file provides a priority worklist. See the class comments for details.
0012 ///
0013 //===----------------------------------------------------------------------===//
0014
0015 #ifndef LLVM_ADT_PRIORITYWORKLIST_H
0016 #define LLVM_ADT_PRIORITYWORKLIST_H
0017
0018 #include "llvm/ADT/DenseMap.h"
0019 #include "llvm/ADT/STLExtras.h"
0020 #include "llvm/ADT/SmallVector.h"
0021 #include "llvm/Support/Compiler.h"
0022 #include <cassert>
0023 #include <cstddef>
0024 #include <iterator>
0025 #include <type_traits>
0026 #include <vector>
0027
0028 namespace llvm {
0029
0030 /// A FILO worklist that prioritizes on re-insertion without duplication.
0031 ///
0032 /// This is very similar to a \c SetVector with the primary difference that
0033 /// while re-insertion does not create a duplicate, it does adjust the
0034 /// visitation order to respect the last insertion point. This can be useful
0035 /// when the visit order needs to be prioritized based on insertion point
0036 /// without actually having duplicate visits.
0037 ///
0038 /// Note that this doesn't prevent re-insertion of elements which have been
0039 /// visited -- if you need to break cycles, a set will still be necessary.
0040 ///
0041 /// The type \c T must be default constructable to a null value that will be
0042 /// ignored. It is an error to insert such a value, and popping elements will
0043 /// never produce such a value. It is expected to be used with common nullable
0044 /// types like pointers or optionals.
0045 ///
0046 /// Internally this uses a vector to store the worklist and a map to identify
0047 /// existing elements in the worklist. Both of these may be customized, but the
0048 /// map must support the basic DenseMap API for mapping from a T to an integer
0049 /// index into the vector.
0050 ///
0051 /// A partial specialization is provided to automatically select a SmallVector
0052 /// and a SmallDenseMap if custom data structures are not provided.
0053 template <typename T, typename VectorT = std::vector<T>,
0054           typename MapT = DenseMap<T, ptrdiff_t>>
0055 class PriorityWorklist {
0056 public:
0057   using value_type = T;
0058   using key_type = T;
0059   using reference = T&;
0060   using const_reference = const T&;
0061   using size_type = typename MapT::size_type;
0062
0063   /// Construct an empty PriorityWorklist
0064   PriorityWorklist() = default;
0065
0066   /// Determine if the PriorityWorklist is empty or not.
0067   bool empty() const {
0068     return V.empty();
0069   }
0070
0071   /// Returns the number of elements in the worklist.
0072   size_type size() const {
0073     return M.size();
0074   }
0075
0076   /// Count the number of elements of a given key in the PriorityWorklist.
0077   /// \returns 0 if the element is not in the PriorityWorklist, 1 if it is.
0078   size_type count(const key_type &key) const {
0079     return M.count(key);
0080   }
0081
0082   /// Return the last element of the PriorityWorklist.
0083   const T &back() const {
0084     assert(!empty() && "Cannot call back() on empty PriorityWorklist!");
0085     return V.back();
0086   }
0087
0088   /// Insert a new element into the PriorityWorklist.
0089   /// \returns true if the element was inserted into the PriorityWorklist.
0090   bool insert(const T &X) {
0091     assert(X != T() && "Cannot insert a null (default constructed) value!");
0092     auto InsertResult = M.insert({X, V.size()});
0093     if (InsertResult.second) {
0094       // Fresh value, just append it to the vector.
0095       V.push_back(X);
0096       return true;
0097     }
0098
0099     auto &Index = InsertResult.first->second;
0100     assert(V[Index] == X && "Value not actually at index in map!");
0101     if (Index != (ptrdiff_t)(V.size() - 1)) {
0102       // If the element isn't at the back, null it out and append a fresh one.
0103       V[Index] = T();
0104       Index = (ptrdiff_t)V.size();
0105       V.push_back(X);
0106     }
0107     return false;
0108   }
0109
0110   /// Insert a sequence of new elements into the PriorityWorklist.
0111   template <typename SequenceT>
0112   std::enable_if_t<!std::is_convertible<SequenceT, T>::value>
0113   insert(SequenceT &&Input) {
0114     if (std::begin(Input) == std::end(Input))
0115       // Nothing to do for an empty input sequence.
0116       return;
0117
0118     // First pull the input sequence into the vector as a bulk append
0119     // operation.
0120     ptrdiff_t StartIndex = V.size();
0121     V.insert(V.end(), std::begin(Input), std::end(Input));
0122     // Now walk backwards fixing up the index map and deleting any duplicates.
0123     for (ptrdiff_t i = V.size() - 1; i >= StartIndex; --i) {
0124       auto InsertResult = M.insert({V[i], i});
0125       if (InsertResult.second)
0126         continue;
0127
0128       // If the existing index is before this insert's start, nuke that one and
0129       // move it up.
0130       ptrdiff_t &Index = InsertResult.first->second;
0131       if (Index < StartIndex) {
0132         V[Index] = T();
0133         Index = i;
0134         continue;
0135       }
0136
0137       // Otherwise the existing one comes first so just clear out the value in
0138       // this slot.
0139       V[i] = T();
0140     }
0141   }
0142
0143   /// Remove the last element of the PriorityWorklist.
0144   void pop_back() {
0145     assert(!empty() && "Cannot remove an element when empty!");
0146     assert(back() != T() && "Cannot have a null element at the back!");
0147     M.erase(back());
0148     do {
0149       V.pop_back();
0150     } while (!V.empty() && V.back() == T());
0151   }
0152
0153   [[nodiscard]] T pop_back_val() {
0154     T Ret = back();
0155     pop_back();
0156     return Ret;
0157   }
0158
0159   /// Erase an item from the worklist.
0160   ///
0161   /// Note that this is constant time due to the nature of the worklist implementation.
0162   bool erase(const T& X) {
0163     auto I = M.find(X);
0164     if (I == M.end())
0165       return false;
0166
0167     assert(V[I->second] == X && "Value not actually at index in map!");
0168     if (I->second == (ptrdiff_t)(V.size() - 1)) {
0169       do {
0170         V.pop_back();
0171       } while (!V.empty() && V.back() == T());
0172     } else {
0173       V[I->second] = T();
0174     }
0175     M.erase(I);
0176     return true;
0177   }
0178
0179   /// Erase items from the set vector based on a predicate function.
0180   ///
0181   /// This is intended to be equivalent to the following code, if we could
0182   /// write it:
0183   ///
0184   /// \code
0185   ///   V.erase(remove_if(V, P), V.end());
0186   /// \endcode
0187   ///
0188   /// However, PriorityWorklist doesn't expose non-const iterators, making any
0189   /// algorithm like remove_if impossible to use.
0190   ///
0191   /// \returns true if any element is removed.
0192   template <typename UnaryPredicate>
0193   bool erase_if(UnaryPredicate P) {
0194     typename VectorT::iterator E =
0195         remove_if(V, TestAndEraseFromMap<UnaryPredicate>(P, M));
0196     if (E == V.end())
0197       return false;
0198     for (auto I = V.begin(); I != E; ++I)
0199       if (*I != T())
0200         M[*I] = I - V.begin();
0201     V.erase(E, V.end());
0202     return true;
0203   }
0204
0205   /// Reverse the items in the PriorityWorklist.
0206   ///
0207   /// This does an in-place reversal. Other kinds of reverse aren't easy to
0208   /// support in the face of the worklist semantics.
0209
0210   /// Completely clear the PriorityWorklist
0211   void clear() {
0212     M.clear();
0213     V.clear();
0214   }
0215
0216 private:
0217   /// A wrapper predicate designed for use with std::remove_if.
0218   ///
0219   /// This predicate wraps a predicate suitable for use with std::remove_if to
0220   /// call M.erase(x) on each element which is slated for removal. This just
0221   /// allows the predicate to be move only which we can't do with lambdas
0222   /// today.
0223   template <typename UnaryPredicateT>
0224   class TestAndEraseFromMap {
0225     UnaryPredicateT P;
0226     MapT &M;
0227
0228   public:
0229     TestAndEraseFromMap(UnaryPredicateT P, MapT &M)
0230         : P(std::move(P)), M(M) {}
0231
0232     bool operator()(const T &Arg) {
0233       if (Arg == T())
0234         // Skip null values in the PriorityWorklist.
0235         return false;
0236
0237       if (P(Arg)) {
0238         M.erase(Arg);
0239         return true;
0240       }
0241       return false;
0242     }
0243   };
0244
0245   /// The map from value to index in the vector.
0246   MapT M;
0247
0248   /// The vector of elements in insertion order.
0249   VectorT V;
0250 };
0251
0252 /// A version of \c PriorityWorklist that selects small size optimized data
0253 /// structures for the vector and map.
0254 template <typename T, unsigned N>
0255 class SmallPriorityWorklist
0256     : public PriorityWorklist<T, SmallVector<T, N>,
0257                               SmallDenseMap<T, ptrdiff_t>> {
0258 public:
0259   SmallPriorityWorklist() = default;
0260 };
0261
0262 } // end namespace llvm
0263
0264 #endif // LLVM_ADT_PRIORITYWORKLIST_H