include/llvm/ADT/BreadthFirstIterator.h

   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
  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
  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
  65 private:
  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
 109 public:
 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.
 147 template <class T> bf_iterator<T> bf_begin(const T &G) {
 148   return bf_iterator<T>::begin(G);
 149 }
 150
 151 template <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.
 156 template <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