BreadthFirstIterator.h 5.0 KB

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