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