1 files changed, 6 insertions, 4 deletions
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index f7d40a3be42..e091928609f 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -1549,10 +1549,12 @@ static void replacePrepare(CallInst *Prepare, CallGraph &CG) {
     // If so, we'll need to update the call graph.
     if (PrepareUserNode) {
       for (auto &Use : Cast->uses()) {
-        auto CS = CallSite(Use.getUser());
-        if (!CS || !CS.isCallee(&Use)) continue;
-        PrepareUserNode->removeCallEdgeFor(CS);
-        PrepareUserNode->addCalledFunction(CS, FnNode);
+        if (auto *CB = dyn_cast<CallBase>(Use.getUser())) {
+          if (!CB->isCallee(&Use))
+            continue;
+          PrepareUserNode->removeCallEdgeFor(*CB);
+          PrepareUserNode->addCalledFunction(CB, FnNode);
+        }
       }
     }