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 = 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 
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 
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"
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 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 
144   unsigned getLevel() const { return Level; }
145 };
146 
147 // Provide global constructors that automatically figure out correct types.
148 template <class T> bf_iterator<T> bf_begin(const T &G) {
149   return bf_iterator<T>::begin(G);
150 }
151 
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.
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