1//===- 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 <set>
29#include <utility>
30
31namespace llvm {
32
33// bf_iterator_storage - A private class which is used to figure out where to
34// store the visited set. We only provide a non-external variant for now.
35template <class SetType> class bf_iterator_storage {
36public:
37  SetType Visited;
38};
39
40// The visited state for the iteration is a simple set.
41template <typename NodeRef, unsigned SmallSize = 8>
42using bf_iterator_default_set = SmallPtrSet<NodeRef, SmallSize>;
43
44// Generic Breadth first search iterator.
45template <class GraphT,
46          class SetType =
47              bf_iterator_default_set<typename GraphTraits<GraphT>::NodeRef>,
48          class GT = GraphTraits<GraphT>>
49class bf_iterator
50    : public std::iterator<std::forward_iterator_tag, typename GT::NodeRef>,
51      public bf_iterator_storage<SetType> {
52  typedef std::iterator<std::forward_iterator_tag, typename GT::NodeRef> super;
53
54  typedef typename GT::NodeRef NodeRef;
55  typedef typename GT::ChildIteratorType ChildItTy;
56
57  // First element is the node reference, second is the next child to visit.
58  typedef std::pair<NodeRef, Optional<ChildItTy>> QueueElement;
59
60  // Visit queue - used to maintain BFS ordering.
61  // Optional<> because we need markers for levels.
62  std::queue<Optional<QueueElement>> VisitQueue;
63
64  // Current level.
65  unsigned Level;
66
67private:
68  inline bf_iterator(NodeRef Node) {
69    this->Visited.insert(Node);
70    Level = 0;
71
72    // Also, insert a dummy node as marker.
73    VisitQueue.push(QueueElement(Node, None));
74    VisitQueue.push(None);
75  }
76
77  inline bf_iterator() = default;
78
79  inline void toNext() {
80    Optional<QueueElement> Head = VisitQueue.front();
81    QueueElement H = Head.getValue();
82    NodeRef Node = H.first;
83    Optional<ChildItTy> &ChildIt = H.second;
84
85    if (!ChildIt)
86      ChildIt.emplace(GT::child_begin(Node));
87    while (*ChildIt != GT::child_end(Node)) {
88      NodeRef Next = *(*ChildIt)++;
89
90      // Already visited?
91      if (this->Visited.insert(Next).second)
92        VisitQueue.push(QueueElement(Next, None));
93    }
94    VisitQueue.pop();
95
96    // Go to the next element skipping markers if needed.
97    if (!VisitQueue.empty()) {
98      Head = VisitQueue.front();
99      if (Head != None)
100        return;
101      Level += 1;
102      VisitQueue.pop();
103
104      // Don't push another marker if this is the last
105      // element.
106      if (!VisitQueue.empty())
107        VisitQueue.push(None);
108    }
109  }
110
111public:
112  typedef typename super::pointer pointer;
113
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 dereferenfces 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.
149template <class T> bf_iterator<T> bf_begin(const T &G) {
150  return bf_iterator<T>::begin(G);
151}
152
153template <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.
158template <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