3 files changed, 44 insertions, 5 deletions

diff --git a/llvm/lib/Analysis/CGSCCPassManager.cpp b/llvm/lib/Analysis/CGSCCPassManager.cpp
index a0b3f83cca6..54886d45a1d 100644
--- a/llvm/lib/Analysis/CGSCCPassManager.cpp
+++ b/llvm/lib/Analysis/CGSCCPassManager.cpp
@@ -449,7 +449,7 @@ LazyCallGraph::SCC &llvm::updateCGAndAnalysisManagerForFunctionPass(
   // irrelevant.
   for (Instruction &I : instructions(F))
     if (auto CS = CallSite(&I))
-      if (Function *Callee = CS.getCalledFunction())
+      for (Function *Callee : CS.getKnownCallees()) {
         if (Visited.insert(Callee).second && !Callee->isDeclaration()) {
           Node &CalleeN = *G.lookup(*Callee);
           Edge *E = N->lookup(CalleeN);
@@ -467,6 +467,7 @@ LazyCallGraph::SCC &llvm::updateCGAndAnalysisManagerForFunctionPass(
           if (!E->isCall())
             PromotedRefTargets.insert(&CalleeN);
         }
+      }
 
   // Now walk all references.
   for (Instruction &I : instructions(F))
diff --git a/llvm/lib/Analysis/CallGraph.cpp b/llvm/lib/Analysis/CallGraph.cpp
index 70aeb1a688e..a7bb16299a5 100644
--- a/llvm/lib/Analysis/CallGraph.cpp
+++ b/llvm/lib/Analysis/CallGraph.cpp
@@ -37,6 +37,7 @@ CallGraph::CallGraph(Module &M)
 
 CallGraph::CallGraph(CallGraph &&Arg)
     : M(Arg.M), FunctionMap(std::move(Arg.FunctionMap)),
+      DummyNodeMap(std::move(Arg.DummyNodeMap)),
       ExternalCallingNode(Arg.ExternalCallingNode),
       CallsExternalNode(std::move(Arg.CallsExternalNode)) {
   Arg.FunctionMap.clear();
@@ -53,6 +54,8 @@ CallGraph::~CallGraph() {
 #ifndef NDEBUG
   for (auto &I : FunctionMap)
     I.second->allReferencesDropped();
+  for (auto &N : DummyNodeMap)
+    N.second->allReferencesDropped();
 #endif
 }
 
@@ -74,7 +77,14 @@ void CallGraph::addToCallGraph(Function *F) {
     for (Instruction &I : BB) {
       if (auto *Call = dyn_cast<CallBase>(&I)) {
         const Function *Callee = Call->getCalledFunction();
-        if (!Callee || !Intrinsic::isLeaf(Callee->getIntrinsicID()))
+        MDNode *CalleesMD = I.getMetadata(LLVMContext::MD_callees);
+        if (!Callee && CalleesMD)
+          // If an indirect call site has !callees metadata indicating its
+          // possible callees, we add an edge from the call site to a dummy
+          // node. When we construct the dummy node, we add edges from it to
+          // the functions indicated in the !callees metadata.
+          Node->addCalledFunction(Call, getOrInsertNodeForCalleesMD(CalleesMD));
+        else if (!Callee || !Intrinsic::isLeaf(Callee->getIntrinsicID()))
           // Indirect calls of intrinsics are not allowed so no need to check.
           // We can be more precise here by using TargetArg returned by
           // Intrinsic::isLeaf.
@@ -105,6 +115,11 @@ void CallGraph::print(raw_ostream &OS) const {
 
   for (CallGraphNode *CN : Nodes)
     CN->print(OS);
+
+  // The iteration order of the DummyNodeMap is deterministic, so we don't need
+  // to sort the nodes. Just print them.
+  for (auto &Entry : DummyNodeMap)
+    Entry.second->print(OS);
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
@@ -120,6 +135,12 @@ LLVM_DUMP_METHOD void CallGraph::dump() const { print(dbgs()); }
 Function *CallGraph::removeFunctionFromModule(CallGraphNode *CGN) {
   assert(CGN->empty() && "Cannot remove function from call "
          "graph if it references other functions!");
+
+  // If any dummy node references the node for the given function, we first
+  // need to remove those edges.
+  for (auto &Entry : DummyNodeMap)
+    Entry.second->removeAnyCallEdgeTo(CGN);
+
   Function *F = CGN->getFunction(); // Get the function for the call graph node
   FunctionMap.erase(F);             // Remove the call graph node from the map
 
@@ -154,6 +175,21 @@ CallGraphNode *CallGraph::getOrInsertFunction(const Function *F) {
   return CGN.get();
 }
 
+CallGraphNode *CallGraph::getOrInsertNodeForCalleesMD(MDNode *Callees) {
+  auto &CGN = DummyNodeMap[Callees];
+  if (CGN)
+    return CGN.get();
+  CGN = llvm::make_unique<CallGraphNode>(nullptr);
+  for (const MDOperand &Op : Callees->operands())
+    if (auto *MDConstant = mdconst::extract_or_null<Constant>(Op)) {
+      auto *F = cast<Function>(MDConstant);
+
+      assert(!F->isIntrinsic());
+      CGN->addCalledFunction(nullptr, getOrInsertFunction(F));
+    }
+  return CGN.get();
+}
+
 //===----------------------------------------------------------------------===//
 // Implementations of the CallGraphNode class methods.
 //
@@ -171,7 +207,7 @@ void CallGraphNode::print(raw_ostream &OS) const {
     if (Function *FI = I.second->getFunction())
       OS << "function '" << FI->getName() <<"'\n";
     else
-      OS << "external node\n";
+      OS << "<<null function>><<" << I.second << ">>\n";
   }
   OS << '\n';
 }
diff --git a/llvm/lib/Analysis/LazyCallGraph.cpp b/llvm/lib/Analysis/LazyCallGraph.cpp
index 08d6e76ea03..c832dc94877 100644
--- a/llvm/lib/Analysis/LazyCallGraph.cpp
+++ b/llvm/lib/Analysis/LazyCallGraph.cpp
@@ -100,12 +100,14 @@ LazyCallGraph::EdgeSequence &LazyCallGraph::Node::populateSlow() {
   for (BasicBlock &BB : *F)
     for (Instruction &I : BB) {
       if (auto CS = CallSite(&I))
-        if (Function *Callee = CS.getCalledFunction())
+        for (Function *Callee : CS.getKnownCallees())
           if (!Callee->isDeclaration())
             if (Callees.insert(Callee).second) {
               Visited.insert(Callee);
+              auto EdgeK = CS.getCalledFunction() ? LazyCallGraph::Edge::Call
+                                                  : LazyCallGraph::Edge::Ref;
               addEdge(Edges->Edges, Edges->EdgeIndexMap, G->get(*Callee),
-                      LazyCallGraph::Edge::Call);
+                      EdgeK);
             }
 
       for (Value *Op : I.operand_values())