Introduce a CallGraph updater helper class

The CallGraphUpdater is a helper that simplifies the process of updating
the call graph, both old and new style, while running an CGSCC pass.

The uses are contained in different commits, e.g. D70767.

More functionality is added as we need it.

Reviewed By: modocache, hfinkel

Differential Revision: https://reviews.llvm.org/D70927

Author: jdoerfert

llvm/include/llvm/Analysis/CallGraph.h

@@ -94,10 +94,6 @@ class CallGraph {
   /// callers from the old function to the new.
   void spliceFunction(const Function *From, const Function *To);
 
-  /// Add a function to the call graph, and link the node to all of the
-  /// functions that it calls.
-  void addToCallGraph(Function *F);
-
 public:
   explicit CallGraph(Module &M);
   CallGraph(CallGraph &&Arg);
@@ -158,6 +154,13 @@ class CallGraph {
   /// Similar to operator[], but this will insert a new CallGraphNode for
   /// \c F if one does not already exist.
   CallGraphNode *getOrInsertFunction(const Function *F);
+
+  /// Populate \p CGN based on the calls inside the associated function.
+  void populateCallGraphNode(CallGraphNode *CGN);
+
+  /// Add a function to the call graph, and link the node to all of the
+  /// functions that it calls.
+  void addToCallGraph(Function *F);
 };
 
 /// A node in the call graph for a module.

llvm/include/llvm/Analysis/LazyCallGraph.h

@@ -1058,6 +1058,9 @@ class LazyCallGraph {
   /// fully visited by the DFS prior to calling this routine.
   void removeDeadFunction(Function &F);
 
+  /// Introduce a node for the function \p NewF in the SCC \p C.
+  void addNewFunctionIntoSCC(Function &NewF, SCC &C);
+
   ///@}
 
   ///@{

llvm/include/llvm/Transforms/Utils/CallGraphUpdater.h

@@ -0,0 +1,106 @@
+//===- CallGraphUpdater.h - A (lazy) call graph update helper ---*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+/// \file
+///
+/// This file provides interfaces used to manipulate a call graph, regardless
+/// if it is a "old style" CallGraph or an "new style" LazyCallGraph.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_UTILS_CALLGRAPHUPDATER_H
+#define LLVM_TRANSFORMS_UTILS_CALLGRAPHUPDATER_H
+
+#include "llvm/Analysis/CGSCCPassManager.h"
+#include "llvm/Analysis/CallGraph.h"
+#include "llvm/Analysis/CallGraphSCCPass.h"
+#include "llvm/Analysis/LazyCallGraph.h"
+
+namespace llvm {
+
+/// Wrapper to unify "old style" CallGraph and "new style" LazyCallGraph. This
+/// simplifies the interface and the call sites, e.g., new and old pass manager
+/// passes can share the same code.
+class CallGraphUpdater {
+  /// Containers for functions which we did replace or want to delete when
+  /// `finalize` is called. This can happen explicitly or as part of the
+  /// destructor. Dead functions in comdat sections are tracked seperatly
+  /// because a function with discardable linakage in a COMDAT should only
+  /// be dropped if the entire COMDAT is dropped, see git ac07703842cf.
+  ///{
+  SmallPtrSet<Function *, 16> ReplacedFunctions;
+  SmallVector<Function *, 16> DeadFunctions;
+  SmallVector<Function *, 16> DeadFunctionsInComdats;
+  ///}
+
+  /// Old PM variables
+  ///{
+  CallGraph *CG = nullptr;
+  CallGraphSCC *CGSCC = nullptr;
+  ///}
+
+  /// New PM variables
+  ///{
+  LazyCallGraph *LCG = nullptr;
+  LazyCallGraph::SCC *SCC = nullptr;
+  CGSCCAnalysisManager *AM = nullptr;
+  CGSCCUpdateResult *UR = nullptr;
+  ///}
+
+public:
+  CallGraphUpdater() {}
+  ~CallGraphUpdater() { finalize(); }
+
+  /// Initializers for usage outside of a CGSCC pass, inside a CGSCC pass in
+  /// the old and new pass manager (PM).
+  ///{
+  void initialize(CallGraph &CG, CallGraphSCC &SCC) {
+    this->CG = &CG;
+    this->CGSCC = &SCC;
+  }
+  void initialize(LazyCallGraph &LCG, LazyCallGraph::SCC &SCC,
+                  CGSCCAnalysisManager &AM, CGSCCUpdateResult &UR) {
+    this->LCG = &LCG;
+    this->SCC = &SCC;
+    this->AM = &AM;
+    this->UR = &UR;
+  }
+  ///}
+
+  /// Finalizer that will trigger actions like function removal from the CG.
+  bool finalize();
+
+  /// Remove \p Fn from the call graph.
+  void removeFunction(Function &Fn);
+
+  /// After an CGSCC pass changes a function in ways that affect the call
+  /// graph, this method can be called to update it.
+  void reanalyzeFunction(Function &Fn);
+
+  /// If a new function was created by outlining, this method can be called
+  /// to update the call graph for the new function. Note that the old one
+  /// still needs to be re-analyzed or manually updated.
+  void registerOutlinedFunction(Function &NewFn);
+
+  /// Replace \p OldFn in the call graph (and SCC) with \p NewFn. The uses
+  /// outside the call graph and the function \p OldFn are not modified.
+  /// Note that \p OldFn is also removed from the call graph
+  /// (\see removeFunction).
+  void replaceFunctionWith(Function &OldFn, Function &NewFn);
+
+  /// Remove the call site \p CS from the call graph.
+  void removeCallSite(CallBase &CS);
+
+  /// Replace \p OldCS with the new call site \p NewCS.
+  /// \return True if the replacement was successful, otherwise False. In the
+  /// latter case the parent function of \p OldCB needs to be re-analyzed.
+  bool replaceCallSite(CallBase &OldCS, CallBase &NewCS);
+};
+
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_UTILS_CALLGRAPHUPDATER_H

llvm/lib/Analysis/CallGraph.cpp

@@ -74,6 +74,12 @@ void CallGraph::addToCallGraph(Function *F) {
   if (!F->hasLocalLinkage() || F->hasAddressTaken())
     ExternalCallingNode->addCalledFunction(nullptr, Node);
 
+  populateCallGraphNode(Node);
+}
+
+void CallGraph::populateCallGraphNode(CallGraphNode *Node) {
+  Function *F = Node->getFunction();
+
   // If this function is not defined in this translation unit, it could call
   // anything.
   if (F->isDeclaration() && !F->isIntrinsic())

llvm/lib/Analysis/CallGraphSCCPass.cpp

@@ -549,7 +549,10 @@ void CallGraphSCC::ReplaceNode(CallGraphNode *Old, CallGraphNode *New) {
   for (unsigned i = 0; ; ++i) {
     assert(i != Nodes.size() && "Node not in SCC");
     if (Nodes[i] != Old) continue;
-    Nodes[i] = New;
+    if (New)
+      Nodes[i] = New;
+    else
+      Nodes.erase(Nodes.begin() + i);
     break;
   }
 

llvm/lib/Analysis/LazyCallGraph.cpp

@@ -1566,6 +1566,15 @@ void LazyCallGraph::removeDeadFunction(Function &F) {
   // allocators.
 }
 
+void LazyCallGraph::addNewFunctionIntoSCC(Function &NewF, SCC &C) {
+  Node &CGNode = get(NewF);
+  CGNode.DFSNumber = CGNode.LowLink = -1;
+  CGNode.populate();
+  C.Nodes.push_back(&CGNode);
+  SCCMap[&CGNode] = &C;
+  NodeMap[&NewF] = &CGNode;
+}
+
 LazyCallGraph::Node &LazyCallGraph::insertInto(Function &F, Node *&MappedN) {
   return *new (MappedN = BPA.Allocate()) Node(*this, F);
 }

llvm/lib/Transforms/Utils/CMakeLists.txt

@@ -7,6 +7,7 @@ add_llvm_component_library(LLVMTransformUtils
   BuildLibCalls.cpp
   BypassSlowDivision.cpp
   CallPromotionUtils.cpp
+  CallGraphUpdater.cpp
   CanonicalizeAliases.cpp
   CloneFunction.cpp
   CloneModule.cpp

llvm/lib/Transforms/Utils/CallGraphUpdater.cpp

@@ -0,0 +1,152 @@
+//===- CallGraphUpdater.cpp - A (lazy) call graph update helper -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+/// \file
+///
+/// This file provides interfaces used to manipulate a call graph, regardless
+/// if it is a "old style" CallGraph or an "new style" LazyCallGraph.
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/CallGraphUpdater.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
+
+using namespace llvm;
+
+bool CallGraphUpdater::finalize() {
+  if (!DeadFunctionsInComdats.empty()) {
+    filterDeadComdatFunctions(*DeadFunctionsInComdats.front()->getParent(),
+                              DeadFunctionsInComdats);
+    DeadFunctions.append(DeadFunctionsInComdats.begin(),
+                         DeadFunctionsInComdats.end());
+  }
+
+  for (Function *DeadFn : DeadFunctions) {
+    DeadFn->removeDeadConstantUsers();
+
+    if (CG) {
+      CallGraphNode *OldCGN = CG->getOrInsertFunction(DeadFn);
+      CG->getExternalCallingNode()->removeAnyCallEdgeTo(OldCGN);
+      OldCGN->removeAllCalledFunctions();
+      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:
+      FunctionAnalysisManager &FAM =
+          AM->getResult<FunctionAnalysisManagerCGSCCProxy>(*SCC, *LCG)
+              .getManager();
+
+      LazyCallGraph::Node &N = LCG->get(*DeadFn);
+      auto *DeadSCC = LCG->lookupSCC(N);
+      assert(DeadSCC && DeadSCC->size() == 1 &&
+             &DeadSCC->begin()->getFunction() == DeadFn);
+      auto &DeadRC = DeadSCC->getOuterRefSCC();
+
+      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();
+  DeadFunctionsInComdats.clear();
+  DeadFunctions.clear();
+  return Changed;
+}
+
+void CallGraphUpdater::reanalyzeFunction(Function &Fn) {
+  if (CG) {
+    CallGraphNode *OldCGN = CG->getOrInsertFunction(&Fn);
+    OldCGN->removeAllCalledFunctions();
+    CG->populateCallGraphNode(OldCGN);
+  } else if (LCG) {
+    LazyCallGraph::Node &N = LCG->get(Fn);
+    LazyCallGraph::SCC *C = LCG->lookupSCC(N);
+    updateCGAndAnalysisManagerForCGSCCPass(*LCG, *C, N, *AM, *UR);
+  }
+}
+
+void CallGraphUpdater::registerOutlinedFunction(Function &NewFn) {
+  if (CG)
+    CG->addToCallGraph(&NewFn);
+  else if (LCG)
+    LCG->addNewFunctionIntoSCC(NewFn, *SCC);
+}
+
+void CallGraphUpdater::removeFunction(Function &DeadFn) {
+  DeadFn.deleteBody();
+  DeadFn.setLinkage(GlobalValue::ExternalLinkage);
+  if (DeadFn.hasComdat())
+    DeadFunctionsInComdats.push_back(&DeadFn);
+  else
+    DeadFunctions.push_back(&DeadFn);
+}
+
+void CallGraphUpdater::replaceFunctionWith(Function &OldFn, Function &NewFn) {
+  ReplacedFunctions.insert(&OldFn);
+  if (CG) {
+    // Update the call graph for the newly promoted function.
+    // CG->spliceFunction(&OldFn, &NewFn);
+    CallGraphNode *OldCGN = (*CG)[&OldFn];
+    CallGraphNode *NewCGN = CG->getOrInsertFunction(&NewFn);
+    NewCGN->stealCalledFunctionsFrom(OldCGN);
+
+    // And update the SCC we're iterating as well.
+    CGSCC->ReplaceNode(OldCGN, NewCGN);
+  } else if (LCG) {
+    // Directly substitute the functions in the call graph.
+    LazyCallGraph::Node &OldLCGN = LCG->get(OldFn);
+    SCC->getOuterRefSCC().replaceNodeFunction(OldLCGN, NewFn);
+  }
+  removeFunction(OldFn);
+}
+
+bool CallGraphUpdater::replaceCallSite(CallBase &OldCS, CallBase &NewCS) {
+  // This is only necessary in the (old) CG.
+  if (!CG)
+    return true;
+
+  Function *Caller = OldCS.getCaller();
+  CallGraphNode *NewCalleeNode =
+      CG->getOrInsertFunction(NewCS.getCalledFunction());
+  CallGraphNode *CallerNode = (*CG)[Caller];
+  if (llvm::none_of(*CallerNode, [&OldCS](const CallGraphNode::CallRecord &CR) {
+        return CR.first == &OldCS;
+      }))
+    return false;
+  CallerNode->replaceCallEdge(OldCS, NewCS, NewCalleeNode);
+  return true;
+}
+
+void CallGraphUpdater::removeCallSite(CallBase &CS) {
+  // This is only necessary in the (old) CG.
+  if (!CG)
+    return;
+
+  Function *Caller = CS.getCaller();
+  CallGraphNode *CallerNode = (*CG)[Caller];
+  CallerNode->removeCallEdgeFor(CS);
+}