1 //===- llvm/ADT/BreadthFirstIterator.h - Breadth First iterator -*- 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 // This file builds on the ADT/GraphTraits.h file to build a generic breadth
10 // first graph iterator.  This file exposes the following functions/types:
11 //
12 // bf_begin/bf_end/bf_iterator
13 //   * Normal breadth-first iteration - visit a graph level-by-level.
14 //
15 //===----------------------------------------------------------------------===//
16 
17 #ifndef LLVM_ADT_BREADTHFIRSTITERATOR_H
18 #define LLVM_ADT_BREADTHFIRSTITERATOR_H
19 
20 #include "llvm/ADT/GraphTraits.h"
21 #include "llvm/ADT/None.h"
22 #include "llvm/ADT/Optional.h"
23 #include "llvm/ADT/SmallPtrSet.h"
24 #include "llvm/ADT/iterator_range.h"
25 #include <iterator>
26 #include <queue>
27 #include <utility>
28 
29 namespace llvm {
30 
31 // bf_iterator_storage - A private class which is used to figure out where to
32 // store the visited set. We only provide a non-external variant for now.
33 template <class SetType> class bf_iterator_storage {
34 public:
35   SetType Visited;
36 };
37 
38 // The visited state for the iteration is a simple set.
39 template <typename NodeRef, unsigned SmallSize = 8>
40 using bf_iterator_default_set = SmallPtrSet<NodeRef, SmallSize>;
41 
42 // Generic Breadth first search iterator.
43 template <class GraphT,
44           class SetType =
45               bf_iterator_default_set<typename GraphTraits<GraphT>::NodeRef>,
46           class GT = GraphTraits<GraphT>>
47 class bf_iterator
48     : public std::iterator<std::forward_iterator_tag, typename GT::NodeRef>,
49       public bf_iterator_storage<SetType> {
50   using super = std::iterator<std::forward_iterator_tag, typename GT::NodeRef>;
51 
52   using NodeRef = typename GT::NodeRef;
53   using ChildItTy = typename GT::ChildIteratorType;
54 
55   // First element is the node reference, second is the next child to visit.
56   using QueueElement = std::pair<NodeRef, Optional<ChildItTy>>;
57 
58   // Visit queue - used to maintain BFS ordering.
59   // Optional<> because we need markers for levels.
60   std::queue<Optional<QueueElement>> VisitQueue;
61 
62   // Current level.
63   unsigned Level;
64 
65 private:
66   inline bf_iterator(NodeRef Node) {
67     this->Visited.insert(Node);
68     Level = 0;
69 
70     // Also, insert a dummy node as marker.
71     VisitQueue.push(QueueElement(Node, None));
72     VisitQueue.push(None);
73   }
74 
75   inline bf_iterator() = default;
76 
77   inline void toNext() {
78     Optional<QueueElement> Head = VisitQueue.front();
79     QueueElement H = Head.getValue();
80     NodeRef Node = H.first;
81     Optional<ChildItTy> &ChildIt = H.second;
82 
83     if (!ChildIt)
84       ChildIt.emplace(GT::child_begin(Node));
85     while (*ChildIt != GT::child_end(Node)) {
86       NodeRef Next = *(*ChildIt)++;
87 
88       // Already visited?
89       if (this->Visited.insert(Next).second)
90         VisitQueue.push(QueueElement(Next, None));
91     }
92     VisitQueue.pop();
93 
94     // Go to the next element skipping markers if needed.
95     if (!VisitQueue.empty()) {
96       Head = VisitQueue.front();
97       if (Head != None)
98         return;
99       Level += 1;
100       VisitQueue.pop();
101 
102       // Don't push another marker if this is the last
103       // element.
104       if (!VisitQueue.empty())
105         VisitQueue.push(None);
106     }
107   }
108 
109 public:
110   using pointer = typename super::pointer;
111 
112   // Provide static begin and end methods as our public "constructors"
113   static bf_iterator begin(const GraphT &G) {
114     return bf_iterator(GT::getEntryNode(G));
115   }
116 
117   static bf_iterator end(const GraphT &G) { return bf_iterator(); }
118 
119   bool operator==(const bf_iterator &RHS) const {
120     return VisitQueue == RHS.VisitQueue;
121   }
122 
123   bool operator!=(const bf_iterator &RHS) const { return !(*this == RHS); }
124 
125   const NodeRef &operator*() const { return VisitQueue.front()->first; }
126 
127   // This is a nonstandard operator-> that dereferences the pointer an extra
128   // time so that you can actually call methods on the node, because the
129   // contained type is a pointer.
130   NodeRef operator->() const { return **this; }
131 
132   bf_iterator &operator++() { // Pre-increment
133     toNext();
134     return *this;
135   }
136 
137   bf_iterator operator++(int) { // Post-increment
138     bf_iterator ItCopy = *this;
139     ++*this;
140     return ItCopy;
141   }
142 
143   unsigned getLevel() const { return Level; }
144 };
145 
146 // Provide global constructors that automatically figure out correct types.
147 template <class T> bf_iterator<T> bf_begin(const T &G) {
148   return bf_iterator<T>::begin(G);
149 }
150 
151 template <class T> bf_iterator<T> bf_end(const T &G) {
152   return bf_iterator<T>::end(G);
153 }
154 
155 // Provide an accessor method to use them in range-based patterns.
156 template <class T> iterator_range<bf_iterator<T>> breadth_first(const T &G) {
157   return make_range(bf_begin(G), bf_end(G));
158 }
159 
160 } // end namespace llvm
161 
162 #endif // LLVM_ADT_BREADTHFIRSTITERATOR_H
163