[SCCIterator] Fix an issue in scc_member_iterator sorting

[hoy]

Summary

Members in an scc are supposed to be sorted in a top-down or topological order based on edge weights. Previously this is achived by building a MST out of the SCC and enforcing an BFS walk on the MST. A BFS on a tree does give a top-down topological order, however, the MST built here isn't really a tree. This is becuase of a trick done to avoid expansive detection of a cycle on a directed graph when an edge is added. When the MST is built, its edges are considered undirected. But in reality they are directed, thus a BST walk doesn't necessarily give a topological order. I'm tweaking the BFS walk slightly to yield a topological order.

Basically I'm using Kahn's algorithm on MST to compute a topological traversal order. The algorithm starts from nodes that have no incoming edge. These nodes are "roots" of the MST forest. This ensures that nodes are visited before their descendants are, thus ensures a topological traversal order of the MST.

[hoy]

Updating D130717: [SCCIterator] Fix an issue in scc_member_iterator sorting

[wenlei]

I'm tweaking the BFS walk slightly to yield a topological order.

Would be good to spell out the tweak in the description.

llvm/include/llvm/ADT/SCCIterator.h

346–349

nit: this comment seems slightly misplaced? It describes the initialization of the queue on line 353. And slightly duplicated with the comment on line 354. Maybe move there and merge the comments.

370

Do we need to check the visited flag still? If not, we can remove the flag?

[hoy]

llvm/include/llvm/ADT/SCCIterator.h

346–349

The comment was supposed to explain a bit on the Kahn's algorithm. The next comment was mainly for getting a deterministic walk. I just tweaked them a little bit. Let me know if it looks good now.

370

Good point. It's used any more up since here but used when initializing the queue. I'm removing the setting of it in this loop.

[hoy]

Addressing comments.

[wenlei]

lgtm, thanks.

[nikic]

Shouldn't there be a test for this change?

[hoy]

In D130717#4229609, @nikic wrote:

Shouldn't there be a test for this change?

Good point. I'll add a test.

llvm/include/llvm/ADT/SCCIterator.h

 /// This requires some simple wrappers and is not supported yet.)
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ADT_SCCITERATOR_H
 #define LLVM_ADT_SCCITERATOR_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/ADT/iterator.h"
 #include <cassert>
 #include <cstddef>
 #include <iterator>
 #include <queue>
 #include <set>
 #include <unordered_map>
class scc_member_iterator {
   using NodeType = typename GT::NodeType;
   using EdgeType = typename GT::EdgeType;
   using NodesType = std::vector<NodeType *>;
 
   // Auxilary node information used during the MST calculation.
   struct NodeInfo {
     NodeInfo *Group = this;
     uint32_t Rank = 0;
-    bool Visited = true;
+    bool Visited = false;
+    DenseSet<const EdgeType *> IncomingMSTEdges;
   };
 
   // Find the root group of the node and compress the path from node to the
   // root.
   NodeInfo *find(NodeInfo *Node) {
     if (Node->Group != Node)
       Node->Group = find(Node->Group);
     return Node->Group;
scc_member_iterator<GraphT, GT>::scc_member_iterator(
   // Traverse all the edges and compute the Maximum Weight Spanning Tree
   // using Kruskal's algorithm.
   std::unordered_set<const EdgeType *> MSTEdges;
   for (auto *Edge : SortedEdges) {
     if (unionGroups(Edge))
       MSTEdges.insert(Edge);
   }
 
-  // Do BFS on MST, starting from nodes that have no incoming edge. These nodes
-  // are "roots" of the MST forest. This ensures that nodes are visited before
-  // their decsendents are, thus ensures hot edges are processed before cold
-  // edges, based on how MST is computed.
+  // Run Kahn's algorithm on MST to compute a topological traversal order.
+  // The algorithm starts from nodes that have no incoming edge. These nodes are
+  // "roots" of the MST forest. This ensures that nodes are visited before their
+  // descendants are, thus ensures hot edges are processed before cold edges,
+  // based on how MST is computed.
+  std::queue<NodeType *> Queue;
   for (const auto *Edge : MSTEdges)
-    NodeInfoMap[Edge->Target].Visited = false;
+    NodeInfoMap[Edge->Target].IncomingMSTEdges.insert(Edge);
 
-  std::queue<NodeType *> Queue;
-  // Initialze the queue with MST roots. Note that walking through SortedEdges
-  // instead of NodeInfoMap ensures an ordered deterministic push.
+  // Walk through SortedEdges to initialize the queue, instead of using NodeInfoMap
+  // to ensure an ordered deterministic push.
   for (auto *Edge : SortedEdges) {
-    if (NodeInfoMap[Edge->Source].Visited) {
+    if (!NodeInfoMap[Edge->Source].Visited &&
+        NodeInfoMap[Edge->Source].IncomingMSTEdges.empty()) {
       Queue.push(Edge->Source);
-      NodeInfoMap[Edge->Source].Visited = false;
+      NodeInfoMap[Edge->Source].Visited = true;
     }
   }
 
   while (!Queue.empty()) {
     auto *Node = Queue.front();
     Queue.pop();
     Nodes.push_back(Node);
     for (auto &Edge : Node->Edges) {
-      if (MSTEdges.count(&Edge) && !NodeInfoMap[Edge.Target].Visited) {
-        NodeInfoMap[Edge.Target].Visited = true;
+      NodeInfoMap[Edge.Target].IncomingMSTEdges.erase(&Edge);
+      if (MSTEdges.count(&Edge) &&
+          NodeInfoMap[Edge.Target].IncomingMSTEdges.empty()) {
         Queue.push(Edge.Target);
       }
     }
   }
 
   assert(InputNodes.size() == Nodes.size() && "missing nodes in MST");
   std::reverse(Nodes.begin(), Nodes.end());
 }
 } // end namespace llvm
 
 #endif // LLVM_ADT_SCCITERATOR_H