llvm/ADT/BreadthFirstIterator.h

0001 //===- llvm/ADT/BreadthFirstIterator.h - Breadth First iterator -*- 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 /// This file builds on the ADT/GraphTraits.h file to build a generic breadth
0011 /// first graph iterator.  This file exposes the following functions/types:
0012 ///
0013 /// bf_begin/bf_end/bf_iterator
0014 ///   * Normal breadth-first iteration - visit a graph level-by-level.
0015 ///
0016 //===----------------------------------------------------------------------===//
0017
0018 #ifndef LLVM_ADT_BREADTHFIRSTITERATOR_H
0019 #define LLVM_ADT_BREADTHFIRSTITERATOR_H
0020
0021 #include "llvm/ADT/GraphTraits.h"
0022 #include "llvm/ADT/SmallPtrSet.h"
0023 #include "llvm/ADT/iterator_range.h"
0024 #include <iterator>
0025 #include <optional>
0026 #include <queue>
0027 #include <utility>
0028
0029 namespace llvm {
0030
0031 // bf_iterator_storage - A private class which is used to figure out where to
0032 // store the visited set. We only provide a non-external variant for now.
0033 template <class SetType> class bf_iterator_storage {
0034 public:
0035   SetType Visited;
0036 };
0037
0038 // The visited state for the iteration is a simple set.
0039 template <typename NodeRef, unsigned SmallSize = 8>
0040 using bf_iterator_default_set = SmallPtrSet<NodeRef, SmallSize>;
0041
0042 // Generic Breadth first search iterator.
0043 template <class GraphT,
0044           class SetType =
0045               bf_iterator_default_set<typename GraphTraits<GraphT>::NodeRef>,
0046           class GT = GraphTraits<GraphT>>
0047 class bf_iterator : public bf_iterator_storage<SetType> {
0048 public:
0049   using iterator_category = std::forward_iterator_tag;
0050   using value_type = typename GT::NodeRef;
0051   using difference_type = std::ptrdiff_t;
0052   using pointer = value_type *;
0053   using reference = const value_type &;
0054
0055 private:
0056   using NodeRef = typename GT::NodeRef;
0057   using ChildItTy = typename GT::ChildIteratorType;
0058
0059   // First element is the node reference, second is the next child to visit.
0060   using QueueElement = std::pair<NodeRef, std::optional<ChildItTy>>;
0061
0062   // Visit queue - used to maintain BFS ordering.
0063   // std::optional<> because we need markers for levels.
0064   std::queue<std::optional<QueueElement>> VisitQueue;
0065
0066   // Current level.
0067   unsigned Level = 0;
0068
0069   inline bf_iterator(NodeRef Node) {
0070     this->Visited.insert(Node);
0071     Level = 0;
0072
0073     // Also, insert a dummy node as marker.
0074     VisitQueue.push(QueueElement(Node, std::nullopt));
0075     VisitQueue.push(std::nullopt);
0076   }
0077
0078   inline bf_iterator() = default;
0079
0080   inline void toNext() {
0081     std::optional<QueueElement> Head = VisitQueue.front();
0082     QueueElement H = *Head;
0083     NodeRef Node = H.first;
0084     std::optional<ChildItTy> &ChildIt = H.second;
0085
0086     if (!ChildIt)
0087       ChildIt.emplace(GT::child_begin(Node));
0088     while (*ChildIt != GT::child_end(Node)) {
0089       NodeRef Next = *(*ChildIt)++;
0090
0091       // Already visited?
0092       if (this->Visited.insert(Next).second)
0093         VisitQueue.push(QueueElement(Next, std::nullopt));
0094     }
0095     VisitQueue.pop();
0096
0097     // Go to the next element skipping markers if needed.
0098     if (!VisitQueue.empty()) {
0099       Head = VisitQueue.front();
0100       if (Head != std::nullopt)
0101         return;
0102       Level += 1;
0103       VisitQueue.pop();
0104
0105       // Don't push another marker if this is the last
0106       // element.
0107       if (!VisitQueue.empty())
0108         VisitQueue.push(std::nullopt);
0109     }
0110   }
0111
0112 public:
0113   // Provide static begin and end methods as our public "constructors"
0114   static bf_iterator begin(const GraphT &G) {
0115     return bf_iterator(GT::getEntryNode(G));
0116   }
0117
0118   static bf_iterator end(const GraphT &G) { return bf_iterator(); }
0119
0120   bool operator==(const bf_iterator &RHS) const {
0121     return VisitQueue == RHS.VisitQueue;
0122   }
0123
0124   bool operator!=(const bf_iterator &RHS) const { return !(*this == RHS); }
0125
0126   reference operator*() const { return VisitQueue.front()->first; }
0127
0128   // This is a nonstandard operator-> that dereferences the pointer an extra
0129   // time so that you can actually call methods on the node, because the
0130   // contained type is a pointer.
0131   NodeRef operator->() const { return **this; }
0132
0133   bf_iterator &operator++() { // Pre-increment
0134     toNext();
0135     return *this;
0136   }
0137
0138   bf_iterator operator++(int) { // Post-increment
0139     bf_iterator ItCopy = *this;
0140     ++*this;
0141     return ItCopy;
0142   }
0143
0144   unsigned getLevel() const { return Level; }
0145 };
0146
0147 // Provide global constructors that automatically figure out correct types.
0148 template <class T> bf_iterator<T> bf_begin(const T &G) {
0149   return bf_iterator<T>::begin(G);
0150 }
0151
0152 template <class T> bf_iterator<T> bf_end(const T &G) {
0153   return bf_iterator<T>::end(G);
0154 }
0155
0156 // Provide an accessor method to use them in range-based patterns.
0157 template <class T> iterator_range<bf_iterator<T>> breadth_first(const T &G) {
0158   return make_range(bf_begin(G), bf_end(G));
0159 }
0160
0161 } // end namespace llvm
0162
0163 #endif // LLVM_ADT_BREADTHFIRSTITERATOR_H