[LCG] Rotate logic applied to the top of the DFSStack to instead be

applied prior to pushing a node onto the DFSStack. This is the first
step toward avoiding the stack entirely for leaf nodes. It also
simplifies things a bit and I think is pointing the way toward factoring
some more of the shared logic out of the two implementations.

It is also making it more obvious how to restructure the loops
themselves to be a bit easier to read (although no different in terms of
functionality).

llvm-svn: 207095

llvm/lib/Analysis/LazyCallGraph.cpp

@@ -183,7 +183,7 @@ LazyCallGraph::SCC::removeInternalEdge(LazyCallGraph &G, Node &Caller,
   ResultSCCs.push_back(this);
 
   // We're going to do a full mini-Tarjan's walk using a local stack here.
-  int NextDFSNumber = 1;
+  int NextDFSNumber;
   SmallVector<std::pair<Node *, Node::iterator>, 4> DFSStack;
 
   // The worklist is every node in the original SCC. FIXME: switch the SCC to
@@ -210,24 +210,15 @@ LazyCallGraph::SCC::removeInternalEdge(LazyCallGraph &G, Node &Caller,
       if (Worklist.empty())
         break;
       Node *N = Worklist.pop_back_val();
+      N->LowLink = N->DFSNumber = 1;
+      NextDFSNumber = 2;
       DFSStack.push_back(std::make_pair(N, N->begin()));
     }
 
     Node *N = DFSStack.back().first;
 
-    // Check if we have reached a node in the new (known connected) set. If so,
+    assert(N->DFSNumber != 0 && "We should always assign a DFS number "
-    // the entire stack is necessarily in that set and we can re-start.
+                                "before placing a node onto the stack.");
-    if (NewNodes.count(N)) {
-      DFSStack.pop_back();
-      while (!DFSStack.empty())
-        NewNodes.insert(DFSStack.pop_back_val().first);
-      continue;
-    }
-
-    if (N->DFSNumber == 0) {
-      N->LowLink = N->DFSNumber = NextDFSNumber++;
-      Worklist.remove(N);
-    }
 
     auto SI = DFSStack.rbegin();
     bool PushedChildNode = false;
@@ -243,6 +234,14 @@ LazyCallGraph::SCC::removeInternalEdge(LazyCallGraph &G, Node &Caller,
           continue;
         }
 
+        // Check if we have reached a node in the new (known connected) set. If
+        // so, the entire stack is necessarily in that set and we can re-start.
+        if (NewNodes.count(&ChildN)) {
+          while (!DFSStack.empty())
+            NewNodes.insert(DFSStack.pop_back_val().first);
+          continue;
+        }
+
         if (ChildN.DFSNumber == 0) {
           // Mark that we should start at this child when next this node is the
           // top of the stack. We don't start at the next child to ensure this
@@ -250,6 +249,8 @@ LazyCallGraph::SCC::removeInternalEdge(LazyCallGraph &G, Node &Caller,
           SI->second = I;
 
           // Recurse onto this node via a tail call.
+          ChildN.LowLink = ChildN.DFSNumber = NextDFSNumber++;
+          Worklist.remove(&ChildN);
           DFSStack.push_back(std::make_pair(&ChildN, ChildN.begin()));
           PushedChildNode = true;
           break;
@@ -272,7 +273,7 @@ LazyCallGraph::SCC::removeInternalEdge(LazyCallGraph &G, Node &Caller,
     } while (!PushedChildNode && N->LowLink != N->DFSNumber &&
              SI != DFSStack.rend());
 
-    if (PushedChildNode)
+    if (DFSStack.empty() || PushedChildNode)
       continue;
 
     // Form the new SCC out of the top of the DFS stack.
@@ -421,21 +422,15 @@ LazyCallGraph::SCC *LazyCallGraph::getNextSCCInPostOrder() {
     if (SCCEntryNodes.empty())
       return nullptr;
 
-    // Reset the DFS numbering.
-    NextDFSNumber = 1;
     Node &N = get(*SCCEntryNodes.pop_back_val());
+    N.LowLink = N.DFSNumber = 1;
+    NextDFSNumber = 2;
     DFSStack.push_back(std::make_pair(&N, N.begin()));
   }
 
   auto SI = DFSStack.rbegin();
-  if (SI->first->DFSNumber == 0) {
+  assert(SI->first->DFSNumber != 0 && "We should always assign a DFS number "
-    // This node hasn't been visited before, assign it a DFS number and remove
+                                      "before placing a node onto the stack.");
-    // it from the entry set.
-    assert(!SCCMap.count(SI->first) &&
-           "Found a node with 0 DFS number but already in an SCC!");
-    SI->first->LowLink = SI->first->DFSNumber = NextDFSNumber++;
-    SCCEntryNodes.remove(&SI->first->getFunction());
-  }
 
   do {
     Node *N = SI->first;
@@ -448,6 +443,10 @@ LazyCallGraph::SCC *LazyCallGraph::getNextSCCInPostOrder() {
         SI->second = I;
 
         // Recurse onto this node via a tail call.
+        assert(!SCCMap.count(&ChildN) &&
+               "Found a node with 0 DFS number but already in an SCC!");
+        ChildN.LowLink = ChildN.DFSNumber = NextDFSNumber++;
+        SCCEntryNodes.remove(&ChildN.getFunction());
         DFSStack.push_back(std::make_pair(&ChildN, ChildN.begin()));
         return LazyCallGraph::getNextSCCInPostOrder();
       }