[Coroutines] [LazyCallGraph] Make all new functions reference eachother.

Author: tnowicki

llvm/include/llvm/Analysis/LazyCallGraph.h

@@ -1081,12 +1081,12 @@ class LazyCallGraph {
 
   /// Add new ref-recursive functions split/outlined from an existing function.
   ///
-  /// The new functions may only reference other functions that the original
-  /// function did. The new functions may reference (not call) the original
-  /// function.
+  /// The new functions may only reference the original function or other
+  /// functions that the original function did. New functions must not call
+  /// other new functions.
   ///
-  /// The original function must reference (not call) all new functions.
-  /// All new functions must reference (not call) each other.
+  /// Mark the original function as referencing all new functions.
+  /// Mark all new functions as referencing each other.
   void addSplitRefRecursiveFunctions(Function &OriginalFunction,
                                      ArrayRef<Function *> NewFunctions);
 

llvm/lib/Analysis/LazyCallGraph.cpp

@@ -1720,6 +1720,7 @@ void LazyCallGraph::addSplitRefRecursiveFunctions(
   for (Function *NewFunction : NewFunctions) {
     Node &NewN = initNode(*NewFunction);
 
+    // Make the original function reference each new function
     OriginalN->insertEdgeInternal(NewN, Edge::Kind::Ref);
 
     // Check if there is any edge from any new function back to any function in
@@ -1732,6 +1733,23 @@ void LazyCallGraph::addSplitRefRecursiveFunctions(
     }
   }
 
+  for (Function *NewFunction : NewFunctions) {
+    Node &NewN = get(*NewFunction);
+    for (Function *OtherNewFunction : NewFunctions) {
+      if (NewFunction == OtherNewFunction)
+        continue;
+
+      Node &OtherNewN = get(*OtherNewFunction);
+
+      // Don't add an edge if one already exists.
+      if (NewN->lookup(OtherNewN))
+        continue;
+
+      // Make the new function reference each other new function
+      NewN->insertEdgeInternal(OtherNewN, Edge::Kind::Ref);
+    }
+  }
+
   RefSCC *NewRC;
   if (ExistsRefToOriginalRefSCC) {
     // If there is any edge from any new function to any function in the
@@ -1775,7 +1793,7 @@ void LazyCallGraph::addSplitRefRecursiveFunctions(
       if (F1 == F2)
         continue;
       Node &N2 = get(*F2);
-      assert(!N1->lookup(N2) ||
+      assert(N1->lookup(N2) &&
              (!N1->lookup(N2)->isCall() &&
               "Edges between new functions must be ref edges"));
     }