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