From f637334df93f231460a12adc35cbf8c5e8cf35c5 Mon Sep 17 00:00:00 2001
From: Sergey Dmitriev <serguei.n.dmitriev@intel.com>
Date: Mon, 4 May 2020 18:58:59 -0700
Subject: [PATCH] [CallGraphUpdater] Removed references to calles when deleting
 function

Summary: Otherwise we can get unaccounted references to call graph nodes.

Reviewers: jdoerfert, sstefan1

Reviewed By: jdoerfert

Subscribers: hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D79382
---
 .../lib/Transforms/Utils/CallGraphUpdater.cpp |  7 +++--
 .../Util/cg-updater-dead-function-callees.ll  | 31 +++++++++++++++++++
 llvm/unittests/IR/LegacyPassManagerTest.cpp   |  6 ++--
 3 files changed, 39 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/Transforms/Util/cg-updater-dead-function-callees.ll

diff --git a/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp b/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp
index 8cdf01e80935..8060e50abee9 100644
--- a/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp
+++ b/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp
@@ -112,8 +112,11 @@ void CallGraphUpdater::removeFunction(Function &DeadFn) {
     DeadFunctions.push_back(&DeadFn);
 
   // For the old call graph we remove the function from the SCC right away.
-  if (CG && !ReplacedFunctions.count(&DeadFn))
-    CGSCC->DeleteNode((*CG)[&DeadFn]);
+  if (CG && !ReplacedFunctions.count(&DeadFn)) {
+    CallGraphNode *DeadCGN = (*CG)[&DeadFn];
+    DeadCGN->removeAllCalledFunctions();
+    CGSCC->DeleteNode(DeadCGN);
+  }
 }
 
 void CallGraphUpdater::replaceFunctionWith(Function &OldFn, Function &NewFn) {
diff --git a/llvm/test/Transforms/Util/cg-updater-dead-function-callees.ll b/llvm/test/Transforms/Util/cg-updater-dead-function-callees.ll
new file mode 100644
index 000000000000..8ae0b8e3c5d4
--- /dev/null
+++ b/llvm/test/Transforms/Util/cg-updater-dead-function-callees.ll
@@ -0,0 +1,31 @@
+; RUN: opt -inline -attributor-cgscc -tailcallelim -S %s | FileCheck %s
+;
+; CHECK: define void @foo()
+; CHECK: declare i32 @baz()
+; CHECK-NOT: void @goo()
+; CHECK-NOT: void @bar()
+
+define void @foo() {
+  call fastcc void @bar()
+  ret void
+}
+
+define internal fastcc void @goo() {
+  call fastcc void @bar()
+  ret void
+}
+
+define internal fastcc void @bar() {
+  %call = call i32 @baz()
+  %cond = icmp eq i32 %call, 0
+  br i1 %cond, label %if.then, label %if.end
+
+if.then:
+  call fastcc void @goo()
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+declare i32 @baz()
diff --git a/llvm/unittests/IR/LegacyPassManagerTest.cpp b/llvm/unittests/IR/LegacyPassManagerTest.cpp
index 1d4b3b835b69..b7801b52481d 100644
--- a/llvm/unittests/IR/LegacyPassManagerTest.cpp
+++ b/llvm/unittests/IR/LegacyPassManagerTest.cpp
@@ -687,9 +687,9 @@ namespace llvm {
       Passes.add(P);
       Passes.run(*M);
       ASSERT_EQ(P->SetupWorked, 1U);
-      ASSERT_EQ(P->NumSCCs, 5U);
-      ASSERT_EQ(P->NumFns, 8U);
-      ASSERT_EQ(P->NumFnDecls, 3U);
+      ASSERT_EQ(P->NumSCCs, 4U);
+      ASSERT_EQ(P->NumFns, 6U);
+      ASSERT_EQ(P->NumFnDecls, 1U);
       ASSERT_EQ(M->getFunctionList().size(), 3U);
       ASSERT_EQ(P->NumExtCalledBefore, /* test1, 2a, 2b, 3, 4 */ 5U);
       ASSERT_EQ(P->NumExtCalledAfter, /* test1, 3repl, 4 */ 3U);