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