[llvm][NFC] Replace CallSite with CallBase in Inliner

[mtrofin]

Summary

*Almost* all uses are replaced. Left FIXMEs for the two sites that
require refactoring outside of Inliner, to scope this patch.

Depends on D77762.

[mtrofin]

abandoning - arc diff snafu. real revision is D77817

llvm/include/llvm/Transforms/IPO/Inliner.h

struct LegacyInlinerBase : public CallGraphSCCPass {
 
   /// Remove now-dead linkonce functions at the end of processing to avoid
   /// breaking the SCC traversal.
   bool doFinalization(CallGraph &CG) override;
 
   /// This method must be implemented by the subclass to determine the cost of
   /// inlining the specified call site.  If the cost returned is greater than
   /// the current inline threshold, the call site is not inlined.
+  // FIXME(mtrofin): remove this in favor of the CallBase-based one
   virtual InlineCost getInlineCost(CallSite CS) = 0;
 
+  /// This method must be implemented by the subclass to determine the cost of
+  /// inlining the specified call site.  If the cost returned is greater than
+  /// the current inline threshold, the call site is not inlined.
+  virtual InlineCost getInlineCost(CallBase *CB) {
+    return getInlineCost(CallSite(CB));
+  }
+
   /// Remove dead functions.
   ///
   /// This also includes a hack in the form of the 'AlwaysInlineOnly' flag
   /// which restricts it to deleting functions with an 'AlwaysInline'
   /// attribute. This is useful for the InlineAlways pass that only wants to
   /// deal with that subset of the functions.
   bool removeDeadFunctions(CallGraph &CG, bool AlwaysInlineOnly = false);
 

llvm/lib/Transforms/IPO/Inliner.cpp

 /// do so and update the CallGraph for this operation.
 ///
 /// This function also does some basic book-keeping to update the IR.  The
 /// InlinedArrayAllocas map keeps track of any allocas that are already
 /// available from other functions inlined into the caller.  If we are able to
 /// inline this call site we attempt to reuse already available allocas or add
 /// any new allocas to the set if not possible.
 static InlineResult InlineCallIfPossible(
-    CallSite CS, InlineFunctionInfo &IFI,
+    CallBase *CS, InlineFunctionInfo &IFI,
     InlinedArrayAllocasTy &InlinedArrayAllocas, int InlineHistory,
     bool InsertLifetime, function_ref<AAResults &(Function &)> &AARGetter,
     ImportedFunctionsInliningStatistics &ImportedFunctionsStats) {
-  Function *Callee = CS.getCalledFunction();
-  Function *Caller = CS.getCaller();
+  Function *Callee = CS->getCalledFunction();
+  Function *Caller = CS->getCaller();
 
   AAResults &AAR = AARGetter(*Callee);
 
   // Try to inline the function.  Get the list of static allocas that were
   // inlined.
   InlineResult IR = InlineFunction(CS, IFI, &AAR, InsertLifetime);
   if (!IR.isSuccess())
     return IR;

 }
 
 /// Return true if inlining of CS can block the caller from being
 /// inlined which is proved to be more beneficial. \p IC is the
 /// estimated inline cost associated with callsite \p CS.
 /// \p TotalSecondaryCost will be set to the estimated cost of inlining the
 /// caller if \p CS is suppressed for inlining.
 static bool
-shouldBeDeferred(Function *Caller, CallSite CS, InlineCost IC,
+shouldBeDeferred(Function *Caller, CallBase *CS, InlineCost IC,
                  int &TotalSecondaryCost,
-                 function_ref<InlineCost(CallSite CS)> GetInlineCost) {
+                 function_ref<InlineCost(CallBase *CS)> GetInlineCost) {
   // For now we only handle local or inline functions.
   if (!Caller->hasLocalLinkage() && !Caller->hasLinkOnceODRLinkage())
     return false;
   // If the cost of inlining CS is non-positive, it is not going to prevent the
   // caller from being inlined into its callers and hence we don't need to
   // defer.
   if (IC.getCost() <= 0)
     return false;
shouldBeDeferred(Function *Caller, CallBase *CS, InlineCost IC,
   // This bool tracks what happens if we DO inline C into B.
   bool inliningPreventsSomeOuterInline = false;
   for (User *U : Caller->users()) {
     // If the caller will not be removed (either because it does not have a
     // local linkage or because the LastCallToStaticBonus has been already
     // applied), then we can exit the loop early.
     if (!ApplyLastCallBonus && TotalSecondaryCost >= IC.getCost())
       return false;
-    CallSite CS2(U);
+    CallBase *CS2 = dyn_cast<CallBase>(U);
 
     // If this isn't a call to Caller (it could be some other sort
     // of reference) skip it.  Such references will prevent the caller
     // from being removed.
-    if (!CS2 || CS2.getCalledFunction() != Caller) {
+    if (!CS2 || CS2->getCalledFunction() != Caller) {
       ApplyLastCallBonus = false;
       continue;
     }
 
     InlineCost IC2 = GetInlineCost(CS2);
     ++NumCallerCallersAnalyzed;
     if (!IC2) {
       ApplyLastCallBonus = false;
static std::string inlineCostStr(const InlineCost &IC) {
   Remark << IC;
   return Remark.str();
 }
 
 /// Return the cost only if the inliner should attempt to inline at the given
 /// CallSite. If we return the cost, we will emit an optimisation remark later
 /// using that cost, so we won't do so from this function.
 static Optional<InlineCost>
-shouldInline(CallSite CS, function_ref<InlineCost(CallSite CS)> GetInlineCost,
+shouldInline(CallBase *CS, function_ref<InlineCost(CallBase *CS)> GetInlineCost,
              OptimizationRemarkEmitter &ORE) {
   using namespace ore;
 
   InlineCost IC = GetInlineCost(CS);
-  Instruction *Call = CS.getInstruction();
-  Function *Callee = CS.getCalledFunction();
-  Function *Caller = CS.getCaller();
+  Instruction *Call = CS;
+  Function *Callee = CS->getCalledFunction();
+  Function *Caller = CS->getCaller();
 
   if (IC.isAlways()) {
     LLVM_DEBUG(dbgs() << "    Inlining " << inlineCostStr(IC)
-                      << ", Call: " << *CS.getInstruction() << "\n");
+                      << ", Call: " << *CS << "\n");
     return IC;
   }
 
   if (IC.isNever()) {
     LLVM_DEBUG(dbgs() << "    NOT Inlining " << inlineCostStr(IC)
-                      << ", Call: " << *CS.getInstruction() << "\n");
+                      << ", Call: " << *CS << "\n");
     ORE.emit([&]() {
       return OptimizationRemarkMissed(DEBUG_TYPE, "NeverInline", Call)
              << NV("Callee", Callee) << " not inlined into "
              << NV("Caller", Caller) << " because it should never be inlined "
              << IC;
     });
     return IC;
   }
 
   if (!IC) {
     LLVM_DEBUG(dbgs() << "    NOT Inlining " << inlineCostStr(IC)
-                      << ", Call: " << *CS.getInstruction() << "\n");
+                      << ", Call: " << *CS << "\n");
     ORE.emit([&]() {
       return OptimizationRemarkMissed(DEBUG_TYPE, "TooCostly", Call)
              << NV("Callee", Callee) << " not inlined into "
              << NV("Caller", Caller) << " because too costly to inline " << IC;
     });
     return IC;
   }
 
   int TotalSecondaryCost = 0;
   if (shouldBeDeferred(Caller, CS, IC, TotalSecondaryCost, GetInlineCost)) {
-    LLVM_DEBUG(dbgs() << "    NOT Inlining: " << *CS.getInstruction()
+    LLVM_DEBUG(dbgs() << "    NOT Inlining: " << *CS
                       << " Cost = " << IC.getCost()
                       << ", outer Cost = " << TotalSecondaryCost << '\n');
     ORE.emit([&]() {
       return OptimizationRemarkMissed(DEBUG_TYPE, "IncreaseCostInOtherContexts",
                                       Call)
              << "Not inlining. Cost of inlining " << NV("Callee", Callee)
              << " increases the cost of inlining " << NV("Caller", Caller)
              << " in other contexts";
     });
 
     // IC does not bool() to false, so get an InlineCost that will.
     // This will not be inspected to make an error message.
     return None;
   }
 
-  LLVM_DEBUG(dbgs() << "    Inlining " << inlineCostStr(IC)
-                    << ", Call: " << *CS.getInstruction() << '\n');
+  LLVM_DEBUG(dbgs() << "    Inlining " << inlineCostStr(IC) << ", Call: " << *CS
+                    << '\n');
   return IC;
 }
 
 /// Return true if the specified inline history ID
 /// indicates an inline history that includes the specified function.
 static bool InlineHistoryIncludes(
     Function *F, int InlineHistoryID,
     const SmallVectorImpl<std::pair<Function *, int>> &InlineHistory) {
ORE.emit([&]() {
     bool AlwaysInline = IC.isAlways();
     StringRef RemarkName = AlwaysInline ? "AlwaysInline" : "Inlined";
     return OptimizationRemark(DEBUG_TYPE, RemarkName, DLoc, Block)
            << ore::NV("Callee", &Callee) << " inlined into "
            << ore::NV("Caller", &Caller) << " with " << IC;
   });
 }
 
-static void setInlineRemark(CallSite &CS, StringRef message) {
+static void setInlineRemark(CallBase *CS, StringRef message) {
   if (!InlineRemarkAttribute)
     return;
 
   Attribute attr = Attribute::get(CS->getContext(), "inline-remark", message);
-  CS.addAttribute(AttributeList::FunctionIndex, attr);
+  CS->addAttribute(AttributeList::FunctionIndex, attr);
 }
 
 static bool
 inlineCallsImpl(CallGraphSCC &SCC, CallGraph &CG,
                 std::function<AssumptionCache &(Function &)> GetAssumptionCache,
                 ProfileSummaryInfo *PSI,
                 std::function<const TargetLibraryInfo &(Function &)> GetTLI,
                 bool InsertLifetime,
-                function_ref<InlineCost(CallSite CS)> GetInlineCost,
+                function_ref<InlineCost(CallBase *CS)> GetInlineCost,
                 function_ref<AAResults &(Function &)> AARGetter,
                 ImportedFunctionsInliningStatistics &ImportedFunctionsStats) {
   SmallPtrSet<Function *, 8> SCCFunctions;
   LLVM_DEBUG(dbgs() << "Inliner visiting SCC:");
   for (CallGraphNode *Node : SCC) {
     Function *F = Node->getFunction();
     if (F)
       SCCFunctions.insert(F);
     LLVM_DEBUG(dbgs() << " " << (F ? F->getName() : "INDIRECTNODE"));
   }
 
   // Scan through and identify all call sites ahead of time so that we only
   // inline call sites in the original functions, not call sites that result
   // from inlining other functions.
-  SmallVector<std::pair<CallSite, int>, 16> CallSites;
+  SmallVector<std::pair<CallBase *, int>, 16> CallSites;
 
   // When inlining a callee produces new call sites, we want to keep track of
   // the fact that they were inlined from the callee.  This allows us to avoid
   // infinite inlining in some obscure cases.  To represent this, we use an
   // index into the InlineHistory vector.
   SmallVector<std::pair<Function *, int>, 8> InlineHistory;
 
   for (CallGraphNode *Node : SCC) {
     Function *F = Node->getFunction();
     if (!F || F->isDeclaration())
       continue;
 
     OptimizationRemarkEmitter ORE(F);
     for (BasicBlock &BB : *F)
       for (Instruction &I : BB) {
-        CallSite CS(cast<Value>(&I));
+        auto *CS = dyn_cast<CallBase>(&I);
         // If this isn't a call, or it is a call to an intrinsic, it can
         // never be inlined.
         if (!CS || isa<IntrinsicInst>(I))
           continue;
 
         // If this is a direct call to an external function, we can never inline
         // it.  If it is an indirect call, inlining may resolve it to be a
         // direct call, so we keep it.
-        if (Function *Callee = CS.getCalledFunction())
+        if (Function *Callee = CS->getCalledFunction())
           if (Callee->isDeclaration()) {
             using namespace ore;
 
             setInlineRemark(CS, "unavailable definition");
             ORE.emit([&]() {
               return OptimizationRemarkMissed(DEBUG_TYPE, "NoDefinition", &I)
                      << NV("Callee", Callee) << " will not be inlined into "
-                     << NV("Caller", CS.getCaller())
+                     << NV("Caller", CS->getCaller())
                      << " because its definition is unavailable"
                      << setIsVerbose();
             });
             continue;
           }
 
         CallSites.push_back(std::make_pair(CS, -1));
       }
   }
 
   LLVM_DEBUG(dbgs() << ": " << CallSites.size() << " call sites.\n");
 
   // If there are no calls in this function, exit early.
   if (CallSites.empty())
     return false;
 
   // Now that we have all of the call sites, move the ones to functions in the
   // current SCC to the end of the list.
   unsigned FirstCallInSCC = CallSites.size();
   for (unsigned i = 0; i < FirstCallInSCC; ++i)
-    if (Function *F = CallSites[i].first.getCalledFunction())
+    if (Function *F = CallSites[i].first->getCalledFunction())
       if (SCCFunctions.count(F))
         std::swap(CallSites[i--], CallSites[--FirstCallInSCC]);
 
   InlinedArrayAllocasTy InlinedArrayAllocas;
   InlineFunctionInfo InlineInfo(&CG, &GetAssumptionCache, PSI);
 
   // Now that we have all of the call sites, loop over them and inline them if
   // it looks profitable to do so.
   bool Changed = false;
   bool LocalChange;
   do {
     LocalChange = false;
     // Iterate over the outer loop because inlining functions can cause indirect
     // calls to become direct calls.
     // CallSites may be modified inside so ranged for loop can not be used.
     for (unsigned CSi = 0; CSi != CallSites.size(); ++CSi) {
-      CallSite CS = CallSites[CSi].first;
+      CallBase *CS = CallSites[CSi].first;
 
-      Function *Caller = CS.getCaller();
-      Function *Callee = CS.getCalledFunction();
+      Function *Caller = CS->getCaller();
+      Function *Callee = CS->getCalledFunction();
 
       // We can only inline direct calls to non-declarations.
       if (!Callee || Callee->isDeclaration())
         continue;
 
-      Instruction *Instr = CS.getInstruction();
+      Instruction *Instr = CS;
 
       bool IsTriviallyDead =
           isInstructionTriviallyDead(Instr, &GetTLI(*Caller));
 
       int InlineHistoryID;
       if (!IsTriviallyDead) {
         // If this call site was obtained by inlining another function, verify
         // that the include path for the function did not include the callee
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(*cast<CallBase>(CS.getInstruction()));
+        CG[Caller]->removeCallEdgeFor(*CS);
         Instr->eraseFromParent();
         ++NumCallsDeleted;
       } else {
         // Get DebugLoc to report. CS will be invalid after Inliner.
         DebugLoc DLoc = CS->getDebugLoc();
-        BasicBlock *Block = CS.getParent();
+        BasicBlock *Block = CS->getParent();
 
         // Attempt to inline the function.
         using namespace ore;
 
         InlineResult IR = InlineCallIfPossible(
             CS, InlineInfo, InlinedArrayAllocas, InlineHistoryID,
             InsertLifetime, AARGetter, ImportedFunctionsStats);
         if (!IR.isSuccess()) {
for (unsigned CSi = 0; CSi != CallSites.size(); ++CSi) {
           int NewHistoryID = InlineHistory.size();
           InlineHistory.push_back(std::make_pair(Callee, InlineHistoryID));
 
 #ifndef NDEBUG
           // Make sure no dupplicates in the inline candidates. This could
           // happen when a callsite is simpilfied to reusing the return value
           // of another callsite during function cloning, thus the other
           // callsite will be reconsidered here.
-          DenseSet<CallSite> DbgCallSites;
+          DenseSet<CallBase *> DbgCallSites;
           for (auto &II : CallSites)
             DbgCallSites.insert(II.first);
 #endif
 
           for (Value *Ptr : InlineInfo.InlinedCalls) {
 #ifndef NDEBUG
-            assert(DbgCallSites.count(CallSite(Ptr)) == 0);
+            assert(DbgCallSites.count(dyn_cast<CallBase>(Ptr)) == 0);
 #endif
-            CallSites.push_back(std::make_pair(CallSite(Ptr), NewHistoryID));
+            CallSites.push_back(
+                std::make_pair(dyn_cast<CallBase>(Ptr), NewHistoryID));
           }
         }
       }
 
       // If we inlined or deleted the last possible call site to the function,
       // delete the function body now.
       if (Callee && Callee->use_empty() && Callee->hasLocalLinkage() &&
           // TODO: Can remove if in SCC now.
bool LegacyInlinerBase::inlineCalls(CallGraphSCC &SCC) {
   GetTLI = [&](Function &F) -> const TargetLibraryInfo & {
     return getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
   };
   auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
     return ACT->getAssumptionCache(F);
   };
   return inlineCallsImpl(
       SCC, CG, GetAssumptionCache, PSI, GetTLI, InsertLifetime,
-      [this](CallSite CS) { return getInlineCost(CS); }, LegacyAARGetter(*this),
-      ImportedFunctionsStats);
+      [this](CallBase *CS) { return getInlineCost(CS); },
+      LegacyAARGetter(*this), ImportedFunctionsStats);
 }
 
 /// Remove now-dead linkonce functions at the end of
 /// processing to avoid breaking the SCC traversal.
 bool LegacyInlinerBase::doFinalization(CallGraph &CG) {
   if (InlinerFunctionImportStats != InlinerFunctionImportStatsOpts::No)
     ImportedFunctionsStats.dump(InlinerFunctionImportStats ==
                                 InlinerFunctionImportStatsOpts::Verbose);
PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
   // don't have a topological ordering to prevent explosions.
   //
   // To compensate for this, we don't process transitive edges made immediate
   // by inlining until we've done one pass of inlining across the entire SCC.
   // Large, highly connected SCCs still lead to some amount of code bloat in
   // this model, but it is uniformly spread across all the functions in the SCC
   // and eventually they all become too large to inline, rather than
   // incrementally maknig a single function grow in a super linear fashion.
-  SmallVector<std::pair<CallSite, int>, 16> Calls;
+  SmallVector<std::pair<CallBase *, int>, 16> Calls;
 
   FunctionAnalysisManager &FAM =
       AM.getResult<FunctionAnalysisManagerCGSCCProxy>(InitialC, CG)
           .getManager();
 
   // Populate the initial list of calls in this SCC.
   for (auto &N : InitialC) {
     auto &ORE =
         FAM.getResult<OptimizationRemarkEmitterAnalysis>(N.getFunction());
     // We want to generally process call sites top-down in order for
     // simplifications stemming from replacing the call with the returned value
     // after inlining to be visible to subsequent inlining decisions.
     // FIXME: Using instructions sequence is a really bad way to do this.
     // Instead we should do an actual RPO walk of the function body.
     for (Instruction &I : instructions(N.getFunction()))
-      if (auto CS = CallSite(&I))
-        if (Function *Callee = CS.getCalledFunction()) {
+      if (auto CS = dyn_cast<CallBase>(&I))
+        if (Function *Callee = CS->getCalledFunction()) {
           if (!Callee->isDeclaration())
             Calls.push_back({CS, -1});
           else if (!isa<IntrinsicInst>(I)) {
             using namespace ore;
             setInlineRemark(CS, "unavailable definition");
             ORE.emit([&]() {
               return OptimizationRemarkMissed(DEBUG_TYPE, "NoDefinition", &I)
                      << NV("Callee", Callee) << " will not be inlined into "
-                     << NV("Caller", CS.getCaller())
+                     << NV("Caller", CS->getCaller())
                      << " because its definition is unavailable"
                      << setIsVerbose();
             });
           }
         }
   }
   if (Calls.empty())
     return PreservedAnalyses::all();
PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
 
   // Loop forward over all of the calls. Note that we cannot cache the size as
   // inlining can introduce new calls that need to be processed.
   for (int i = 0; i < (int)Calls.size(); ++i) {
     // We expect the calls to typically be batched with sequences of calls that
     // have the same caller, so we first set up some shared infrastructure for
     // this caller. We also do any pruning we can at this layer on the caller
     // alone.
-    Function &F = *Calls[i].first.getCaller();
+    Function &F = *Calls[i].first->getCaller();
     LazyCallGraph::Node &N = *CG.lookup(F);
     if (CG.lookupSCC(N) != C)
       continue;
     if (F.hasOptNone()) {
       setInlineRemark(Calls[i].first, "optnone attribute");
       continue;
     }
 
for (int i = 0; i < (int)Calls.size(); ++i) {
     };
     auto GetBFI = [&](Function &F) -> BlockFrequencyInfo & {
       return FAM.getResult<BlockFrequencyAnalysis>(F);
     };
     auto GetTLI = [&](Function &F) -> const TargetLibraryInfo & {
       return FAM.getResult<TargetLibraryAnalysis>(F);
     };
 
-    auto GetInlineCost = [&](CallSite CS) {
-      Function &Callee = *CS.getCalledFunction();
+    auto GetInlineCost = [&](CallBase *CS) {
+      Function &Callee = *CS->getCalledFunction();
       auto &CalleeTTI = FAM.getResult<TargetIRAnalysis>(Callee);
       bool RemarksEnabled =
           Callee.getContext().getDiagHandlerPtr()->isMissedOptRemarkEnabled(
               DEBUG_TYPE);
-      return getInlineCost(cast<CallBase>(*CS.getInstruction()), Params,
-                           CalleeTTI, GetAssumptionCache, {GetBFI}, GetTLI, PSI,
-                           RemarksEnabled ? &ORE : nullptr);
+      return getInlineCost(*CS, Params, CalleeTTI, GetAssumptionCache, {GetBFI},
+                           GetTLI, PSI, RemarksEnabled ? &ORE : nullptr);
     };
 
     // Now process as many calls as we have within this caller in the sequnece.
     // We bail out as soon as the caller has to change so we can update the
     // call graph and prepare the context of that new caller.
     bool DidInline = false;
-    for (; i < (int)Calls.size() && Calls[i].first.getCaller() == &F; ++i) {
+    for (; i < (int)Calls.size() && Calls[i].first->getCaller() == &F; ++i) {
       int InlineHistoryID;
-      CallSite CS;
+      CallBase *CS;
       std::tie(CS, InlineHistoryID) = Calls[i];
-      Function &Callee = *CS.getCalledFunction();
+      Function &Callee = *CS->getCalledFunction();
 
       if (InlineHistoryID != -1 &&
           InlineHistoryIncludes(&Callee, InlineHistoryID, InlineHistory)) {
         setInlineRemark(CS, "recursive");
         continue;
       }
 
       // Check if this inlining may repeat breaking an SCC apart that has
for (; i < (int)Calls.size() && Calls[i].first->getCaller() == &F; ++i) {
         setInlineRemark(CS, inlineCostStr(*OIC));
         continue;
       }
 
       // Setup the data structure used to plumb customization into the
       // `InlineFunction` routine.
       InlineFunctionInfo IFI(
           /*cg=*/nullptr, &GetAssumptionCache, PSI,
-          &FAM.getResult<BlockFrequencyAnalysis>(*(CS.getCaller())),
+          &FAM.getResult<BlockFrequencyAnalysis>(*(CS->getCaller())),
           &FAM.getResult<BlockFrequencyAnalysis>(Callee));
 
       // Get DebugLoc to report. CS will be invalid after Inliner.
       DebugLoc DLoc = CS->getDebugLoc();
-      BasicBlock *Block = CS.getParent();
+      BasicBlock *Block = CS->getParent();
 
       using namespace ore;
 
       InlineResult IR = InlineFunction(CS, IFI);
       if (!IR.isSuccess()) {
         setInlineRemark(CS, std::string(IR.getFailureReason()) + "; " +
                                 inlineCostStr(*OIC));
         ORE.emit([&]() {
for (; i < (int)Calls.size() && Calls[i].first->getCaller() == &F; ++i) {
       ++NumInlined;
 
       emit_inlined_into(ORE, DLoc, Block, Callee, F, *OIC);
 
       // Add any new callsites to defined functions to the worklist.
       if (!IFI.InlinedCallSites.empty()) {
         int NewHistoryID = InlineHistory.size();
         InlineHistory.push_back({&Callee, InlineHistoryID});
+
+        // FIXME(mtrofin): refactor IFI.InlinedCallSites to be CallBase-based
         for (CallSite &CS : reverse(IFI.InlinedCallSites)) {
           Function *NewCallee = CS.getCalledFunction();
           if (!NewCallee) {
             // Try to promote an indirect (virtual) call without waiting for the
             // post-inline cleanup and the next DevirtSCCRepeatedPass iteration
             // because the next iteration may not happen and we may miss
             // inlining it.
             if (tryPromoteCall(CS))
               NewCallee = CS.getCalledFunction();
           }
           if (NewCallee)
             if (!NewCallee->isDeclaration())
-              Calls.push_back({CS, NewHistoryID});
+              Calls.push_back(
+                  {cast<CallBase>(CS.getInstruction()), NewHistoryID});
         }
       }
 
       if (InlinerFunctionImportStats != InlinerFunctionImportStatsOpts::No)
         ImportedFunctionsStats->recordInline(F, Callee);
 
       // Merge the attributes based on the inlining.
       AttributeFuncs::mergeAttributesForInlining(F, Callee);
 
       // For local functions, check whether this makes the callee trivially
       // dead. In that case, we can drop the body of the function eagerly
       // which may reduce the number of callers of other functions to one,
       // changing inline cost thresholds.
       if (Callee.hasLocalLinkage()) {
         // To check this we also need to nuke any dead constant uses (perhaps
         // made dead by this operation on other functions).
         Callee.removeDeadConstantUsers();
         if (Callee.use_empty() && !CG.isLibFunction(Callee)) {
           Calls.erase(
               std::remove_if(Calls.begin() + i + 1, Calls.end(),
-                             [&Callee](const std::pair<CallSite, int> &Call) {
-                               return Call.first.getCaller() == &Callee;
+                             [&Callee](const std::pair<CallBase *, int> &Call) {
+                               return Call.first->getCaller() == &Callee;
                              }),
               Calls.end());
           // Clear the body and queue the function itself for deletion when we
           // finish inlining and call graph updates.
           // Note that after this point, it is an error to do anything other
           // than use the callee's address or delete it.
           Callee.dropAllReferences();
           assert(find(DeadFunctions, &Callee) == DeadFunctions.end() &&