[NFCI][IR] Introduce CallBase::Create() wrapper

Summary:
It is reasonably common to want to clone some call with different bundles.
Let's actually provide an interface to do that.

Reviewers: chandlerc, jdoerfert, dblaikie, nickdesaulniers

Reviewed By: nickdesaulniers

Subscribers: llvm-commits, hiraditya

Tags: #llvm

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

5 files changed, 32 insertions, 29 deletions

diff --git a/llvm/include/llvm/IR/InstrTypes.h b/llvm/include/llvm/IR/InstrTypes.h
index 06119f3..770d318 100644
--- a/llvm/include/llvm/IR/InstrTypes.h
+++ b/llvm/include/llvm/IR/InstrTypes.h
@@ -1135,6 +1135,15 @@ protected:
 public:
   using Instruction::getContext;
 
+  /// Create a clone of \p CB with a different set of operand bundles and
+  /// insert it before \p InsertPt.
+  ///
+  /// The returned call instruction is identical \p CB in every way except that
+  /// the operand bundles for the new instruction are set to the operand bundles
+  /// in \p Bundles.
+  static CallBase *Create(CallBase *CB, ArrayRef<OperandBundleDef> Bundles,
+                          Instruction *InsertPt = nullptr);
+
   static bool classof(const Instruction *I) {
     return I->getOpcode() == Instruction::Call ||
            I->getOpcode() == Instruction::Invoke ||
diff --git a/llvm/lib/IR/Instructions.cpp b/llvm/lib/IR/Instructions.cpp
index 78887a6..e22f609 100644
--- a/llvm/lib/IR/Instructions.cpp
+++ b/llvm/lib/IR/Instructions.cpp
@@ -247,6 +247,20 @@ void LandingPadInst::addClause(Constant *Val) {
 //                        CallBase Implementation
 //===----------------------------------------------------------------------===//
 
+CallBase *CallBase::Create(CallBase *CB, ArrayRef<OperandBundleDef> Bundles,
+                           Instruction *InsertPt) {
+  switch (CB->getOpcode()) {
+  case Instruction::Call:
+    return CallInst::Create(cast<CallInst>(CB), Bundles, InsertPt);
+  case Instruction::Invoke:
+    return InvokeInst::Create(cast<InvokeInst>(CB), Bundles, InsertPt);
+  case Instruction::CallBr:
+    return CallBrInst::Create(cast<CallBrInst>(CB), Bundles, InsertPt);
+  default:
+    llvm_unreachable("Unknown CallBase sub-class!");
+  }
+}
+
 Function *CallBase::getCaller() { return getParent()->getParent(); }
 
 unsigned CallBase::getNumSubclassExtraOperandsDynamic() const {
diff --git a/llvm/lib/Transforms/CFGuard/CFGuard.cpp b/llvm/lib/Transforms/CFGuard/CFGuard.cpp
index e4f338b..96c083a 100644
--- a/llvm/lib/Transforms/CFGuard/CFGuard.cpp
+++ b/llvm/lib/Transforms/CFGuard/CFGuard.cpp
@@ -204,14 +204,9 @@ void CFGuard::insertCFGuardDispatch(CallBase *CB) {
   Bundles.emplace_back("cfguardtarget", CalledOperand);
 
   // Create a copy of the call/invoke instruction and add the new bundle.
-  CallBase *NewCB;
-  if (CallInst *CI = dyn_cast<CallInst>(CB)) {
-    NewCB = CallInst::Create(CI, Bundles, CB);
-  } else {
-    assert(isa<InvokeInst>(CB) && "Unknown indirect call type");
-    InvokeInst *II = cast<InvokeInst>(CB);
-    NewCB = llvm::InvokeInst::Create(II, Bundles, CB);
-  }
+  assert((isa<CallInst>(CB) || isa<InvokeInst>(CB)) &&
+         "Unknown indirect call type");
+  CallBase *NewCB = CallBase::Create(CB, Bundles, CB);
 
   // Change the target of the call to be the guard dispatch function.
   NewCB->setCalledOperand(GuardDispatchLoad);
diff --git a/llvm/lib/Transforms/IPO/GlobalOpt.cpp b/llvm/lib/Transforms/IPO/GlobalOpt.cpp
index 437451b..ed98952 100644
--- a/llvm/lib/Transforms/IPO/GlobalOpt.cpp
+++ b/llvm/lib/Transforms/IPO/GlobalOpt.cpp
@@ -2321,13 +2321,10 @@ static void RemovePreallocated(Function *F) {
     assert(PreallocatedSetup && "Did not find preallocated bundle");
     uint64_t ArgCount =
         cast<ConstantInt>(PreallocatedSetup->getArgOperand(0))->getZExtValue();
-    CallBase *NewCB = nullptr;
-    if (InvokeInst *II = dyn_cast<InvokeInst>(CB)) {
-      NewCB = InvokeInst::Create(II, OpBundles, CB);
-    } else {
-      CallInst *CI = cast<CallInst>(CB);
-      NewCB = CallInst::Create(CI, OpBundles, CB);
-    }
+
+    assert((isa<CallInst>(CB) || isa<InvokeInst>(CB)) &&
+           "Unknown indirect call type");
+    CallBase *NewCB = CallBase::Create(CB, OpBundles, CB);
     CB->replaceAllUsesWith(NewCB);
     NewCB->takeName(CB);
     CB->eraseFromParent();
diff --git a/llvm/lib/Transforms/Utils/InlineFunction.cpp b/llvm/lib/Transforms/Utils/InlineFunction.cpp
index 203e812..b0b7ca4 100644
--- a/llvm/lib/Transforms/Utils/InlineFunction.cpp
+++ b/llvm/lib/Transforms/Utils/InlineFunction.cpp
@@ -1864,13 +1864,7 @@ llvm::InlineResult llvm::InlineFunction(CallBase &CB, InlineFunctionInfo &IFI,
           OpDefs.emplace_back("deopt", std::move(MergedDeoptArgs));
         }
 
-        Instruction *NewI = nullptr;
-        if (isa<CallInst>(ICS))
-          NewI = CallInst::Create(cast<CallInst>(ICS), OpDefs, ICS);
-        else if (isa<CallBrInst>(ICS))
-          NewI = CallBrInst::Create(cast<CallBrInst>(ICS), OpDefs, ICS);
-        else
-          NewI = InvokeInst::Create(cast<InvokeInst>(ICS), OpDefs, ICS);
+        Instruction *NewI = CallBase::Create(ICS, OpDefs, ICS);
 
         // Note: the RAUW does the appropriate fixup in VMap, so we need to do
         // this even if the call returns void.
@@ -2166,13 +2160,7 @@ llvm::InlineResult llvm::InlineFunction(CallBase &CB, InlineFunctionInfo &IFI,
         I->getOperandBundlesAsDefs(OpBundles);
         OpBundles.emplace_back("funclet", CallSiteEHPad);
 
-        Instruction *NewInst;
-        if (auto *CallI = dyn_cast<CallInst>(I))
-          NewInst = CallInst::Create(CallI, OpBundles, CallI);
-        else if (auto *CallBrI = dyn_cast<CallBrInst>(I))
-          NewInst = CallBrInst::Create(CallBrI, OpBundles, CallBrI);
-        else
-          NewInst = InvokeInst::Create(cast<InvokeInst>(I), OpBundles, I);
+        Instruction *NewInst = CallBase::Create(I, OpBundles, I);
         NewInst->takeName(I);
         I->replaceAllUsesWith(NewInst);
         I->eraseFromParent();