BreadthFirstIterator.h 4.8 KB

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