BreadthFirstIterator.h 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  1. #pragma once
  2. #ifdef __GNUC__
  3. #pragma GCC diagnostic push
  4. #pragma GCC diagnostic ignored "-Wunused-parameter"
  5. #endif
  6. //===- llvm/ADT/BreadthFirstIterator.h - Breadth First iterator -*- C++ -*-===//
  7. //
  8. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
  9. // See https://llvm.org/LICENSE.txt for license information.
  10. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
  11. //
  12. //===----------------------------------------------------------------------===//
  13. ///
  14. /// \file
  15. /// This file builds on the ADT/GraphTraits.h file to build a generic breadth
  16. /// first graph iterator. This file exposes the following functions/types:
  17. ///
  18. /// bf_begin/bf_end/bf_iterator
  19. /// * Normal breadth-first iteration - visit a graph level-by-level.
  20. ///
  21. //===----------------------------------------------------------------------===//
  22. #ifndef LLVM_ADT_BREADTHFIRSTITERATOR_H
  23. #define LLVM_ADT_BREADTHFIRSTITERATOR_H
  24. #include "llvm/ADT/GraphTraits.h"
  25. #include "llvm/ADT/None.h"
  26. #include "llvm/ADT/Optional.h"
  27. #include "llvm/ADT/SmallPtrSet.h"
  28. #include "llvm/ADT/iterator_range.h"
  29. #include <iterator>
  30. #include <queue>
  31. #include <utility>
  32. namespace llvm {
  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.
  35. template <class SetType> class bf_iterator_storage {
  36. public:
  37. SetType Visited;
  38. };
  39. // The visited state for the iteration is a simple set.
  40. template <typename NodeRef, unsigned SmallSize = 8>
  41. using bf_iterator_default_set = SmallPtrSet<NodeRef, SmallSize>;
  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. private:
  55. using NodeRef = typename GT::NodeRef;
  56. using ChildItTy = typename GT::ChildIteratorType;
  57. // First element is the node reference, second is the next child to visit.
  58. using QueueElement = std::pair<NodeRef, Optional<ChildItTy>>;
  59. // Visit queue - used to maintain BFS ordering.
  60. // Optional<> because we need markers for levels.
  61. std::queue<Optional<QueueElement>> VisitQueue;
  62. // Current level.
  63. unsigned Level = 0;
  64. inline bf_iterator(NodeRef Node) {
  65. this->Visited.insert(Node);
  66. Level = 0;
  67. // Also, insert a dummy node as marker.
  68. VisitQueue.push(QueueElement(Node, None));
  69. VisitQueue.push(None);
  70. }
  71. inline bf_iterator() = default;
  72. inline void toNext() {
  73. Optional<QueueElement> Head = VisitQueue.front();
  74. QueueElement H = Head.getValue();
  75. NodeRef Node = H.first;
  76. Optional<ChildItTy> &ChildIt = H.second;
  77. if (!ChildIt)
  78. ChildIt.emplace(GT::child_begin(Node));
  79. while (*ChildIt != GT::child_end(Node)) {
  80. NodeRef Next = *(*ChildIt)++;
  81. // Already visited?
  82. if (this->Visited.insert(Next).second)
  83. VisitQueue.push(QueueElement(Next, None));
  84. }
  85. VisitQueue.pop();
  86. // Go to the next element skipping markers if needed.
  87. if (!VisitQueue.empty()) {
  88. Head = VisitQueue.front();
  89. if (Head != None)
  90. return;
  91. Level += 1;
  92. VisitQueue.pop();
  93. // Don't push another marker if this is the last
  94. // element.
  95. if (!VisitQueue.empty())
  96. VisitQueue.push(None);
  97. }
  98. }
  99. public:
  100. // Provide static begin and end methods as our public "constructors"
  101. static bf_iterator begin(const GraphT &G) {
  102. return bf_iterator(GT::getEntryNode(G));
  103. }
  104. static bf_iterator end(const GraphT &G) { return bf_iterator(); }
  105. bool operator==(const bf_iterator &RHS) const {
  106. return VisitQueue == RHS.VisitQueue;
  107. }
  108. bool operator!=(const bf_iterator &RHS) const { return !(*this == RHS); }
  109. const NodeRef &operator*() const { return VisitQueue.front()->first; }
  110. // This is a nonstandard operator-> that dereferences the pointer an extra
  111. // time so that you can actually call methods on the node, because the
  112. // contained type is a pointer.
  113. NodeRef operator->() const { return **this; }
  114. bf_iterator &operator++() { // Pre-increment
  115. toNext();
  116. return *this;
  117. }
  118. bf_iterator operator++(int) { // Post-increment
  119. bf_iterator ItCopy = *this;
  120. ++*this;
  121. return ItCopy;
  122. }
  123. unsigned getLevel() const { return Level; }
  124. };
  125. // Provide global constructors that automatically figure out correct types.
  126. template <class T> bf_iterator<T> bf_begin(const T &G) {
  127. return bf_iterator<T>::begin(G);
  128. }
  129. template <class T> bf_iterator<T> bf_end(const T &G) {
  130. return bf_iterator<T>::end(G);
  131. }
  132. // Provide an accessor method to use them in range-based patterns.
  133. template <class T> iterator_range<bf_iterator<T>> breadth_first(const T &G) {
  134. return make_range(bf_begin(G), bf_end(G));
  135. }
  136. } // end namespace llvm
  137. #endif // LLVM_ADT_BREADTHFIRSTITERATOR_H
  138. #ifdef __GNUC__
  139. #pragma GCC diagnostic pop
  140. #endif