blob: 6bc63c283b0975be2b883a68366147b988275998 [file] [log] [blame]
Andrew Scull5e1ddfa2018-08-14 10:06:54 +01001//===- llvm/ADT/BreadthFirstIterator.h - Breadth First iterator -*- C++ -*-===//
2//
3// The LLVM Compiler Infrastructure
4//
5// This file is distributed under the University of Illinois Open Source
6// License. See LICENSE.TXT for details.
7//
8//===----------------------------------------------------------------------===//
9//
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/None.h"
23#include "llvm/ADT/Optional.h"
24#include "llvm/ADT/SmallPtrSet.h"
25#include "llvm/ADT/iterator_range.h"
26#include <iterator>
27#include <queue>
28#include <utility>
29
30namespace llvm {
31
32// bf_iterator_storage - A private class which is used to figure out where to
33// store the visited set. We only provide a non-external variant for now.
34template <class SetType> class bf_iterator_storage {
35public:
36 SetType Visited;
37};
38
39// The visited state for the iteration is a simple set.
40template <typename NodeRef, unsigned SmallSize = 8>
41using bf_iterator_default_set = SmallPtrSet<NodeRef, SmallSize>;
42
43// Generic Breadth first search iterator.
44template <class GraphT,
45 class SetType =
46 bf_iterator_default_set<typename GraphTraits<GraphT>::NodeRef>,
47 class GT = GraphTraits<GraphT>>
48class bf_iterator
49 : public std::iterator<std::forward_iterator_tag, typename GT::NodeRef>,
50 public bf_iterator_storage<SetType> {
51 using super = std::iterator<std::forward_iterator_tag, typename GT::NodeRef>;
52
53 using NodeRef = typename GT::NodeRef;
54 using ChildItTy = typename GT::ChildIteratorType;
55
56 // First element is the node reference, second is the next child to visit.
57 using QueueElement = std::pair<NodeRef, Optional<ChildItTy>>;
58
59 // Visit queue - used to maintain BFS ordering.
60 // Optional<> because we need markers for levels.
61 std::queue<Optional<QueueElement>> VisitQueue;
62
63 // Current level.
64 unsigned Level;
65
66private:
67 inline bf_iterator(NodeRef Node) {
68 this->Visited.insert(Node);
69 Level = 0;
70
71 // Also, insert a dummy node as marker.
72 VisitQueue.push(QueueElement(Node, None));
73 VisitQueue.push(None);
74 }
75
76 inline bf_iterator() = default;
77
78 inline void toNext() {
79 Optional<QueueElement> Head = VisitQueue.front();
80 QueueElement H = Head.getValue();
81 NodeRef Node = H.first;
82 Optional<ChildItTy> &ChildIt = H.second;
83
84 if (!ChildIt)
85 ChildIt.emplace(GT::child_begin(Node));
86 while (*ChildIt != GT::child_end(Node)) {
87 NodeRef Next = *(*ChildIt)++;
88
89 // Already visited?
90 if (this->Visited.insert(Next).second)
91 VisitQueue.push(QueueElement(Next, None));
92 }
93 VisitQueue.pop();
94
95 // Go to the next element skipping markers if needed.
96 if (!VisitQueue.empty()) {
97 Head = VisitQueue.front();
98 if (Head != None)
99 return;
100 Level += 1;
101 VisitQueue.pop();
102
103 // Don't push another marker if this is the last
104 // element.
105 if (!VisitQueue.empty())
106 VisitQueue.push(None);
107 }
108 }
109
110public:
111 using pointer = typename super::pointer;
112
113 // Provide static begin and end methods as our public "constructors"
114 static bf_iterator begin(const GraphT &G) {
115 return bf_iterator(GT::getEntryNode(G));
116 }
117
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 const NodeRef &operator*() const { return VisitQueue.front()->first; }
127
128 // This is a nonstandard operator-> that dereferenfces 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
144 unsigned getLevel() const { return Level; }
145};
146
147// Provide global constructors that automatically figure out correct types.
148template <class T> bf_iterator<T> bf_begin(const T &G) {
149 return bf_iterator<T>::begin(G);
150}
151
152template <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.
157template <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