[CallSite removal] Move the legacy PM, call graph, and some inliner code to `CallBase`.

[chandlerc]

Summary

This patch focuses on the legacy PM, call graph, and some of inliner and legacy
passes interacting with those APIs from CallSite to the new CallBase class.
No interesting changes.

[fedor.sergeev]

LGTM, thanks!

llvm/include/llvm/Analysis/CallGraph.h

 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_CALLGRAPH_H
 #define LLVM_ANALYSIS_CALLGRAPH_H
 
 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/IR/CallSite.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Pass.h"
 #include <cassert>
 #include <map>
 #include <memory>
 #include <utility>
public:
   /// Moves all the callee information from N to this node.
   void stealCalledFunctionsFrom(CallGraphNode *N) {
     assert(CalledFunctions.empty() &&
            "Cannot steal callsite information if I already have some");
     std::swap(CalledFunctions, N->CalledFunctions);
   }
 
   /// Adds a function to the list of functions called by this one.
-  void addCalledFunction(CallSite CS, CallGraphNode *M) {
-    assert(!CS.getInstruction() || !CS.getCalledFunction() ||
-           !CS.getCalledFunction()->isIntrinsic() ||
-           !Intrinsic::isLeaf(CS.getCalledFunction()->getIntrinsicID()));
-    CalledFunctions.emplace_back(CS.getInstruction(), M);
+  void addCalledFunction(CallBase *Call, CallGraphNode *M) {
+    assert(!Call || !Call->getCalledFunction() ||
+           !Call->getCalledFunction()->isIntrinsic() ||
+           !Intrinsic::isLeaf(Call->getCalledFunction()->getIntrinsicID()));
+    CalledFunctions.emplace_back(Call, M);
     M->AddRef();
   }
 
   void removeCallEdge(iterator I) {
     I->second->DropRef();
     *I = CalledFunctions.back();
     CalledFunctions.pop_back();
   }
 
   /// Removes the edge in the node for the specified call site.
   ///
   /// Note that this method takes linear time, so it should be used sparingly.
-  void removeCallEdgeFor(CallSite CS);
+  void removeCallEdgeFor(CallBase &Call);
 
   /// Removes all call edges from this node to the specified callee
   /// function.
   ///
   /// This takes more time to execute than removeCallEdgeTo, so it should not
   /// be used unless necessary.
   void removeAnyCallEdgeTo(CallGraphNode *Callee);
 
   /// Removes one edge associated with a null callsite from this node to
   /// the specified callee function.
   void removeOneAbstractEdgeTo(CallGraphNode *Callee);
 
   /// Replaces the edge in the node for the specified call site with a
   /// new one.
   ///
   /// Note that this method takes linear time, so it should be used sparingly.
-  void replaceCallEdge(CallSite CS, CallSite NewCS, CallGraphNode *NewNode);
+  void replaceCallEdge(CallBase &Call, CallBase &NewCall,
+                       CallGraphNode *NewNode);
 
 private:
   friend class CallGraph;
 
   Function *F;
 
   std::vector<CallRecord> CalledFunctions;
 

llvm/lib/Analysis/CallGraph.cpp

 //===- CallGraph.cpp - Build a Module's call graph ------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Config/llvm-config.h"
-#include "llvm/IR/CallSite.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"

 }
 
 void CallGraph::addToCallGraph(Function *F) {
   CallGraphNode *Node = getOrInsertFunction(F);
 
   // If this function has external linkage or has its address taken, anything
   // could call it.
   if (!F->hasLocalLinkage() || F->hasAddressTaken())
-    ExternalCallingNode->addCalledFunction(CallSite(), Node);
+    ExternalCallingNode->addCalledFunction(nullptr, Node);
 
   // If this function is not defined in this translation unit, it could call
   // anything.
   if (F->isDeclaration() && !F->isIntrinsic())
-    Node->addCalledFunction(CallSite(), CallsExternalNode.get());
+    Node->addCalledFunction(nullptr, CallsExternalNode.get());
 
   // Look for calls by this function.
   for (BasicBlock &BB : *F)
     for (Instruction &I : BB) {
-      if (auto CS = CallSite(&I)) {
-        const Function *Callee = CS.getCalledFunction();
+      if (auto *Call = dyn_cast<CallBase>(&I)) {
+        const Function *Callee = Call->getCalledFunction();
         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.
-          Node->addCalledFunction(CS, CallsExternalNode.get());
+          Node->addCalledFunction(Call, CallsExternalNode.get());
         else if (!Callee->isIntrinsic())
-          Node->addCalledFunction(CS, getOrInsertFunction(Callee));
+          Node->addCalledFunction(Call, getOrInsertFunction(Callee));
       }
     }
 }
 
 void CallGraph::print(raw_ostream &OS) const {
   // Print in a deterministic order by sorting CallGraphNodes by name.  We do
   // this here to avoid slowing down the non-printing fast path.
 

 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD void CallGraphNode::dump() const { print(dbgs()); }
 #endif
 
 /// removeCallEdgeFor - This method removes the edge in the node for the
 /// specified call site.  Note that this method takes linear time, so it
 /// should be used sparingly.
-void CallGraphNode::removeCallEdgeFor(CallSite CS) {
+void CallGraphNode::removeCallEdgeFor(CallBase &Call) {
   for (CalledFunctionsVector::iterator I = CalledFunctions.begin(); ; ++I) {
     assert(I != CalledFunctions.end() && "Cannot find callsite to remove!");
-    if (I->first == CS.getInstruction()) {
+    if (I->first == &Call) {
       I->second->DropRef();
       *I = CalledFunctions.back();
       CalledFunctions.pop_back();
       return;
     }
   }
 }
 
if (CR.second == Callee && CR.first == nullptr) {
       return;
     }
   }
 }
 
 /// replaceCallEdge - This method replaces the edge in the node for the
 /// specified call site with a new one.  Note that this method takes linear
 /// time, so it should be used sparingly.
-void CallGraphNode::replaceCallEdge(CallSite CS,
-                                    CallSite NewCS, CallGraphNode *NewNode){
+void CallGraphNode::replaceCallEdge(CallBase &Call, CallBase &NewCall,
+                                    CallGraphNode *NewNode) {
   for (CalledFunctionsVector::iterator I = CalledFunctions.begin(); ; ++I) {
     assert(I != CalledFunctions.end() && "Cannot find callsite to remove!");
-    if (I->first == CS.getInstruction()) {
+    if (I->first == &Call) {
       I->second->DropRef();
-      I->first = NewCS.getInstruction();
+      I->first = &NewCall;
       I->second = NewNode;
       NewNode->AddRef();
       return;
     }
   }
 }
 
 // Provide an explicit template instantiation for the static ID.

llvm/lib/Analysis/CallGraphSCCPass.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/CallGraph.h"
-#include "llvm/IR/CallSite.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRPrintingPasses.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LegacyPassManagers.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/OptBisect.h"
 #include "llvm/IR/PassTimingInfo.h"

 /// CallGraphSCC passes to verify that they correctly updated the callgraph.
 ///
 /// This function returns true if it devirtualized an existing function call,
 /// meaning it turned an indirect call into a direct call.  This happens when
 /// a function pass like GVN optimizes away stuff feeding the indirect call.
 /// This never happens in checking mode.
 bool CGPassManager::RefreshCallGraph(const CallGraphSCC &CurSCC, CallGraph &CG,
                                      bool CheckingMode) {
-  DenseMap<Value*, CallGraphNode*> CallSites;
+  DenseMap<Value *, CallGraphNode *> Calls;
 
   LLVM_DEBUG(dbgs() << "CGSCCPASSMGR: Refreshing SCC with " << CurSCC.size()
                     << " nodes:\n";
              for (CallGraphNode *CGN
                   : CurSCC) CGN->dump(););
 
   bool MadeChange = false;
   bool DevirtualizedCall = false;
for (CallGraphSCC::iterator SCCIdx = CurSCC.begin(), E = CurSCC.end();
     // Keep track of the number of direct and indirect calls that were
     // invalidated and removed.
     unsigned NumDirectRemoved = 0, NumIndirectRemoved = 0;
 
     // Get the set of call sites currently in the function.
     for (CallGraphNode::iterator I = CGN->begin(), E = CGN->end(); I != E; ) {
       // If this call site is null, then the function pass deleted the call
       // entirely and the WeakTrackingVH nulled it out.
+      auto *Call = dyn_cast_or_null<CallBase>(I->first);
       if (!I->first ||
           // If we've already seen this call site, then the FunctionPass RAUW'd
           // one call with another, which resulted in two "uses" in the edge
           // list of the same call.
-          CallSites.count(I->first) ||
+          Calls.count(I->first) ||
 
           // If the call edge is not from a call or invoke, or it is a
           // instrinsic call, then the function pass RAUW'd a call with
           // another value. This can happen when constant folding happens
           // of well known functions etc.
-          !CallSite(I->first) ||
-          (CallSite(I->first).getCalledFunction() &&
-           CallSite(I->first).getCalledFunction()->isIntrinsic() &&
-           Intrinsic::isLeaf(
-               CallSite(I->first).getCalledFunction()->getIntrinsicID()))) {
+          !Call ||
+          (Call->getCalledFunction() &&
+           Call->getCalledFunction()->isIntrinsic() &&
+           Intrinsic::isLeaf(Call->getCalledFunction()->getIntrinsicID()))) {
         assert(!CheckingMode &&
                "CallGraphSCCPass did not update the CallGraph correctly!");
 
         // If this was an indirect call site, count it.
         if (!I->second->getFunction())
           ++NumIndirectRemoved;
         else
           ++NumDirectRemoved;
 
         // Just remove the edge from the set of callees, keep track of whether
         // I points to the last element of the vector.
         bool WasLast = I + 1 == E;
         CGN->removeCallEdge(I);
 
         // If I pointed to the last element of the vector, we have to bail out:
         // iterator checking rejects comparisons of the resultant pointer with
         // end.
         if (WasLast)
           break;
         E = CGN->end();
         continue;
       }
 
-      assert(!CallSites.count(I->first) &&
+      assert(!Calls.count(I->first) &&
              "Call site occurs in node multiple times");
 
-      CallSite CS(I->first);
-      if (CS) {
-        Function *Callee = CS.getCalledFunction();
+      if (Call) {
+        Function *Callee = Call->getCalledFunction();
         // Ignore intrinsics because they're not really function calls.
         if (!Callee || !(Callee->isIntrinsic()))
-          CallSites.insert(std::make_pair(I->first, I->second));
+          Calls.insert(std::make_pair(I->first, I->second));
       }
       ++I;
     }
 
     // Loop over all of the instructions in the function, getting the callsites.
     // Keep track of the number of direct/indirect calls added.
     unsigned NumDirectAdded = 0, NumIndirectAdded = 0;
 
     for (BasicBlock &BB : *F)
       for (Instruction &I : BB) {
-        CallSite CS(&I);
-        if (!CS) continue;
-        Function *Callee = CS.getCalledFunction();
-        if (Callee && Callee->isIntrinsic()) continue;
+        auto *Call = dyn_cast<CallBase>(&I);
+        if (!Call)
+          continue;
+        Function *Callee = Call->getCalledFunction();
+        if (Callee && Callee->isIntrinsic())
+          continue;
 
         // If this call site already existed in the callgraph, just verify it
-        // matches up to expectations and remove it from CallSites.
-        DenseMap<Value*, CallGraphNode*>::iterator ExistingIt =
-          CallSites.find(CS.getInstruction());
-        if (ExistingIt != CallSites.end()) {
+        // matches up to expectations and remove it from Calls.
+        DenseMap<Value *, CallGraphNode *>::iterator ExistingIt =
+            Calls.find(Call);
+        if (ExistingIt != Calls.end()) {
           CallGraphNode *ExistingNode = ExistingIt->second;
 
-          // Remove from CallSites since we have now seen it.
-          CallSites.erase(ExistingIt);
+          // Remove from Calls since we have now seen it.
+          Calls.erase(ExistingIt);
 
           // Verify that the callee is right.
-          if (ExistingNode->getFunction() == CS.getCalledFunction())
+          if (ExistingNode->getFunction() == Call->getCalledFunction())
             continue;
 
           // If we are in checking mode, we are not allowed to actually mutate
           // the callgraph.  If this is a case where we can infer that the
           // callgraph is less precise than it could be (e.g. an indirect call
           // site could be turned direct), don't reject it in checking mode, and
           // don't tweak it to be more precise.
-          if (CheckingMode && CS.getCalledFunction() &&
+          if (CheckingMode && Call->getCalledFunction() &&
               ExistingNode->getFunction() == nullptr)
             continue;
 
           assert(!CheckingMode &&
                  "CallGraphSCCPass did not update the CallGraph correctly!");
 
           // If not, we either went from a direct call to indirect, indirect to
           // direct, or direct to different direct.
           CallGraphNode *CalleeNode;
-          if (Function *Callee = CS.getCalledFunction()) {
+          if (Function *Callee = Call->getCalledFunction()) {
             CalleeNode = CG.getOrInsertFunction(Callee);
             // Keep track of whether we turned an indirect call into a direct
             // one.
             if (!ExistingNode->getFunction()) {
               DevirtualizedCall = true;
               LLVM_DEBUG(dbgs() << "  CGSCCPASSMGR: Devirtualized call to '"
                                 << Callee->getName() << "'\n");
             }
           } else {
             CalleeNode = CG.getCallsExternalNode();
           }
 
           // Update the edge target in CGN.
-          CGN->replaceCallEdge(CS, CS, CalleeNode);
+          CGN->replaceCallEdge(*Call, *Call, CalleeNode);
           MadeChange = true;
           continue;
         }
 
         assert(!CheckingMode &&
                "CallGraphSCCPass did not update the CallGraph correctly!");
 
         // If the call site didn't exist in the CGN yet, add it.
         CallGraphNode *CalleeNode;
-        if (Function *Callee = CS.getCalledFunction()) {
+        if (Function *Callee = Call->getCalledFunction()) {
           CalleeNode = CG.getOrInsertFunction(Callee);
           ++NumDirectAdded;
         } else {
           CalleeNode = CG.getCallsExternalNode();
           ++NumIndirectAdded;
         }
 
-        CGN->addCalledFunction(CS, CalleeNode);
+        CGN->addCalledFunction(Call, CalleeNode);
         MadeChange = true;
       }
 
     // We scanned the old callgraph node, removing invalidated call sites and
     // then added back newly found call sites.  One thing that can happen is
     // that an old indirect call site was deleted and replaced with a new direct
     // call.  In this case, we have devirtualized a call, and CGSCCPM would like
     // to iteratively optimize the new code.  Unfortunately, we don't really
     // have a great way to detect when this happens.  As an approximation, we
     // just look at whether the number of indirect calls is reduced and the
     // number of direct calls is increased.  There are tons of ways to fool this
     // (e.g. DCE'ing an indirect call and duplicating an unrelated block with a
     // direct call) but this is close enough.
     if (NumIndirectRemoved > NumIndirectAdded &&
         NumDirectRemoved < NumDirectAdded)
       DevirtualizedCall = true;
 
     // After scanning this function, if we still have entries in callsites, then
     // they are dangling pointers.  WeakTrackingVH should save us for this, so
     // abort if
     // this happens.
-    assert(CallSites.empty() && "Dangling pointers found in call sites map");
+    assert(Calls.empty() && "Dangling pointers found in call sites map");
 
     // Periodically do an explicit clear to remove tombstones when processing
     // large scc's.
     if ((FunctionNo & 15) == 15)
-      CallSites.clear();
+      Calls.clear();
   }
 
   LLVM_DEBUG(if (MadeChange) {
     dbgs() << "CGSCCPASSMGR: Refreshed SCC is now:\n";
     for (CallGraphNode *CGN : CurSCC)
       CGN->dump();
     if (DevirtualizedCall)
       dbgs() << "CGSCCPASSMGR: Refresh devirtualized a call!\n";

llvm/lib/Transforms/Coroutines/Coroutines.cpp

 // FIXME: This code is stolen from CallGraph::addToCallGraph(Function *F), which
 // happens to be private. It is better for this functionality exposed by the
 // CallGraph.
 static void buildCGN(CallGraph &CG, CallGraphNode *Node) {
   Function *F = Node->getFunction();
 
   // Look for calls by this function.
   for (Instruction &I : instructions(F))
-    if (CallSite CS = CallSite(cast<Value>(&I))) {
-      const Function *Callee = CS.getCalledFunction();
+    if (auto *Call = dyn_cast<CallBase>(&I)) {
+      const Function *Callee = Call->getCalledFunction();
       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.
-        Node->addCalledFunction(CS, CG.getCallsExternalNode());
+        Node->addCalledFunction(Call, CG.getCallsExternalNode());
       else if (!Callee->isIntrinsic())
-        Node->addCalledFunction(CS, CG.getOrInsertFunction(Callee));
+        Node->addCalledFunction(Call, CG.getOrInsertFunction(Callee));
     }
 }
 
 // Rebuild CGN after we extracted parts of the code from ParentFunc into
 // NewFuncs. Builds CGNs for the NewFuncs and adds them to the current SCC.
 void coro::updateCallGraph(Function &ParentFunc, ArrayRef<Function *> NewFuncs,
                            CallGraph &CG, CallGraphSCC &SCC) {
   // Rebuild CGN from scratch for the ParentFunc

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

for (CallGraphNode *OldNode : SCC) {
       if (!OldF)
         continue;
 
       auto ReplaceCallSite = [&](CallSite OldCS, CallSite NewCS) {
         Function *Caller = OldCS.getInstruction()->getParent()->getParent();
         CallGraphNode *NewCalleeNode =
             CG.getOrInsertFunction(NewCS.getCalledFunction());
         CallGraphNode *CallerNode = CG[Caller];
-        CallerNode->replaceCallEdge(OldCS, NewCS, NewCalleeNode);
+        CallerNode->replaceCallEdge(*cast<CallBase>(OldCS.getInstruction()),
+                                    *cast<CallBase>(NewCS.getInstruction()),
+                                    NewCalleeNode);
       };
 
       const TargetTransformInfo &TTI =
           getAnalysis<TargetTransformInfoWrapperPass>().getTTI(*OldF);
       if (Function *NewF = promoteArguments(OldF, AARGetter, MaxElements,
                                             {ReplaceCallSite}, TTI)) {
         LocalChange = true;
 

llvm/lib/Transforms/IPO/Inliner.cpp

for (unsigned CSi = 0; CSi != CallSites.size(); ++CSi) {
       // If this call site is dead and it is to a readonly function, we should
       // just delete the call instead of trying to inline it, regardless of
       // size.  This happens because IPSCCP propagates the result out of the
       // call and then we're left with the dead call.
       if (IsTriviallyDead) {
         LLVM_DEBUG(dbgs() << "    -> Deleting dead call: " << *Instr << "\n");
         // Update the call graph by deleting the edge from Callee to Caller.
         setInlineRemark(CS, "trivially dead");
-        CG[Caller]->removeCallEdgeFor(CS);
+        CG[Caller]->removeCallEdgeFor(*cast<CallBase>(CS.getInstruction()));
         Instr->eraseFromParent();
         ++NumCallsDeleted;
       } else {
         // Get DebugLoc to report. CS will be invalid after Inliner.
         DebugLoc DLoc = CS->getDebugLoc();
         BasicBlock *Block = CS.getParent();
 
         // Attempt to inline the function.

llvm/lib/Transforms/IPO/PruneEH.cpp

static void DeleteBasicBlock(BasicBlock *BB, CallGraph &CG) {
   for (BasicBlock::iterator I = BB->end(), E = BB->begin(); I != E; ) {
     --I;
 
     if (I->getType()->isTokenTy()) {
       TokenInst = &*I;
       break;
     }
 
-    if (auto CS = CallSite (&*I)) {
-      const Function *Callee = CS.getCalledFunction();
+    if (auto *Call = dyn_cast<CallBase>(&*I)) {
+      const Function *Callee = Call->getCalledFunction();
       if (!Callee || !Intrinsic::isLeaf(Callee->getIntrinsicID()))
-        CGN->removeCallEdgeFor(CS);
+        CGN->removeCallEdgeFor(*Call);
       else if (!Callee->isIntrinsic())
-        CGN->removeCallEdgeFor(CS);
+        CGN->removeCallEdgeFor(*Call);
     }
 
     if (!I->use_empty())
       I->replaceAllUsesWith(UndefValue::get(I->getType()));
   }
 
   if (TokenInst) {
     if (!TokenInst->isTerminator())

llvm/lib/Transforms/Utils/InlineFunction.cpp

for (; I != E; ++I) {
 
     ValueToValueMapTy::iterator VMI = VMap.find(OrigCall);
     // Only copy the edge if the call was inlined!
     if (VMI == VMap.end() || VMI->second == nullptr)
       continue;
 
     // If the call was inlined, but then constant folded, there is no edge to
     // add.  Check for this case.
-    Instruction *NewCall = dyn_cast<Instruction>(VMI->second);
+    auto *NewCall = dyn_cast<CallBase>(VMI->second);
     if (!NewCall)
       continue;
 
     // We do not treat intrinsic calls like real function calls because we
     // expect them to become inline code; do not add an edge for an intrinsic.
-    CallSite CS = CallSite(NewCall);
-    if (CS && CS.getCalledFunction() && CS.getCalledFunction()->isIntrinsic())
+    if (NewCall->getCalledFunction() &&
+        NewCall->getCalledFunction()->isIntrinsic())
       continue;
 
     // Remember that this call site got inlined for the client of
     // InlineFunction.
     IFI.InlinedCalls.push_back(NewCall);
 
     // It's possible that inlining the callsite will cause it to go from an
     // indirect to a direct call by resolving a function pointer.  If this
     // happens, set the callee of the new call site to a more precise
     // destination.  This can also happen if the call graph node of the caller
     // was just unnecessarily imprecise.
     if (!I->second->getFunction())
-      if (Function *F = CallSite(NewCall).getCalledFunction()) {
+      if (Function *F = NewCall->getCalledFunction()) {
         // Indirect call site resolved to direct call.
-        CallerNode->addCalledFunction(CallSite(NewCall), CG[F]);
+        CallerNode->addCalledFunction(NewCall, CG[F]);
 
         continue;
       }
 
-    CallerNode->addCalledFunction(CallSite(NewCall), I->second);
+    CallerNode->addCalledFunction(NewCall, I->second);
   }
 
   // Update the call graph by deleting the edge from Callee to Caller.  We must
   // do this after the loop above in case Caller and Callee are the same.
-  CallerNode->removeCallEdgeFor(CS);
+  CallerNode->removeCallEdgeFor(*cast<CallBase>(CS.getInstruction()));
 }
 
 static void HandleByValArgumentInit(Value *Dst, Value *Src, Module *M,
                                     BasicBlock *InsertBlock,
                                     InlineFunctionInfo &IFI) {
   Type *AggTy = cast<PointerType>(Src->getType())->getElementType();
   IRBuilder<> Builder(InsertBlock, InsertBlock->begin());