1 files changed, 45 insertions, 38 deletions

diff --git a/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp b/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp
index 4de7124..00c6007 100644
--- a/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp
+++ b/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp
@@ -26,49 +26,56 @@ bool CallGraphUpdater::finalize() {
                          DeadFunctionsInComdats.end());
   }
 
-  for (Function *DeadFn : DeadFunctions) {
-    DeadFn->removeDeadConstantUsers();
-
-    if (CG) {
-      CallGraphNode *OldCGN = CG->getOrInsertFunction(DeadFn);
-      CG->getExternalCallingNode()->removeAnyCallEdgeTo(OldCGN);
-      OldCGN->removeAllCalledFunctions();
+  if (CG) {
+    // First remove all references, e.g., outgoing via called functions. This is
+    // necessary as we can delete functions that have circular references.
+    for (Function *DeadFn : DeadFunctions) {
+      DeadFn->removeDeadConstantUsers();
+      CallGraphNode *DeadCGN = (*CG)[DeadFn];
+      DeadCGN->removeAllCalledFunctions();
+      CG->getExternalCallingNode()->removeAnyCallEdgeTo(DeadCGN);
       DeadFn->replaceAllUsesWith(UndefValue::get(DeadFn->getType()));
-
-      assert(OldCGN->getNumReferences() == 0);
-
-      delete CG->removeFunctionFromModule(OldCGN);
-      continue;
     }
 
-    // The old style call graph (CG) has a value handle we do not want to
-    // replace with undef so we do this here.
-    DeadFn->replaceAllUsesWith(UndefValue::get(DeadFn->getType()));
-
-    if (LCG && !ReplacedFunctions.count(DeadFn)) {
-      // Taken mostly from the inliner:
-      LazyCallGraph::Node &N = LCG->get(*DeadFn);
-      auto *DeadSCC = LCG->lookupSCC(N);
-      assert(DeadSCC && DeadSCC->size() == 1 &&
-             &DeadSCC->begin()->getFunction() == DeadFn);
-      auto &DeadRC = DeadSCC->getOuterRefSCC();
-
-      FunctionAnalysisManager &FAM =
-          AM->getResult<FunctionAnalysisManagerCGSCCProxy>(*DeadSCC, *LCG)
-              .getManager();
-
-      FAM.clear(*DeadFn, DeadFn->getName());
-      AM->clear(*DeadSCC, DeadSCC->getName());
-      LCG->removeDeadFunction(*DeadFn);
-
-      // Mark the relevant parts of the call graph as invalid so we don't visit
-      // them.
-      UR->InvalidatedSCCs.insert(DeadSCC);
-      UR->InvalidatedRefSCCs.insert(&DeadRC);
+    // Then remove the node and function from the module.
+    for (Function *DeadFn : DeadFunctions) {
+      CallGraphNode *DeadCGN = CG->getOrInsertFunction(DeadFn);
+      assert(DeadCGN->getNumReferences() == 0 &&
+             "References should have been handled by now");
+      delete CG->removeFunctionFromModule(DeadCGN);
     }
+  } else {
+    // This is the code path for the new lazy call graph and for the case were
+    // no call graph was provided.
+    for (Function *DeadFn : DeadFunctions) {
+      DeadFn->removeDeadConstantUsers();
+      DeadFn->replaceAllUsesWith(UndefValue::get(DeadFn->getType()));
 
-    // The function is now really dead and de-attached from everything.
-    DeadFn->eraseFromParent();
+      if (LCG && !ReplacedFunctions.count(DeadFn)) {
+        // Taken mostly from the inliner:
+        LazyCallGraph::Node &N = LCG->get(*DeadFn);
+        auto *DeadSCC = LCG->lookupSCC(N);
+        assert(DeadSCC && DeadSCC->size() == 1 &&
+               &DeadSCC->begin()->getFunction() == DeadFn);
+        auto &DeadRC = DeadSCC->getOuterRefSCC();
+
+        FunctionAnalysisManager &FAM =
+            AM->getResult<FunctionAnalysisManagerCGSCCProxy>(*DeadSCC, *LCG)
+                .getManager();
+
+        FAM.clear(*DeadFn, DeadFn->getName());
+        AM->clear(*DeadSCC, DeadSCC->getName());
+        LCG->removeDeadFunction(*DeadFn);
+
+        // Mark the relevant parts of the call graph as invalid so we don't
+        // visit them.
+        UR->InvalidatedSCCs.insert(DeadSCC);
+        UR->InvalidatedRefSCCs.insert(&DeadRC);
+      }
+
+      // The function is now really dead and de-attached from everything.
+      DeadFn->eraseFromParent();
+    }
   }
 
   bool Changed = !DeadFunctions.empty();