SCCIterator.h 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383
  1. #pragma once
  2. #ifdef __GNUC__
  3. #pragma GCC diagnostic push
  4. #pragma GCC diagnostic ignored "-Wunused-parameter"
  5. #endif
  6. //===- ADT/SCCIterator.h - Strongly Connected Comp. Iter. -------*- 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. /// \file
  14. ///
  15. /// This builds on the llvm/ADT/GraphTraits.h file to find the strongly
  16. /// connected components (SCCs) of a graph in O(N+E) time using Tarjan's DFS
  17. /// algorithm.
  18. ///
  19. /// The SCC iterator has the important property that if a node in SCC S1 has an
  20. /// edge to a node in SCC S2, then it visits S1 *after* S2.
  21. ///
  22. /// To visit S1 *before* S2, use the scc_iterator on the Inverse graph. (NOTE:
  23. /// This requires some simple wrappers and is not supported yet.)
  24. ///
  25. //===----------------------------------------------------------------------===//
  26. #ifndef LLVM_ADT_SCCITERATOR_H
  27. #define LLVM_ADT_SCCITERATOR_H
  28. #include "llvm/ADT/DenseMap.h"
  29. #include "llvm/ADT/GraphTraits.h"
  30. #include "llvm/ADT/iterator.h"
  31. #include <cassert>
  32. #include <cstddef>
  33. #include <iterator>
  34. #include <queue>
  35. #include <set>
  36. #include <unordered_map>
  37. #include <unordered_set>
  38. #include <vector>
  39. namespace llvm {
  40. /// Enumerate the SCCs of a directed graph in reverse topological order
  41. /// of the SCC DAG.
  42. ///
  43. /// This is implemented using Tarjan's DFS algorithm using an internal stack to
  44. /// build up a vector of nodes in a particular SCC. Note that it is a forward
  45. /// iterator and thus you cannot backtrack or re-visit nodes.
  46. template <class GraphT, class GT = GraphTraits<GraphT>>
  47. class scc_iterator : public iterator_facade_base<
  48. scc_iterator<GraphT, GT>, std::forward_iterator_tag,
  49. const std::vector<typename GT::NodeRef>, ptrdiff_t> {
  50. using NodeRef = typename GT::NodeRef;
  51. using ChildItTy = typename GT::ChildIteratorType;
  52. using SccTy = std::vector<NodeRef>;
  53. using reference = typename scc_iterator::reference;
  54. /// Element of VisitStack during DFS.
  55. struct StackElement {
  56. NodeRef Node; ///< The current node pointer.
  57. ChildItTy NextChild; ///< The next child, modified inplace during DFS.
  58. unsigned MinVisited; ///< Minimum uplink value of all children of Node.
  59. StackElement(NodeRef Node, const ChildItTy &Child, unsigned Min)
  60. : Node(Node), NextChild(Child), MinVisited(Min) {}
  61. bool operator==(const StackElement &Other) const {
  62. return Node == Other.Node &&
  63. NextChild == Other.NextChild &&
  64. MinVisited == Other.MinVisited;
  65. }
  66. };
  67. /// The visit counters used to detect when a complete SCC is on the stack.
  68. /// visitNum is the global counter.
  69. ///
  70. /// nodeVisitNumbers are per-node visit numbers, also used as DFS flags.
  71. unsigned visitNum;
  72. DenseMap<NodeRef, unsigned> nodeVisitNumbers;
  73. /// Stack holding nodes of the SCC.
  74. std::vector<NodeRef> SCCNodeStack;
  75. /// The current SCC, retrieved using operator*().
  76. SccTy CurrentSCC;
  77. /// DFS stack, Used to maintain the ordering. The top contains the current
  78. /// node, the next child to visit, and the minimum uplink value of all child
  79. std::vector<StackElement> VisitStack;
  80. /// A single "visit" within the non-recursive DFS traversal.
  81. void DFSVisitOne(NodeRef N);
  82. /// The stack-based DFS traversal; defined below.
  83. void DFSVisitChildren();
  84. /// Compute the next SCC using the DFS traversal.
  85. void GetNextSCC();
  86. scc_iterator(NodeRef entryN) : visitNum(0) {
  87. DFSVisitOne(entryN);
  88. GetNextSCC();
  89. }
  90. /// End is when the DFS stack is empty.
  91. scc_iterator() = default;
  92. public:
  93. static scc_iterator begin(const GraphT &G) {
  94. return scc_iterator(GT::getEntryNode(G));
  95. }
  96. static scc_iterator end(const GraphT &) { return scc_iterator(); }
  97. /// Direct loop termination test which is more efficient than
  98. /// comparison with \c end().
  99. bool isAtEnd() const {
  100. assert(!CurrentSCC.empty() || VisitStack.empty());
  101. return CurrentSCC.empty();
  102. }
  103. bool operator==(const scc_iterator &x) const {
  104. return VisitStack == x.VisitStack && CurrentSCC == x.CurrentSCC;
  105. }
  106. scc_iterator &operator++() {
  107. GetNextSCC();
  108. return *this;
  109. }
  110. reference operator*() const {
  111. assert(!CurrentSCC.empty() && "Dereferencing END SCC iterator!");
  112. return CurrentSCC;
  113. }
  114. /// Test if the current SCC has a cycle.
  115. ///
  116. /// If the SCC has more than one node, this is trivially true. If not, it may
  117. /// still contain a cycle if the node has an edge back to itself.
  118. bool hasCycle() const;
  119. /// This informs the \c scc_iterator that the specified \c Old node
  120. /// has been deleted, and \c New is to be used in its place.
  121. void ReplaceNode(NodeRef Old, NodeRef New) {
  122. assert(nodeVisitNumbers.count(Old) && "Old not in scc_iterator?");
  123. // Do the assignment in two steps, in case 'New' is not yet in the map, and
  124. // inserting it causes the map to grow.
  125. auto tempVal = nodeVisitNumbers[Old];
  126. nodeVisitNumbers[New] = tempVal;
  127. nodeVisitNumbers.erase(Old);
  128. }
  129. };
  130. template <class GraphT, class GT>
  131. void scc_iterator<GraphT, GT>::DFSVisitOne(NodeRef N) {
  132. ++visitNum;
  133. nodeVisitNumbers[N] = visitNum;
  134. SCCNodeStack.push_back(N);
  135. VisitStack.push_back(StackElement(N, GT::child_begin(N), visitNum));
  136. #if 0 // Enable if needed when debugging.
  137. dbgs() << "TarjanSCC: Node " << N <<
  138. " : visitNum = " << visitNum << "\n";
  139. #endif
  140. }
  141. template <class GraphT, class GT>
  142. void scc_iterator<GraphT, GT>::DFSVisitChildren() {
  143. assert(!VisitStack.empty());
  144. while (VisitStack.back().NextChild != GT::child_end(VisitStack.back().Node)) {
  145. // TOS has at least one more child so continue DFS
  146. NodeRef childN = *VisitStack.back().NextChild++;
  147. typename DenseMap<NodeRef, unsigned>::iterator Visited =
  148. nodeVisitNumbers.find(childN);
  149. if (Visited == nodeVisitNumbers.end()) {
  150. // this node has never been seen.
  151. DFSVisitOne(childN);
  152. continue;
  153. }
  154. unsigned childNum = Visited->second;
  155. if (VisitStack.back().MinVisited > childNum)
  156. VisitStack.back().MinVisited = childNum;
  157. }
  158. }
  159. template <class GraphT, class GT> void scc_iterator<GraphT, GT>::GetNextSCC() {
  160. CurrentSCC.clear(); // Prepare to compute the next SCC
  161. while (!VisitStack.empty()) {
  162. DFSVisitChildren();
  163. // Pop the leaf on top of the VisitStack.
  164. NodeRef visitingN = VisitStack.back().Node;
  165. unsigned minVisitNum = VisitStack.back().MinVisited;
  166. assert(VisitStack.back().NextChild == GT::child_end(visitingN));
  167. VisitStack.pop_back();
  168. // Propagate MinVisitNum to parent so we can detect the SCC starting node.
  169. if (!VisitStack.empty() && VisitStack.back().MinVisited > minVisitNum)
  170. VisitStack.back().MinVisited = minVisitNum;
  171. #if 0 // Enable if needed when debugging.
  172. dbgs() << "TarjanSCC: Popped node " << visitingN <<
  173. " : minVisitNum = " << minVisitNum << "; Node visit num = " <<
  174. nodeVisitNumbers[visitingN] << "\n";
  175. #endif
  176. if (minVisitNum != nodeVisitNumbers[visitingN])
  177. continue;
  178. // A full SCC is on the SCCNodeStack! It includes all nodes below
  179. // visitingN on the stack. Copy those nodes to CurrentSCC,
  180. // reset their minVisit values, and return (this suspends
  181. // the DFS traversal till the next ++).
  182. do {
  183. CurrentSCC.push_back(SCCNodeStack.back());
  184. SCCNodeStack.pop_back();
  185. nodeVisitNumbers[CurrentSCC.back()] = ~0U;
  186. } while (CurrentSCC.back() != visitingN);
  187. return;
  188. }
  189. }
  190. template <class GraphT, class GT>
  191. bool scc_iterator<GraphT, GT>::hasCycle() const {
  192. assert(!CurrentSCC.empty() && "Dereferencing END SCC iterator!");
  193. if (CurrentSCC.size() > 1)
  194. return true;
  195. NodeRef N = CurrentSCC.front();
  196. for (ChildItTy CI = GT::child_begin(N), CE = GT::child_end(N); CI != CE;
  197. ++CI)
  198. if (*CI == N)
  199. return true;
  200. return false;
  201. }
  202. /// Construct the begin iterator for a deduced graph type T.
  203. template <class T> scc_iterator<T> scc_begin(const T &G) {
  204. return scc_iterator<T>::begin(G);
  205. }
  206. /// Construct the end iterator for a deduced graph type T.
  207. template <class T> scc_iterator<T> scc_end(const T &G) {
  208. return scc_iterator<T>::end(G);
  209. }
  210. /// Sort the nodes of a directed SCC in the decreasing order of the edge
  211. /// weights. The instantiating GraphT type should have weighted edge type
  212. /// declared in its graph traits in order to use this iterator.
  213. ///
  214. /// This is implemented using Kruskal's minimal spanning tree algorithm followed
  215. /// by a BFS walk. First a maximum spanning tree (forest) is built based on all
  216. /// edges within the SCC collection. Then a BFS walk is initiated on tree nodes
  217. /// that do not have a predecessor. Finally, the BFS order computed is the
  218. /// traversal order of the nodes of the SCC. Such order ensures that
  219. /// high-weighted edges are visited first during the tranversal.
  220. template <class GraphT, class GT = GraphTraits<GraphT>>
  221. class scc_member_iterator {
  222. using NodeType = typename GT::NodeType;
  223. using EdgeType = typename GT::EdgeType;
  224. using NodesType = std::vector<NodeType *>;
  225. // Auxilary node information used during the MST calculation.
  226. struct NodeInfo {
  227. NodeInfo *Group = this;
  228. uint32_t Rank = 0;
  229. bool Visited = true;
  230. };
  231. // Find the root group of the node and compress the path from node to the
  232. // root.
  233. NodeInfo *find(NodeInfo *Node) {
  234. if (Node->Group != Node)
  235. Node->Group = find(Node->Group);
  236. return Node->Group;
  237. }
  238. // Union the source and target node into the same group and return true.
  239. // Returns false if they are already in the same group.
  240. bool unionGroups(const EdgeType *Edge) {
  241. NodeInfo *G1 = find(&NodeInfoMap[Edge->Source]);
  242. NodeInfo *G2 = find(&NodeInfoMap[Edge->Target]);
  243. // If the edge forms a cycle, do not add it to MST
  244. if (G1 == G2)
  245. return false;
  246. // Make the smaller rank tree a direct child or the root of high rank tree.
  247. if (G1->Rank < G1->Rank)
  248. G1->Group = G2;
  249. else {
  250. G2->Group = G1;
  251. // If the ranks are the same, increment root of one tree by one.
  252. if (G1->Rank == G2->Rank)
  253. G2->Rank++;
  254. }
  255. return true;
  256. }
  257. std::unordered_map<NodeType *, NodeInfo> NodeInfoMap;
  258. NodesType Nodes;
  259. public:
  260. scc_member_iterator(const NodesType &InputNodes);
  261. NodesType &operator*() { return Nodes; }
  262. };
  263. template <class GraphT, class GT>
  264. scc_member_iterator<GraphT, GT>::scc_member_iterator(
  265. const NodesType &InputNodes) {
  266. if (InputNodes.size() <= 1) {
  267. Nodes = InputNodes;
  268. return;
  269. }
  270. // Initialize auxilary node information.
  271. NodeInfoMap.clear();
  272. for (auto *Node : InputNodes) {
  273. // This is specifically used to construct a `NodeInfo` object in place. An
  274. // insert operation will involve a copy construction which invalidate the
  275. // initial value of the `Group` field which should be `this`.
  276. (void)NodeInfoMap[Node].Group;
  277. }
  278. // Sort edges by weights.
  279. struct EdgeComparer {
  280. bool operator()(const EdgeType *L, const EdgeType *R) const {
  281. return L->Weight > R->Weight;
  282. }
  283. };
  284. std::multiset<const EdgeType *, EdgeComparer> SortedEdges;
  285. for (auto *Node : InputNodes) {
  286. for (auto &Edge : Node->Edges) {
  287. if (NodeInfoMap.count(Edge.Target))
  288. SortedEdges.insert(&Edge);
  289. }
  290. }
  291. // Traverse all the edges and compute the Maximum Weight Spanning Tree
  292. // using Kruskal's algorithm.
  293. std::unordered_set<const EdgeType *> MSTEdges;
  294. for (auto *Edge : SortedEdges) {
  295. if (unionGroups(Edge))
  296. MSTEdges.insert(Edge);
  297. }
  298. // Do BFS on MST, starting from nodes that have no incoming edge. These nodes
  299. // are "roots" of the MST forest. This ensures that nodes are visited before
  300. // their decsendents are, thus ensures hot edges are processed before cold
  301. // edges, based on how MST is computed.
  302. for (const auto *Edge : MSTEdges)
  303. NodeInfoMap[Edge->Target].Visited = false;
  304. std::queue<NodeType *> Queue;
  305. for (auto &Node : NodeInfoMap)
  306. if (Node.second.Visited)
  307. Queue.push(Node.first);
  308. while (!Queue.empty()) {
  309. auto *Node = Queue.front();
  310. Queue.pop();
  311. Nodes.push_back(Node);
  312. for (auto &Edge : Node->Edges) {
  313. if (MSTEdges.count(&Edge) && !NodeInfoMap[Edge.Target].Visited) {
  314. NodeInfoMap[Edge.Target].Visited = true;
  315. Queue.push(Edge.Target);
  316. }
  317. }
  318. }
  319. assert(InputNodes.size() == Nodes.size() && "missing nodes in MST");
  320. std::reverse(Nodes.begin(), Nodes.end());
  321. }
  322. } // end namespace llvm
  323. #endif // LLVM_ADT_SCCITERATOR_H
  324. #ifdef __GNUC__
  325. #pragma GCC diagnostic pop
  326. #endif