1 //===- CoroSplit.cpp - Converts a coroutine into a state machine ----------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 // This pass builds the coroutine frame and outlines resume and destroy parts
9 // of the coroutine into separate functions.
10 //
11 // We present a coroutine to an LLVM as an ordinary function with suspension
12 // points marked up with intrinsics. We let the optimizer party on the coroutine
13 // as a single function for as long as possible. Shortly before the coroutine is
14 // eligible to be inlined into its callers, we split up the coroutine into parts
15 // corresponding to an initial, resume and destroy invocations of the coroutine,
16 // add them to the current SCC and restart the IPO pipeline to optimize the
17 // coroutine subfunctions we extracted before proceeding to the caller of the
18 // coroutine.
19 //===----------------------------------------------------------------------===//
20 
21 #include "llvm/Transforms/Coroutines/CoroSplit.h"
22 #include "CoroInstr.h"
23 #include "CoroInternal.h"
24 #include "llvm/ADT/DenseMap.h"
25 #include "llvm/ADT/SmallPtrSet.h"
26 #include "llvm/ADT/SmallVector.h"
27 #include "llvm/ADT/StringRef.h"
28 #include "llvm/ADT/Twine.h"
29 #include "llvm/Analysis/CallGraph.h"
30 #include "llvm/Analysis/CallGraphSCCPass.h"
31 #include "llvm/IR/Argument.h"
32 #include "llvm/IR/Attributes.h"
33 #include "llvm/IR/BasicBlock.h"
34 #include "llvm/IR/CFG.h"
35 #include "llvm/IR/CallingConv.h"
36 #include "llvm/IR/Constants.h"
37 #include "llvm/IR/DataLayout.h"
38 #include "llvm/IR/DerivedTypes.h"
39 #include "llvm/IR/Function.h"
40 #include "llvm/IR/GlobalValue.h"
41 #include "llvm/IR/GlobalVariable.h"
42 #include "llvm/IR/IRBuilder.h"
43 #include "llvm/IR/InstIterator.h"
44 #include "llvm/IR/InstrTypes.h"
45 #include "llvm/IR/Instruction.h"
46 #include "llvm/IR/Instructions.h"
47 #include "llvm/IR/IntrinsicInst.h"
48 #include "llvm/IR/LLVMContext.h"
49 #include "llvm/IR/LegacyPassManager.h"
50 #include "llvm/IR/Module.h"
51 #include "llvm/IR/Type.h"
52 #include "llvm/IR/Value.h"
53 #include "llvm/IR/Verifier.h"
54 #include "llvm/InitializePasses.h"
55 #include "llvm/Pass.h"
56 #include "llvm/Support/Casting.h"
57 #include "llvm/Support/Debug.h"
58 #include "llvm/Support/PrettyStackTrace.h"
59 #include "llvm/Support/raw_ostream.h"
60 #include "llvm/Transforms/Scalar.h"
61 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
62 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
63 #include "llvm/Transforms/Utils/Cloning.h"
64 #include "llvm/Transforms/Utils/Local.h"
65 #include "llvm/Transforms/Utils/ValueMapper.h"
66 #include <cassert>
67 #include <cstddef>
68 #include <cstdint>
69 #include <initializer_list>
70 #include <iterator>
71 
72 using namespace llvm;
73 
74 #define DEBUG_TYPE "coro-split"
75 
76 namespace {
77 
78 /// A little helper class for building
79 class CoroCloner {
80 public:
81   enum class Kind {
82     /// The shared resume function for a switch lowering.
83     SwitchResume,
84 
85     /// The shared unwind function for a switch lowering.
86     SwitchUnwind,
87 
88     /// The shared cleanup function for a switch lowering.
89     SwitchCleanup,
90 
91     /// An individual continuation function.
92     Continuation,
93   };
94 private:
95   Function &OrigF;
96   Function *NewF;
97   const Twine &Suffix;
98   coro::Shape &Shape;
99   Kind FKind;
100   ValueToValueMapTy VMap;
101   IRBuilder<> Builder;
102   Value *NewFramePtr = nullptr;
103   Value *SwiftErrorSlot = nullptr;
104 
105   /// The active suspend instruction; meaningful only for continuation ABIs.
106   AnyCoroSuspendInst *ActiveSuspend = nullptr;
107 
108 public:
109   /// Create a cloner for a switch lowering.
110   CoroCloner(Function &OrigF, const Twine &Suffix, coro::Shape &Shape,
111              Kind FKind)
112     : OrigF(OrigF), NewF(nullptr), Suffix(Suffix), Shape(Shape),
113       FKind(FKind), Builder(OrigF.getContext()) {
114     assert(Shape.ABI == coro::ABI::Switch);
115   }
116 
117   /// Create a cloner for a continuation lowering.
118   CoroCloner(Function &OrigF, const Twine &Suffix, coro::Shape &Shape,
119              Function *NewF, AnyCoroSuspendInst *ActiveSuspend)
120     : OrigF(OrigF), NewF(NewF), Suffix(Suffix), Shape(Shape),
121       FKind(Kind::Continuation), Builder(OrigF.getContext()),
122       ActiveSuspend(ActiveSuspend) {
123     assert(Shape.ABI == coro::ABI::Retcon ||
124            Shape.ABI == coro::ABI::RetconOnce);
125     assert(NewF && "need existing function for continuation");
126     assert(ActiveSuspend && "need active suspend point for continuation");
127   }
128 
129   Function *getFunction() const {
130     assert(NewF != nullptr && "declaration not yet set");
131     return NewF;
132   }
133 
134   void create();
135 
136 private:
137   bool isSwitchDestroyFunction() {
138     switch (FKind) {
139     case Kind::Continuation:
140     case Kind::SwitchResume:
141       return false;
142     case Kind::SwitchUnwind:
143     case Kind::SwitchCleanup:
144       return true;
145     }
146     llvm_unreachable("Unknown CoroCloner::Kind enum");
147   }
148 
149   void createDeclaration();
150   void replaceEntryBlock();
151   Value *deriveNewFramePointer();
152   void replaceRetconSuspendUses();
153   void replaceCoroSuspends();
154   void replaceCoroEnds();
155   void replaceSwiftErrorOps();
156   void handleFinalSuspend();
157   void maybeFreeContinuationStorage();
158 };
159 
160 } // end anonymous namespace
161 
162 static void maybeFreeRetconStorage(IRBuilder<> &Builder,
163                                    const coro::Shape &Shape, Value *FramePtr,
164                                    CallGraph *CG) {
165   assert(Shape.ABI == coro::ABI::Retcon ||
166          Shape.ABI == coro::ABI::RetconOnce);
167   if (Shape.RetconLowering.IsFrameInlineInStorage)
168     return;
169 
170   Shape.emitDealloc(Builder, FramePtr, CG);
171 }
172 
173 /// Replace a non-unwind call to llvm.coro.end.
174 static void replaceFallthroughCoroEnd(CoroEndInst *End,
175                                       const coro::Shape &Shape, Value *FramePtr,
176                                       bool InResume, CallGraph *CG) {
177   // Start inserting right before the coro.end.
178   IRBuilder<> Builder(End);
179 
180   // Create the return instruction.
181   switch (Shape.ABI) {
182   // The cloned functions in switch-lowering always return void.
183   case coro::ABI::Switch:
184     // coro.end doesn't immediately end the coroutine in the main function
185     // in this lowering, because we need to deallocate the coroutine.
186     if (!InResume)
187       return;
188     Builder.CreateRetVoid();
189     break;
190 
191   // In unique continuation lowering, the continuations always return void.
192   // But we may have implicitly allocated storage.
193   case coro::ABI::RetconOnce:
194     maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
195     Builder.CreateRetVoid();
196     break;
197 
198   // In non-unique continuation lowering, we signal completion by returning
199   // a null continuation.
200   case coro::ABI::Retcon: {
201     maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
202     auto RetTy = Shape.getResumeFunctionType()->getReturnType();
203     auto RetStructTy = dyn_cast<StructType>(RetTy);
204     PointerType *ContinuationTy =
205       cast<PointerType>(RetStructTy ? RetStructTy->getElementType(0) : RetTy);
206 
207     Value *ReturnValue = ConstantPointerNull::get(ContinuationTy);
208     if (RetStructTy) {
209       ReturnValue = Builder.CreateInsertValue(UndefValue::get(RetStructTy),
210                                               ReturnValue, 0);
211     }
212     Builder.CreateRet(ReturnValue);
213     break;
214   }
215   }
216 
217   // Remove the rest of the block, by splitting it into an unreachable block.
218   auto *BB = End->getParent();
219   BB->splitBasicBlock(End);
220   BB->getTerminator()->eraseFromParent();
221 }
222 
223 /// Replace an unwind call to llvm.coro.end.
224 static void replaceUnwindCoroEnd(CoroEndInst *End, const coro::Shape &Shape,
225                                  Value *FramePtr, bool InResume, CallGraph *CG){
226   IRBuilder<> Builder(End);
227 
228   switch (Shape.ABI) {
229   // In switch-lowering, this does nothing in the main function.
230   case coro::ABI::Switch:
231     if (!InResume)
232       return;
233     break;
234 
235   // In continuation-lowering, this frees the continuation storage.
236   case coro::ABI::Retcon:
237   case coro::ABI::RetconOnce:
238     maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
239     break;
240   }
241 
242   // If coro.end has an associated bundle, add cleanupret instruction.
243   if (auto Bundle = End->getOperandBundle(LLVMContext::OB_funclet)) {
244     auto *FromPad = cast<CleanupPadInst>(Bundle->Inputs[0]);
245     auto *CleanupRet = Builder.CreateCleanupRet(FromPad, nullptr);
246     End->getParent()->splitBasicBlock(End);
247     CleanupRet->getParent()->getTerminator()->eraseFromParent();
248   }
249 }
250 
251 static void replaceCoroEnd(CoroEndInst *End, const coro::Shape &Shape,
252                            Value *FramePtr, bool InResume, CallGraph *CG) {
253   if (End->isUnwind())
254     replaceUnwindCoroEnd(End, Shape, FramePtr, InResume, CG);
255   else
256     replaceFallthroughCoroEnd(End, Shape, FramePtr, InResume, CG);
257 
258   auto &Context = End->getContext();
259   End->replaceAllUsesWith(InResume ? ConstantInt::getTrue(Context)
260                                    : ConstantInt::getFalse(Context));
261   End->eraseFromParent();
262 }
263 
264 // Create an entry block for a resume function with a switch that will jump to
265 // suspend points.
266 static void createResumeEntryBlock(Function &F, coro::Shape &Shape) {
267   assert(Shape.ABI == coro::ABI::Switch);
268   LLVMContext &C = F.getContext();
269 
270   // resume.entry:
271   //  %index.addr = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0,
272   //  i32 2
273   //  % index = load i32, i32* %index.addr
274   //  switch i32 %index, label %unreachable [
275   //    i32 0, label %resume.0
276   //    i32 1, label %resume.1
277   //    ...
278   //  ]
279 
280   auto *NewEntry = BasicBlock::Create(C, "resume.entry", &F);
281   auto *UnreachBB = BasicBlock::Create(C, "unreachable", &F);
282 
283   IRBuilder<> Builder(NewEntry);
284   auto *FramePtr = Shape.FramePtr;
285   auto *FrameTy = Shape.FrameTy;
286   auto *GepIndex = Builder.CreateStructGEP(
287       FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");
288   auto *Index = Builder.CreateLoad(Shape.getIndexType(), GepIndex, "index");
289   auto *Switch =
290       Builder.CreateSwitch(Index, UnreachBB, Shape.CoroSuspends.size());
291   Shape.SwitchLowering.ResumeSwitch = Switch;
292 
293   size_t SuspendIndex = 0;
294   for (auto *AnyS : Shape.CoroSuspends) {
295     auto *S = cast<CoroSuspendInst>(AnyS);
296     ConstantInt *IndexVal = Shape.getIndex(SuspendIndex);
297 
298     // Replace CoroSave with a store to Index:
299     //    %index.addr = getelementptr %f.frame... (index field number)
300     //    store i32 0, i32* %index.addr1
301     auto *Save = S->getCoroSave();
302     Builder.SetInsertPoint(Save);
303     if (S->isFinal()) {
304       // Final suspend point is represented by storing zero in ResumeFnAddr.
305       auto *GepIndex = Builder.CreateStructGEP(FrameTy, FramePtr,
306                                  coro::Shape::SwitchFieldIndex::Resume,
307                                   "ResumeFn.addr");
308       auto *NullPtr = ConstantPointerNull::get(cast<PointerType>(
309           cast<PointerType>(GepIndex->getType())->getElementType()));
310       Builder.CreateStore(NullPtr, GepIndex);
311     } else {
312       auto *GepIndex = Builder.CreateStructGEP(
313           FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");
314       Builder.CreateStore(IndexVal, GepIndex);
315     }
316     Save->replaceAllUsesWith(ConstantTokenNone::get(C));
317     Save->eraseFromParent();
318 
319     // Split block before and after coro.suspend and add a jump from an entry
320     // switch:
321     //
322     //  whateverBB:
323     //    whatever
324     //    %0 = call i8 @llvm.coro.suspend(token none, i1 false)
325     //    switch i8 %0, label %suspend[i8 0, label %resume
326     //                                 i8 1, label %cleanup]
327     // becomes:
328     //
329     //  whateverBB:
330     //     whatever
331     //     br label %resume.0.landing
332     //
333     //  resume.0: ; <--- jump from the switch in the resume.entry
334     //     %0 = tail call i8 @llvm.coro.suspend(token none, i1 false)
335     //     br label %resume.0.landing
336     //
337     //  resume.0.landing:
338     //     %1 = phi i8[-1, %whateverBB], [%0, %resume.0]
339     //     switch i8 % 1, label %suspend [i8 0, label %resume
340     //                                    i8 1, label %cleanup]
341 
342     auto *SuspendBB = S->getParent();
343     auto *ResumeBB =
344         SuspendBB->splitBasicBlock(S, "resume." + Twine(SuspendIndex));
345     auto *LandingBB = ResumeBB->splitBasicBlock(
346         S->getNextNode(), ResumeBB->getName() + Twine(".landing"));
347     Switch->addCase(IndexVal, ResumeBB);
348 
349     cast<BranchInst>(SuspendBB->getTerminator())->setSuccessor(0, LandingBB);
350     auto *PN = PHINode::Create(Builder.getInt8Ty(), 2, "", &LandingBB->front());
351     S->replaceAllUsesWith(PN);
352     PN->addIncoming(Builder.getInt8(-1), SuspendBB);
353     PN->addIncoming(S, ResumeBB);
354 
355     ++SuspendIndex;
356   }
357 
358   Builder.SetInsertPoint(UnreachBB);
359   Builder.CreateUnreachable();
360 
361   Shape.SwitchLowering.ResumeEntryBlock = NewEntry;
362 }
363 
364 
365 // Rewrite final suspend point handling. We do not use suspend index to
366 // represent the final suspend point. Instead we zero-out ResumeFnAddr in the
367 // coroutine frame, since it is undefined behavior to resume a coroutine
368 // suspended at the final suspend point. Thus, in the resume function, we can
369 // simply remove the last case (when coro::Shape is built, the final suspend
370 // point (if present) is always the last element of CoroSuspends array).
371 // In the destroy function, we add a code sequence to check if ResumeFnAddress
372 // is Null, and if so, jump to the appropriate label to handle cleanup from the
373 // final suspend point.
374 void CoroCloner::handleFinalSuspend() {
375   assert(Shape.ABI == coro::ABI::Switch &&
376          Shape.SwitchLowering.HasFinalSuspend);
377   auto *Switch = cast<SwitchInst>(VMap[Shape.SwitchLowering.ResumeSwitch]);
378   auto FinalCaseIt = std::prev(Switch->case_end());
379   BasicBlock *ResumeBB = FinalCaseIt->getCaseSuccessor();
380   Switch->removeCase(FinalCaseIt);
381   if (isSwitchDestroyFunction()) {
382     BasicBlock *OldSwitchBB = Switch->getParent();
383     auto *NewSwitchBB = OldSwitchBB->splitBasicBlock(Switch, "Switch");
384     Builder.SetInsertPoint(OldSwitchBB->getTerminator());
385     auto *GepIndex = Builder.CreateStructGEP(Shape.FrameTy, NewFramePtr,
386                                        coro::Shape::SwitchFieldIndex::Resume,
387                                              "ResumeFn.addr");
388     auto *Load = Builder.CreateLoad(Shape.getSwitchResumePointerType(),
389                                     GepIndex);
390     auto *Cond = Builder.CreateIsNull(Load);
391     Builder.CreateCondBr(Cond, ResumeBB, NewSwitchBB);
392     OldSwitchBB->getTerminator()->eraseFromParent();
393   }
394 }
395 
396 static Function *createCloneDeclaration(Function &OrigF, coro::Shape &Shape,
397                                         const Twine &Suffix,
398                                         Module::iterator InsertBefore) {
399   Module *M = OrigF.getParent();
400   auto *FnTy = Shape.getResumeFunctionType();
401 
402   Function *NewF =
403       Function::Create(FnTy, GlobalValue::LinkageTypes::InternalLinkage,
404                        OrigF.getName() + Suffix);
405   NewF->addParamAttr(0, Attribute::NonNull);
406   NewF->addParamAttr(0, Attribute::NoAlias);
407 
408   M->getFunctionList().insert(InsertBefore, NewF);
409 
410   return NewF;
411 }
412 
413 /// Replace uses of the active llvm.coro.suspend.retcon call with the
414 /// arguments to the continuation function.
415 ///
416 /// This assumes that the builder has a meaningful insertion point.
417 void CoroCloner::replaceRetconSuspendUses() {
418   assert(Shape.ABI == coro::ABI::Retcon ||
419          Shape.ABI == coro::ABI::RetconOnce);
420 
421   auto NewS = VMap[ActiveSuspend];
422   if (NewS->use_empty()) return;
423 
424   // Copy out all the continuation arguments after the buffer pointer into
425   // an easily-indexed data structure for convenience.
426   SmallVector<Value*, 8> Args;
427   for (auto I = std::next(NewF->arg_begin()), E = NewF->arg_end(); I != E; ++I)
428     Args.push_back(&*I);
429 
430   // If the suspend returns a single scalar value, we can just do a simple
431   // replacement.
432   if (!isa<StructType>(NewS->getType())) {
433     assert(Args.size() == 1);
434     NewS->replaceAllUsesWith(Args.front());
435     return;
436   }
437 
438   // Try to peephole extracts of an aggregate return.
439   for (auto UI = NewS->use_begin(), UE = NewS->use_end(); UI != UE; ) {
440     auto EVI = dyn_cast<ExtractValueInst>((UI++)->getUser());
441     if (!EVI || EVI->getNumIndices() != 1)
442       continue;
443 
444     EVI->replaceAllUsesWith(Args[EVI->getIndices().front()]);
445     EVI->eraseFromParent();
446   }
447 
448   // If we have no remaining uses, we're done.
449   if (NewS->use_empty()) return;
450 
451   // Otherwise, we need to create an aggregate.
452   Value *Agg = UndefValue::get(NewS->getType());
453   for (size_t I = 0, E = Args.size(); I != E; ++I)
454     Agg = Builder.CreateInsertValue(Agg, Args[I], I);
455 
456   NewS->replaceAllUsesWith(Agg);
457 }
458 
459 void CoroCloner::replaceCoroSuspends() {
460   Value *SuspendResult;
461 
462   switch (Shape.ABI) {
463   // In switch lowering, replace coro.suspend with the appropriate value
464   // for the type of function we're extracting.
465   // Replacing coro.suspend with (0) will result in control flow proceeding to
466   // a resume label associated with a suspend point, replacing it with (1) will
467   // result in control flow proceeding to a cleanup label associated with this
468   // suspend point.
469   case coro::ABI::Switch:
470     SuspendResult = Builder.getInt8(isSwitchDestroyFunction() ? 1 : 0);
471     break;
472 
473   // In returned-continuation lowering, the arguments from earlier
474   // continuations are theoretically arbitrary, and they should have been
475   // spilled.
476   case coro::ABI::RetconOnce:
477   case coro::ABI::Retcon:
478     return;
479   }
480 
481   for (AnyCoroSuspendInst *CS : Shape.CoroSuspends) {
482     // The active suspend was handled earlier.
483     if (CS == ActiveSuspend) continue;
484 
485     auto *MappedCS = cast<AnyCoroSuspendInst>(VMap[CS]);
486     MappedCS->replaceAllUsesWith(SuspendResult);
487     MappedCS->eraseFromParent();
488   }
489 }
490 
491 void CoroCloner::replaceCoroEnds() {
492   for (CoroEndInst *CE : Shape.CoroEnds) {
493     // We use a null call graph because there's no call graph node for
494     // the cloned function yet.  We'll just be rebuilding that later.
495     auto NewCE = cast<CoroEndInst>(VMap[CE]);
496     replaceCoroEnd(NewCE, Shape, NewFramePtr, /*in resume*/ true, nullptr);
497   }
498 }
499 
500 static void replaceSwiftErrorOps(Function &F, coro::Shape &Shape,
501                                  ValueToValueMapTy *VMap) {
502   Value *CachedSlot = nullptr;
503   auto getSwiftErrorSlot = [&](Type *ValueTy) -> Value * {
504     if (CachedSlot) {
505       assert(CachedSlot->getType()->getPointerElementType() == ValueTy &&
506              "multiple swifterror slots in function with different types");
507       return CachedSlot;
508     }
509 
510     // Check if the function has a swifterror argument.
511     for (auto &Arg : F.args()) {
512       if (Arg.isSwiftError()) {
513         CachedSlot = &Arg;
514         assert(Arg.getType()->getPointerElementType() == ValueTy &&
515                "swifterror argument does not have expected type");
516         return &Arg;
517       }
518     }
519 
520     // Create a swifterror alloca.
521     IRBuilder<> Builder(F.getEntryBlock().getFirstNonPHIOrDbg());
522     auto Alloca = Builder.CreateAlloca(ValueTy);
523     Alloca->setSwiftError(true);
524 
525     CachedSlot = Alloca;
526     return Alloca;
527   };
528 
529   for (CallInst *Op : Shape.SwiftErrorOps) {
530     auto MappedOp = VMap ? cast<CallInst>((*VMap)[Op]) : Op;
531     IRBuilder<> Builder(MappedOp);
532 
533     // If there are no arguments, this is a 'get' operation.
534     Value *MappedResult;
535     if (Op->getNumArgOperands() == 0) {
536       auto ValueTy = Op->getType();
537       auto Slot = getSwiftErrorSlot(ValueTy);
538       MappedResult = Builder.CreateLoad(ValueTy, Slot);
539     } else {
540       assert(Op->getNumArgOperands() == 1);
541       auto Value = MappedOp->getArgOperand(0);
542       auto ValueTy = Value->getType();
543       auto Slot = getSwiftErrorSlot(ValueTy);
544       Builder.CreateStore(Value, Slot);
545       MappedResult = Slot;
546     }
547 
548     MappedOp->replaceAllUsesWith(MappedResult);
549     MappedOp->eraseFromParent();
550   }
551 
552   // If we're updating the original function, we've invalidated SwiftErrorOps.
553   if (VMap == nullptr) {
554     Shape.SwiftErrorOps.clear();
555   }
556 }
557 
558 void CoroCloner::replaceSwiftErrorOps() {
559   ::replaceSwiftErrorOps(*NewF, Shape, &VMap);
560 }
561 
562 void CoroCloner::replaceEntryBlock() {
563   // In the original function, the AllocaSpillBlock is a block immediately
564   // following the allocation of the frame object which defines GEPs for
565   // all the allocas that have been moved into the frame, and it ends by
566   // branching to the original beginning of the coroutine.  Make this
567   // the entry block of the cloned function.
568   auto *Entry = cast<BasicBlock>(VMap[Shape.AllocaSpillBlock]);
569   auto *OldEntry = &NewF->getEntryBlock();
570   Entry->setName("entry" + Suffix);
571   Entry->moveBefore(OldEntry);
572   Entry->getTerminator()->eraseFromParent();
573 
574   // Clear all predecessors of the new entry block.  There should be
575   // exactly one predecessor, which we created when splitting out
576   // AllocaSpillBlock to begin with.
577   assert(Entry->hasOneUse());
578   auto BranchToEntry = cast<BranchInst>(Entry->user_back());
579   assert(BranchToEntry->isUnconditional());
580   Builder.SetInsertPoint(BranchToEntry);
581   Builder.CreateUnreachable();
582   BranchToEntry->eraseFromParent();
583 
584   // Move any allocas into Entry that weren't moved into the frame.
585   for (auto IT = OldEntry->begin(), End = OldEntry->end(); IT != End;) {
586     Instruction &I = *IT++;
587     if (!isa<AllocaInst>(&I) || I.use_empty())
588       continue;
589 
590     I.moveBefore(*Entry, Entry->getFirstInsertionPt());
591   }
592 
593   // Branch from the entry to the appropriate place.
594   Builder.SetInsertPoint(Entry);
595   switch (Shape.ABI) {
596   case coro::ABI::Switch: {
597     // In switch-lowering, we built a resume-entry block in the original
598     // function.  Make the entry block branch to this.
599     auto *SwitchBB =
600       cast<BasicBlock>(VMap[Shape.SwitchLowering.ResumeEntryBlock]);
601     Builder.CreateBr(SwitchBB);
602     break;
603   }
604 
605   case coro::ABI::Retcon:
606   case coro::ABI::RetconOnce: {
607     // In continuation ABIs, we want to branch to immediately after the
608     // active suspend point.  Earlier phases will have put the suspend in its
609     // own basic block, so just thread our jump directly to its successor.
610     auto MappedCS = cast<CoroSuspendRetconInst>(VMap[ActiveSuspend]);
611     auto Branch = cast<BranchInst>(MappedCS->getNextNode());
612     assert(Branch->isUnconditional());
613     Builder.CreateBr(Branch->getSuccessor(0));
614     break;
615   }
616   }
617 }
618 
619 /// Derive the value of the new frame pointer.
620 Value *CoroCloner::deriveNewFramePointer() {
621   // Builder should be inserting to the front of the new entry block.
622 
623   switch (Shape.ABI) {
624   // In switch-lowering, the argument is the frame pointer.
625   case coro::ABI::Switch:
626     return &*NewF->arg_begin();
627 
628   // In continuation-lowering, the argument is the opaque storage.
629   case coro::ABI::Retcon:
630   case coro::ABI::RetconOnce: {
631     Argument *NewStorage = &*NewF->arg_begin();
632     auto FramePtrTy = Shape.FrameTy->getPointerTo();
633 
634     // If the storage is inline, just bitcast to the storage to the frame type.
635     if (Shape.RetconLowering.IsFrameInlineInStorage)
636       return Builder.CreateBitCast(NewStorage, FramePtrTy);
637 
638     // Otherwise, load the real frame from the opaque storage.
639     auto FramePtrPtr =
640       Builder.CreateBitCast(NewStorage, FramePtrTy->getPointerTo());
641     return Builder.CreateLoad(FramePtrTy, FramePtrPtr);
642   }
643   }
644   llvm_unreachable("bad ABI");
645 }
646 
647 static void addFramePointerAttrs(AttributeList &Attrs, LLVMContext &Context,
648                                  unsigned ParamIndex,
649                                  uint64_t Size, Align Alignment) {
650   AttrBuilder ParamAttrs;
651   ParamAttrs.addAttribute(Attribute::NonNull);
652   ParamAttrs.addAttribute(Attribute::NoAlias);
653   ParamAttrs.addAlignmentAttr(Alignment);
654   ParamAttrs.addDereferenceableAttr(Size);
655   Attrs = Attrs.addParamAttributes(Context, ParamIndex, ParamAttrs);
656 }
657 
658 /// Clone the body of the original function into a resume function of
659 /// some sort.
660 void CoroCloner::create() {
661   // Create the new function if we don't already have one.
662   if (!NewF) {
663     NewF = createCloneDeclaration(OrigF, Shape, Suffix,
664                                   OrigF.getParent()->end());
665   }
666 
667   // Replace all args with undefs. The buildCoroutineFrame algorithm already
668   // rewritten access to the args that occurs after suspend points with loads
669   // and stores to/from the coroutine frame.
670   for (Argument &A : OrigF.args())
671     VMap[&A] = UndefValue::get(A.getType());
672 
673   SmallVector<ReturnInst *, 4> Returns;
674 
675   // Ignore attempts to change certain attributes of the function.
676   // TODO: maybe there should be a way to suppress this during cloning?
677   auto savedVisibility = NewF->getVisibility();
678   auto savedUnnamedAddr = NewF->getUnnamedAddr();
679   auto savedDLLStorageClass = NewF->getDLLStorageClass();
680 
681   // NewF's linkage (which CloneFunctionInto does *not* change) might not
682   // be compatible with the visibility of OrigF (which it *does* change),
683   // so protect against that.
684   auto savedLinkage = NewF->getLinkage();
685   NewF->setLinkage(llvm::GlobalValue::ExternalLinkage);
686 
687   CloneFunctionInto(NewF, &OrigF, VMap, /*ModuleLevelChanges=*/true, Returns);
688 
689   NewF->setLinkage(savedLinkage);
690   NewF->setVisibility(savedVisibility);
691   NewF->setUnnamedAddr(savedUnnamedAddr);
692   NewF->setDLLStorageClass(savedDLLStorageClass);
693 
694   auto &Context = NewF->getContext();
695 
696   // Replace the attributes of the new function:
697   auto OrigAttrs = NewF->getAttributes();
698   auto NewAttrs = AttributeList();
699 
700   switch (Shape.ABI) {
701   case coro::ABI::Switch:
702     // Bootstrap attributes by copying function attributes from the
703     // original function.  This should include optimization settings and so on.
704     NewAttrs = NewAttrs.addAttributes(Context, AttributeList::FunctionIndex,
705                                       OrigAttrs.getFnAttributes());
706 
707     addFramePointerAttrs(NewAttrs, Context, 0,
708                          Shape.FrameSize, Shape.FrameAlign);
709     break;
710 
711   case coro::ABI::Retcon:
712   case coro::ABI::RetconOnce:
713     // If we have a continuation prototype, just use its attributes,
714     // full-stop.
715     NewAttrs = Shape.RetconLowering.ResumePrototype->getAttributes();
716 
717     addFramePointerAttrs(NewAttrs, Context, 0,
718                          Shape.getRetconCoroId()->getStorageSize(),
719                          Shape.getRetconCoroId()->getStorageAlignment());
720     break;
721   }
722 
723   switch (Shape.ABI) {
724   // In these ABIs, the cloned functions always return 'void', and the
725   // existing return sites are meaningless.  Note that for unique
726   // continuations, this includes the returns associated with suspends;
727   // this is fine because we can't suspend twice.
728   case coro::ABI::Switch:
729   case coro::ABI::RetconOnce:
730     // Remove old returns.
731     for (ReturnInst *Return : Returns)
732       changeToUnreachable(Return, /*UseLLVMTrap=*/false);
733     break;
734 
735   // With multi-suspend continuations, we'll already have eliminated the
736   // original returns and inserted returns before all the suspend points,
737   // so we want to leave any returns in place.
738   case coro::ABI::Retcon:
739     break;
740   }
741 
742   NewF->setAttributes(NewAttrs);
743   NewF->setCallingConv(Shape.getResumeFunctionCC());
744 
745   // Set up the new entry block.
746   replaceEntryBlock();
747 
748   Builder.SetInsertPoint(&NewF->getEntryBlock().front());
749   NewFramePtr = deriveNewFramePointer();
750 
751   // Remap frame pointer.
752   Value *OldFramePtr = VMap[Shape.FramePtr];
753   NewFramePtr->takeName(OldFramePtr);
754   OldFramePtr->replaceAllUsesWith(NewFramePtr);
755 
756   // Remap vFrame pointer.
757   auto *NewVFrame = Builder.CreateBitCast(
758       NewFramePtr, Type::getInt8PtrTy(Builder.getContext()), "vFrame");
759   Value *OldVFrame = cast<Value>(VMap[Shape.CoroBegin]);
760   OldVFrame->replaceAllUsesWith(NewVFrame);
761 
762   switch (Shape.ABI) {
763   case coro::ABI::Switch:
764     // Rewrite final suspend handling as it is not done via switch (allows to
765     // remove final case from the switch, since it is undefined behavior to
766     // resume the coroutine suspended at the final suspend point.
767     if (Shape.SwitchLowering.HasFinalSuspend)
768       handleFinalSuspend();
769     break;
770 
771   case coro::ABI::Retcon:
772   case coro::ABI::RetconOnce:
773     // Replace uses of the active suspend with the corresponding
774     // continuation-function arguments.
775     assert(ActiveSuspend != nullptr &&
776            "no active suspend when lowering a continuation-style coroutine");
777     replaceRetconSuspendUses();
778     break;
779   }
780 
781   // Handle suspends.
782   replaceCoroSuspends();
783 
784   // Handle swifterror.
785   replaceSwiftErrorOps();
786 
787   // Remove coro.end intrinsics.
788   replaceCoroEnds();
789 
790   // Eliminate coro.free from the clones, replacing it with 'null' in cleanup,
791   // to suppress deallocation code.
792   if (Shape.ABI == coro::ABI::Switch)
793     coro::replaceCoroFree(cast<CoroIdInst>(VMap[Shape.CoroBegin->getId()]),
794                           /*Elide=*/ FKind == CoroCloner::Kind::SwitchCleanup);
795 }
796 
797 // Create a resume clone by cloning the body of the original function, setting
798 // new entry block and replacing coro.suspend an appropriate value to force
799 // resume or cleanup pass for every suspend point.
800 static Function *createClone(Function &F, const Twine &Suffix,
801                              coro::Shape &Shape, CoroCloner::Kind FKind) {
802   CoroCloner Cloner(F, Suffix, Shape, FKind);
803   Cloner.create();
804   return Cloner.getFunction();
805 }
806 
807 /// Remove calls to llvm.coro.end in the original function.
808 static void removeCoroEnds(const coro::Shape &Shape, CallGraph *CG) {
809   for (auto End : Shape.CoroEnds) {
810     replaceCoroEnd(End, Shape, Shape.FramePtr, /*in resume*/ false, CG);
811   }
812 }
813 
814 static void replaceFrameSize(coro::Shape &Shape) {
815   if (Shape.CoroSizes.empty())
816     return;
817 
818   // In the same function all coro.sizes should have the same result type.
819   auto *SizeIntrin = Shape.CoroSizes.back();
820   Module *M = SizeIntrin->getModule();
821   const DataLayout &DL = M->getDataLayout();
822   auto Size = DL.getTypeAllocSize(Shape.FrameTy);
823   auto *SizeConstant = ConstantInt::get(SizeIntrin->getType(), Size);
824 
825   for (CoroSizeInst *CS : Shape.CoroSizes) {
826     CS->replaceAllUsesWith(SizeConstant);
827     CS->eraseFromParent();
828   }
829 }
830 
831 // Create a global constant array containing pointers to functions provided and
832 // set Info parameter of CoroBegin to point at this constant. Example:
833 //
834 //   @f.resumers = internal constant [2 x void(%f.frame*)*]
835 //                    [void(%f.frame*)* @f.resume, void(%f.frame*)* @f.destroy]
836 //   define void @f() {
837 //     ...
838 //     call i8* @llvm.coro.begin(i8* null, i32 0, i8* null,
839 //                    i8* bitcast([2 x void(%f.frame*)*] * @f.resumers to i8*))
840 //
841 // Assumes that all the functions have the same signature.
842 static void setCoroInfo(Function &F, coro::Shape &Shape,
843                         ArrayRef<Function *> Fns) {
844   // This only works under the switch-lowering ABI because coro elision
845   // only works on the switch-lowering ABI.
846   assert(Shape.ABI == coro::ABI::Switch);
847 
848   SmallVector<Constant *, 4> Args(Fns.begin(), Fns.end());
849   assert(!Args.empty());
850   Function *Part = *Fns.begin();
851   Module *M = Part->getParent();
852   auto *ArrTy = ArrayType::get(Part->getType(), Args.size());
853 
854   auto *ConstVal = ConstantArray::get(ArrTy, Args);
855   auto *GV = new GlobalVariable(*M, ConstVal->getType(), /*isConstant=*/true,
856                                 GlobalVariable::PrivateLinkage, ConstVal,
857                                 F.getName() + Twine(".resumers"));
858 
859   // Update coro.begin instruction to refer to this constant.
860   LLVMContext &C = F.getContext();
861   auto *BC = ConstantExpr::getPointerCast(GV, Type::getInt8PtrTy(C));
862   Shape.getSwitchCoroId()->setInfo(BC);
863 }
864 
865 // Store addresses of Resume/Destroy/Cleanup functions in the coroutine frame.
866 static void updateCoroFrame(coro::Shape &Shape, Function *ResumeFn,
867                             Function *DestroyFn, Function *CleanupFn) {
868   assert(Shape.ABI == coro::ABI::Switch);
869 
870   IRBuilder<> Builder(Shape.FramePtr->getNextNode());
871   auto *ResumeAddr = Builder.CreateStructGEP(
872       Shape.FrameTy, Shape.FramePtr, coro::Shape::SwitchFieldIndex::Resume,
873       "resume.addr");
874   Builder.CreateStore(ResumeFn, ResumeAddr);
875 
876   Value *DestroyOrCleanupFn = DestroyFn;
877 
878   CoroIdInst *CoroId = Shape.getSwitchCoroId();
879   if (CoroAllocInst *CA = CoroId->getCoroAlloc()) {
880     // If there is a CoroAlloc and it returns false (meaning we elide the
881     // allocation, use CleanupFn instead of DestroyFn).
882     DestroyOrCleanupFn = Builder.CreateSelect(CA, DestroyFn, CleanupFn);
883   }
884 
885   auto *DestroyAddr = Builder.CreateStructGEP(
886       Shape.FrameTy, Shape.FramePtr, coro::Shape::SwitchFieldIndex::Destroy,
887       "destroy.addr");
888   Builder.CreateStore(DestroyOrCleanupFn, DestroyAddr);
889 }
890 
891 static void postSplitCleanup(Function &F) {
892   removeUnreachableBlocks(F);
893 
894   // For now, we do a mandatory verification step because we don't
895   // entirely trust this pass.  Note that we don't want to add a verifier
896   // pass to FPM below because it will also verify all the global data.
897   if (verifyFunction(F, &errs()))
898     report_fatal_error("Broken function");
899 
900   legacy::FunctionPassManager FPM(F.getParent());
901 
902   FPM.add(createSCCPPass());
903   FPM.add(createCFGSimplificationPass());
904   FPM.add(createEarlyCSEPass());
905   FPM.add(createCFGSimplificationPass());
906 
907   FPM.doInitialization();
908   FPM.run(F);
909   FPM.doFinalization();
910 }
911 
912 // Assuming we arrived at the block NewBlock from Prev instruction, store
913 // PHI's incoming values in the ResolvedValues map.
914 static void
915 scanPHIsAndUpdateValueMap(Instruction *Prev, BasicBlock *NewBlock,
916                           DenseMap<Value *, Value *> &ResolvedValues) {
917   auto *PrevBB = Prev->getParent();
918   for (PHINode &PN : NewBlock->phis()) {
919     auto V = PN.getIncomingValueForBlock(PrevBB);
920     // See if we already resolved it.
921     auto VI = ResolvedValues.find(V);
922     if (VI != ResolvedValues.end())
923       V = VI->second;
924     // Remember the value.
925     ResolvedValues[&PN] = V;
926   }
927 }
928 
929 // Replace a sequence of branches leading to a ret, with a clone of a ret
930 // instruction. Suspend instruction represented by a switch, track the PHI
931 // values and select the correct case successor when possible.
932 static bool simplifyTerminatorLeadingToRet(Instruction *InitialInst) {
933   DenseMap<Value *, Value *> ResolvedValues;
934   BasicBlock *UnconditionalSucc = nullptr;
935 
936   Instruction *I = InitialInst;
937   while (I->isTerminator() ||
938          (isa<CmpInst>(I) && I->getNextNode()->isTerminator())) {
939     if (isa<ReturnInst>(I)) {
940       if (I != InitialInst) {
941         // If InitialInst is an unconditional branch,
942         // remove PHI values that come from basic block of InitialInst
943         if (UnconditionalSucc)
944           UnconditionalSucc->removePredecessor(InitialInst->getParent(), true);
945         ReplaceInstWithInst(InitialInst, I->clone());
946       }
947       return true;
948     }
949     if (auto *BR = dyn_cast<BranchInst>(I)) {
950       if (BR->isUnconditional()) {
951         BasicBlock *BB = BR->getSuccessor(0);
952         if (I == InitialInst)
953           UnconditionalSucc = BB;
954         scanPHIsAndUpdateValueMap(I, BB, ResolvedValues);
955         I = BB->getFirstNonPHIOrDbgOrLifetime();
956         continue;
957       }
958     } else if (auto *CondCmp = dyn_cast<CmpInst>(I)) {
959       auto *BR = dyn_cast<BranchInst>(I->getNextNode());
960       if (BR && BR->isConditional() && CondCmp == BR->getCondition()) {
961         // If the case number of suspended switch instruction is reduced to
962         // 1, then it is simplified to CmpInst in llvm::ConstantFoldTerminator.
963         // And the comparsion looks like : %cond = icmp eq i8 %V, constant.
964         ConstantInt *CondConst = dyn_cast<ConstantInt>(CondCmp->getOperand(1));
965         if (CondConst && CondCmp->getPredicate() == CmpInst::ICMP_EQ) {
966           Value *V = CondCmp->getOperand(0);
967           auto it = ResolvedValues.find(V);
968           if (it != ResolvedValues.end())
969             V = it->second;
970 
971           if (ConstantInt *Cond0 = dyn_cast<ConstantInt>(V)) {
972             BasicBlock *BB = Cond0->equalsInt(CondConst->getZExtValue())
973                                  ? BR->getSuccessor(0)
974                                  : BR->getSuccessor(1);
975             scanPHIsAndUpdateValueMap(I, BB, ResolvedValues);
976             I = BB->getFirstNonPHIOrDbgOrLifetime();
977             continue;
978           }
979         }
980       }
981     } else if (auto *SI = dyn_cast<SwitchInst>(I)) {
982       Value *V = SI->getCondition();
983       auto it = ResolvedValues.find(V);
984       if (it != ResolvedValues.end())
985         V = it->second;
986       if (ConstantInt *Cond = dyn_cast<ConstantInt>(V)) {
987         BasicBlock *BB = SI->findCaseValue(Cond)->getCaseSuccessor();
988         scanPHIsAndUpdateValueMap(I, BB, ResolvedValues);
989         I = BB->getFirstNonPHIOrDbgOrLifetime();
990         continue;
991       }
992     }
993     return false;
994   }
995   return false;
996 }
997 
998 // Check whether CI obeys the rules of musttail attribute.
999 static bool shouldBeMustTail(const CallInst &CI, const Function &F) {
1000   if (CI.isInlineAsm())
1001     return false;
1002 
1003   // Match prototypes and calling conventions of resume function.
1004   FunctionType *CalleeTy = CI.getFunctionType();
1005   if (!CalleeTy->getReturnType()->isVoidTy() || (CalleeTy->getNumParams() != 1))
1006     return false;
1007 
1008   Type *CalleeParmTy = CalleeTy->getParamType(0);
1009   if (!CalleeParmTy->isPointerTy() ||
1010       (CalleeParmTy->getPointerAddressSpace() != 0))
1011     return false;
1012 
1013   if (CI.getCallingConv() != F.getCallingConv())
1014     return false;
1015 
1016   // CI should not has any ABI-impacting function attributes.
1017   static const Attribute::AttrKind ABIAttrs[] = {
1018       Attribute::StructRet,    Attribute::ByVal,     Attribute::InAlloca,
1019       Attribute::Preallocated, Attribute::InReg,     Attribute::Returned,
1020       Attribute::SwiftSelf,    Attribute::SwiftError};
1021   AttributeList Attrs = CI.getAttributes();
1022   for (auto AK : ABIAttrs)
1023     if (Attrs.hasParamAttribute(0, AK))
1024       return false;
1025 
1026   return true;
1027 }
1028 
1029 // Add musttail to any resume instructions that is immediately followed by a
1030 // suspend (i.e. ret). We do this even in -O0 to support guaranteed tail call
1031 // for symmetrical coroutine control transfer (C++ Coroutines TS extension).
1032 // This transformation is done only in the resume part of the coroutine that has
1033 // identical signature and calling convention as the coro.resume call.
1034 static void addMustTailToCoroResumes(Function &F) {
1035   bool changed = false;
1036 
1037   // Collect potential resume instructions.
1038   SmallVector<CallInst *, 4> Resumes;
1039   for (auto &I : instructions(F))
1040     if (auto *Call = dyn_cast<CallInst>(&I))
1041       if (shouldBeMustTail(*Call, F))
1042         Resumes.push_back(Call);
1043 
1044   // Set musttail on those that are followed by a ret instruction.
1045   for (CallInst *Call : Resumes)
1046     if (simplifyTerminatorLeadingToRet(Call->getNextNode())) {
1047       Call->setTailCallKind(CallInst::TCK_MustTail);
1048       changed = true;
1049     }
1050 
1051   if (changed)
1052     removeUnreachableBlocks(F);
1053 }
1054 
1055 // Coroutine has no suspend points. Remove heap allocation for the coroutine
1056 // frame if possible.
1057 static void handleNoSuspendCoroutine(coro::Shape &Shape) {
1058   auto *CoroBegin = Shape.CoroBegin;
1059   auto *CoroId = CoroBegin->getId();
1060   auto *AllocInst = CoroId->getCoroAlloc();
1061   switch (Shape.ABI) {
1062   case coro::ABI::Switch: {
1063     auto SwitchId = cast<CoroIdInst>(CoroId);
1064     coro::replaceCoroFree(SwitchId, /*Elide=*/AllocInst != nullptr);
1065     if (AllocInst) {
1066       IRBuilder<> Builder(AllocInst);
1067       auto *Frame = Builder.CreateAlloca(Shape.FrameTy);
1068       Frame->setAlignment(Shape.FrameAlign);
1069       auto *VFrame = Builder.CreateBitCast(Frame, Builder.getInt8PtrTy());
1070       AllocInst->replaceAllUsesWith(Builder.getFalse());
1071       AllocInst->eraseFromParent();
1072       CoroBegin->replaceAllUsesWith(VFrame);
1073     } else {
1074       CoroBegin->replaceAllUsesWith(CoroBegin->getMem());
1075     }
1076     break;
1077   }
1078 
1079   case coro::ABI::Retcon:
1080   case coro::ABI::RetconOnce:
1081     CoroBegin->replaceAllUsesWith(UndefValue::get(CoroBegin->getType()));
1082     break;
1083   }
1084 
1085   CoroBegin->eraseFromParent();
1086 }
1087 
1088 // SimplifySuspendPoint needs to check that there is no calls between
1089 // coro_save and coro_suspend, since any of the calls may potentially resume
1090 // the coroutine and if that is the case we cannot eliminate the suspend point.
1091 static bool hasCallsInBlockBetween(Instruction *From, Instruction *To) {
1092   for (Instruction *I = From; I != To; I = I->getNextNode()) {
1093     // Assume that no intrinsic can resume the coroutine.
1094     if (isa<IntrinsicInst>(I))
1095       continue;
1096 
1097     if (isa<CallBase>(I))
1098       return true;
1099   }
1100   return false;
1101 }
1102 
1103 static bool hasCallsInBlocksBetween(BasicBlock *SaveBB, BasicBlock *ResDesBB) {
1104   SmallPtrSet<BasicBlock *, 8> Set;
1105   SmallVector<BasicBlock *, 8> Worklist;
1106 
1107   Set.insert(SaveBB);
1108   Worklist.push_back(ResDesBB);
1109 
1110   // Accumulate all blocks between SaveBB and ResDesBB. Because CoroSaveIntr
1111   // returns a token consumed by suspend instruction, all blocks in between
1112   // will have to eventually hit SaveBB when going backwards from ResDesBB.
1113   while (!Worklist.empty()) {
1114     auto *BB = Worklist.pop_back_val();
1115     Set.insert(BB);
1116     for (auto *Pred : predecessors(BB))
1117       if (Set.count(Pred) == 0)
1118         Worklist.push_back(Pred);
1119   }
1120 
1121   // SaveBB and ResDesBB are checked separately in hasCallsBetween.
1122   Set.erase(SaveBB);
1123   Set.erase(ResDesBB);
1124 
1125   for (auto *BB : Set)
1126     if (hasCallsInBlockBetween(BB->getFirstNonPHI(), nullptr))
1127       return true;
1128 
1129   return false;
1130 }
1131 
1132 static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
1133   auto *SaveBB = Save->getParent();
1134   auto *ResumeOrDestroyBB = ResumeOrDestroy->getParent();
1135 
1136   if (SaveBB == ResumeOrDestroyBB)
1137     return hasCallsInBlockBetween(Save->getNextNode(), ResumeOrDestroy);
1138 
1139   // Any calls from Save to the end of the block?
1140   if (hasCallsInBlockBetween(Save->getNextNode(), nullptr))
1141     return true;
1142 
1143   // Any calls from begging of the block up to ResumeOrDestroy?
1144   if (hasCallsInBlockBetween(ResumeOrDestroyBB->getFirstNonPHI(),
1145                              ResumeOrDestroy))
1146     return true;
1147 
1148   // Any calls in all of the blocks between SaveBB and ResumeOrDestroyBB?
1149   if (hasCallsInBlocksBetween(SaveBB, ResumeOrDestroyBB))
1150     return true;
1151 
1152   return false;
1153 }
1154 
1155 // If a SuspendIntrin is preceded by Resume or Destroy, we can eliminate the
1156 // suspend point and replace it with nornal control flow.
1157 static bool simplifySuspendPoint(CoroSuspendInst *Suspend,
1158                                  CoroBeginInst *CoroBegin) {
1159   Instruction *Prev = Suspend->getPrevNode();
1160   if (!Prev) {
1161     auto *Pred = Suspend->getParent()->getSinglePredecessor();
1162     if (!Pred)
1163       return false;
1164     Prev = Pred->getTerminator();
1165   }
1166 
1167   CallBase *CB = dyn_cast<CallBase>(Prev);
1168   if (!CB)
1169     return false;
1170 
1171   auto *Callee = CB->getCalledOperand()->stripPointerCasts();
1172 
1173   // See if the callsite is for resumption or destruction of the coroutine.
1174   auto *SubFn = dyn_cast<CoroSubFnInst>(Callee);
1175   if (!SubFn)
1176     return false;
1177 
1178   // Does not refer to the current coroutine, we cannot do anything with it.
1179   if (SubFn->getFrame() != CoroBegin)
1180     return false;
1181 
1182   // See if the transformation is safe. Specifically, see if there are any
1183   // calls in between Save and CallInstr. They can potenitally resume the
1184   // coroutine rendering this optimization unsafe.
1185   auto *Save = Suspend->getCoroSave();
1186   if (hasCallsBetween(Save, CB))
1187     return false;
1188 
1189   // Replace llvm.coro.suspend with the value that results in resumption over
1190   // the resume or cleanup path.
1191   Suspend->replaceAllUsesWith(SubFn->getRawIndex());
1192   Suspend->eraseFromParent();
1193   Save->eraseFromParent();
1194 
1195   // No longer need a call to coro.resume or coro.destroy.
1196   if (auto *Invoke = dyn_cast<InvokeInst>(CB)) {
1197     BranchInst::Create(Invoke->getNormalDest(), Invoke);
1198   }
1199 
1200   // Grab the CalledValue from CB before erasing the CallInstr.
1201   auto *CalledValue = CB->getCalledOperand();
1202   CB->eraseFromParent();
1203 
1204   // If no more users remove it. Usually it is a bitcast of SubFn.
1205   if (CalledValue != SubFn && CalledValue->user_empty())
1206     if (auto *I = dyn_cast<Instruction>(CalledValue))
1207       I->eraseFromParent();
1208 
1209   // Now we are good to remove SubFn.
1210   if (SubFn->user_empty())
1211     SubFn->eraseFromParent();
1212 
1213   return true;
1214 }
1215 
1216 // Remove suspend points that are simplified.
1217 static void simplifySuspendPoints(coro::Shape &Shape) {
1218   // Currently, the only simplification we do is switch-lowering-specific.
1219   if (Shape.ABI != coro::ABI::Switch)
1220     return;
1221 
1222   auto &S = Shape.CoroSuspends;
1223   size_t I = 0, N = S.size();
1224   if (N == 0)
1225     return;
1226   while (true) {
1227     auto SI = cast<CoroSuspendInst>(S[I]);
1228     // Leave final.suspend to handleFinalSuspend since it is undefined behavior
1229     // to resume a coroutine suspended at the final suspend point.
1230     if (!SI->isFinal() && simplifySuspendPoint(SI, Shape.CoroBegin)) {
1231       if (--N == I)
1232         break;
1233       std::swap(S[I], S[N]);
1234       continue;
1235     }
1236     if (++I == N)
1237       break;
1238   }
1239   S.resize(N);
1240 }
1241 
1242 /// For every local variable that has lifetime intrinsics markers, we sink
1243 /// their lifetime.start marker to the places where the variable is being
1244 /// used for the first time. Doing so minimizes the lifetime of each variable,
1245 /// hence minimizing the amount of data we end up putting on the frame.
1246 static void sinkLifetimeStartMarkers(Function &F) {
1247   DominatorTree Dom(F);
1248   for (Instruction &I : instructions(F)) {
1249     // We look for this particular pattern:
1250     //   %tmpX = alloca %.., align ...
1251     //   %0 = bitcast %...* %tmpX to i8*
1252     //   call void @llvm.lifetime.start.p0i8(i64 ..., i8* nonnull %0) #2
1253     if (!isa<AllocaInst>(&I))
1254       continue;
1255     // There can be multiple lifetime start markers for the same variable.
1256     SmallPtrSet<IntrinsicInst *, 1> LifetimeStartInsts;
1257     // SinkBarriers stores all instructions that use this local variable.
1258     // When sinking the lifetime start intrinsics, we can never sink past
1259     // these barriers.
1260     SmallPtrSet<Instruction *, 4> SinkBarriers;
1261     bool Valid = true;
1262     auto AddSinkBarrier = [&](Instruction *I) {
1263       // When adding a new barrier to SinkBarriers, we maintain the case
1264       // that no instruction in SinkBarriers dominates another instruction.
1265       SmallPtrSet<Instruction *, 1> ToRemove;
1266       bool ShouldAdd = true;
1267       for (Instruction *S : SinkBarriers) {
1268         if (I == S || Dom.dominates(S, I)) {
1269           ShouldAdd = false;
1270           break;
1271         } else if (Dom.dominates(I, S)) {
1272           ToRemove.insert(S);
1273         }
1274       }
1275       if (ShouldAdd) {
1276         SinkBarriers.insert(I);
1277         for (Instruction *R : ToRemove) {
1278           SinkBarriers.erase(R);
1279         }
1280       }
1281     };
1282     for (User *U : I.users()) {
1283       if (!isa<BitCastInst>(U))
1284         continue;
1285       for (User *CU : U->users()) {
1286         // If we see any user of CastInst that's not lifetime start/end
1287         // intrinsics, give up because it's too complex.
1288         if (auto *CUI = dyn_cast<IntrinsicInst>(CU)) {
1289           if (CUI->getIntrinsicID() == Intrinsic::lifetime_start)
1290             LifetimeStartInsts.insert(CUI);
1291           else if (CUI->getIntrinsicID() == Intrinsic::lifetime_end)
1292             AddSinkBarrier(CUI);
1293           else
1294             Valid = false;
1295         } else {
1296           Valid = false;
1297         }
1298       }
1299     }
1300     if (!Valid || LifetimeStartInsts.empty())
1301       continue;
1302 
1303     for (User *U : I.users()) {
1304       if (isa<BitCastInst>(U))
1305         continue;
1306       // Every user of the variable is also a sink barrier.
1307       AddSinkBarrier(cast<Instruction>(U));
1308     }
1309 
1310     // For each sink barrier, we insert a lifetime start marker right
1311     // before it.
1312     for (Instruction *S : SinkBarriers) {
1313       if (auto *IS = dyn_cast<IntrinsicInst>(S)) {
1314         if (IS->getIntrinsicID() == Intrinsic::lifetime_end) {
1315           // If we have a lifetime end marker in SinkBarriers, meaning it's
1316           // not dominated by any other users, we can safely delete it.
1317           IS->eraseFromParent();
1318           continue;
1319         }
1320       }
1321       // We find an existing lifetime.start marker that domintes the barrier,
1322       // clone it and insert it right before the barrier. We cannot clone an
1323       // arbitrary lifetime.start marker because we want to make sure the
1324       // BitCast instruction referred in the marker also dominates the barrier.
1325       for (const IntrinsicInst *LifetimeStart : LifetimeStartInsts) {
1326         if (Dom.dominates(LifetimeStart, S)) {
1327           LifetimeStart->clone()->insertBefore(S);
1328           break;
1329         }
1330       }
1331     }
1332     // All the old lifetime.start markers are no longer necessary.
1333     for (IntrinsicInst *S : LifetimeStartInsts) {
1334       S->eraseFromParent();
1335     }
1336   }
1337 }
1338 
1339 static void splitSwitchCoroutine(Function &F, coro::Shape &Shape,
1340                                  SmallVectorImpl<Function *> &Clones) {
1341   assert(Shape.ABI == coro::ABI::Switch);
1342 
1343   createResumeEntryBlock(F, Shape);
1344   auto ResumeClone = createClone(F, ".resume", Shape,
1345                                  CoroCloner::Kind::SwitchResume);
1346   auto DestroyClone = createClone(F, ".destroy", Shape,
1347                                   CoroCloner::Kind::SwitchUnwind);
1348   auto CleanupClone = createClone(F, ".cleanup", Shape,
1349                                   CoroCloner::Kind::SwitchCleanup);
1350 
1351   postSplitCleanup(*ResumeClone);
1352   postSplitCleanup(*DestroyClone);
1353   postSplitCleanup(*CleanupClone);
1354 
1355   addMustTailToCoroResumes(*ResumeClone);
1356 
1357   // Store addresses resume/destroy/cleanup functions in the coroutine frame.
1358   updateCoroFrame(Shape, ResumeClone, DestroyClone, CleanupClone);
1359 
1360   assert(Clones.empty());
1361   Clones.push_back(ResumeClone);
1362   Clones.push_back(DestroyClone);
1363   Clones.push_back(CleanupClone);
1364 
1365   // Create a constant array referring to resume/destroy/clone functions pointed
1366   // by the last argument of @llvm.coro.info, so that CoroElide pass can
1367   // determined correct function to call.
1368   setCoroInfo(F, Shape, Clones);
1369 }
1370 
1371 static void splitRetconCoroutine(Function &F, coro::Shape &Shape,
1372                                  SmallVectorImpl<Function *> &Clones) {
1373   assert(Shape.ABI == coro::ABI::Retcon ||
1374          Shape.ABI == coro::ABI::RetconOnce);
1375   assert(Clones.empty());
1376 
1377   // Reset various things that the optimizer might have decided it
1378   // "knows" about the coroutine function due to not seeing a return.
1379   F.removeFnAttr(Attribute::NoReturn);
1380   F.removeAttribute(AttributeList::ReturnIndex, Attribute::NoAlias);
1381   F.removeAttribute(AttributeList::ReturnIndex, Attribute::NonNull);
1382 
1383   // Allocate the frame.
1384   auto *Id = cast<AnyCoroIdRetconInst>(Shape.CoroBegin->getId());
1385   Value *RawFramePtr;
1386   if (Shape.RetconLowering.IsFrameInlineInStorage) {
1387     RawFramePtr = Id->getStorage();
1388   } else {
1389     IRBuilder<> Builder(Id);
1390 
1391     // Determine the size of the frame.
1392     const DataLayout &DL = F.getParent()->getDataLayout();
1393     auto Size = DL.getTypeAllocSize(Shape.FrameTy);
1394 
1395     // Allocate.  We don't need to update the call graph node because we're
1396     // going to recompute it from scratch after splitting.
1397     // FIXME: pass the required alignment
1398     RawFramePtr = Shape.emitAlloc(Builder, Builder.getInt64(Size), nullptr);
1399     RawFramePtr =
1400       Builder.CreateBitCast(RawFramePtr, Shape.CoroBegin->getType());
1401 
1402     // Stash the allocated frame pointer in the continuation storage.
1403     auto Dest = Builder.CreateBitCast(Id->getStorage(),
1404                                       RawFramePtr->getType()->getPointerTo());
1405     Builder.CreateStore(RawFramePtr, Dest);
1406   }
1407 
1408   // Map all uses of llvm.coro.begin to the allocated frame pointer.
1409   {
1410     // Make sure we don't invalidate Shape.FramePtr.
1411     TrackingVH<Instruction> Handle(Shape.FramePtr);
1412     Shape.CoroBegin->replaceAllUsesWith(RawFramePtr);
1413     Shape.FramePtr = Handle.getValPtr();
1414   }
1415 
1416   // Create a unique return block.
1417   BasicBlock *ReturnBB = nullptr;
1418   SmallVector<PHINode *, 4> ReturnPHIs;
1419 
1420   // Create all the functions in order after the main function.
1421   auto NextF = std::next(F.getIterator());
1422 
1423   // Create a continuation function for each of the suspend points.
1424   Clones.reserve(Shape.CoroSuspends.size());
1425   for (size_t i = 0, e = Shape.CoroSuspends.size(); i != e; ++i) {
1426     auto Suspend = cast<CoroSuspendRetconInst>(Shape.CoroSuspends[i]);
1427 
1428     // Create the clone declaration.
1429     auto Continuation =
1430       createCloneDeclaration(F, Shape, ".resume." + Twine(i), NextF);
1431     Clones.push_back(Continuation);
1432 
1433     // Insert a branch to the unified return block immediately before
1434     // the suspend point.
1435     auto SuspendBB = Suspend->getParent();
1436     auto NewSuspendBB = SuspendBB->splitBasicBlock(Suspend);
1437     auto Branch = cast<BranchInst>(SuspendBB->getTerminator());
1438 
1439     // Create the unified return block.
1440     if (!ReturnBB) {
1441       // Place it before the first suspend.
1442       ReturnBB = BasicBlock::Create(F.getContext(), "coro.return", &F,
1443                                     NewSuspendBB);
1444       Shape.RetconLowering.ReturnBlock = ReturnBB;
1445 
1446       IRBuilder<> Builder(ReturnBB);
1447 
1448       // Create PHIs for all the return values.
1449       assert(ReturnPHIs.empty());
1450 
1451       // First, the continuation.
1452       ReturnPHIs.push_back(Builder.CreatePHI(Continuation->getType(),
1453                                              Shape.CoroSuspends.size()));
1454 
1455       // Next, all the directly-yielded values.
1456       for (auto ResultTy : Shape.getRetconResultTypes())
1457         ReturnPHIs.push_back(Builder.CreatePHI(ResultTy,
1458                                                Shape.CoroSuspends.size()));
1459 
1460       // Build the return value.
1461       auto RetTy = F.getReturnType();
1462 
1463       // Cast the continuation value if necessary.
1464       // We can't rely on the types matching up because that type would
1465       // have to be infinite.
1466       auto CastedContinuationTy =
1467         (ReturnPHIs.size() == 1 ? RetTy : RetTy->getStructElementType(0));
1468       auto *CastedContinuation =
1469         Builder.CreateBitCast(ReturnPHIs[0], CastedContinuationTy);
1470 
1471       Value *RetV;
1472       if (ReturnPHIs.size() == 1) {
1473         RetV = CastedContinuation;
1474       } else {
1475         RetV = UndefValue::get(RetTy);
1476         RetV = Builder.CreateInsertValue(RetV, CastedContinuation, 0);
1477         for (size_t I = 1, E = ReturnPHIs.size(); I != E; ++I)
1478           RetV = Builder.CreateInsertValue(RetV, ReturnPHIs[I], I);
1479       }
1480 
1481       Builder.CreateRet(RetV);
1482     }
1483 
1484     // Branch to the return block.
1485     Branch->setSuccessor(0, ReturnBB);
1486     ReturnPHIs[0]->addIncoming(Continuation, SuspendBB);
1487     size_t NextPHIIndex = 1;
1488     for (auto &VUse : Suspend->value_operands())
1489       ReturnPHIs[NextPHIIndex++]->addIncoming(&*VUse, SuspendBB);
1490     assert(NextPHIIndex == ReturnPHIs.size());
1491   }
1492 
1493   assert(Clones.size() == Shape.CoroSuspends.size());
1494   for (size_t i = 0, e = Shape.CoroSuspends.size(); i != e; ++i) {
1495     auto Suspend = Shape.CoroSuspends[i];
1496     auto Clone = Clones[i];
1497 
1498     CoroCloner(F, "resume." + Twine(i), Shape, Clone, Suspend).create();
1499   }
1500 }
1501 
1502 namespace {
1503   class PrettyStackTraceFunction : public PrettyStackTraceEntry {
1504     Function &F;
1505   public:
1506     PrettyStackTraceFunction(Function &F) : F(F) {}
1507     void print(raw_ostream &OS) const override {
1508       OS << "While splitting coroutine ";
1509       F.printAsOperand(OS, /*print type*/ false, F.getParent());
1510       OS << "\n";
1511     }
1512   };
1513 }
1514 
1515 static coro::Shape splitCoroutine(Function &F,
1516                                   SmallVectorImpl<Function *> &Clones) {
1517   PrettyStackTraceFunction prettyStackTrace(F);
1518 
1519   // The suspend-crossing algorithm in buildCoroutineFrame get tripped
1520   // up by uses in unreachable blocks, so remove them as a first pass.
1521   removeUnreachableBlocks(F);
1522 
1523   coro::Shape Shape(F);
1524   if (!Shape.CoroBegin)
1525     return Shape;
1526 
1527   simplifySuspendPoints(Shape);
1528   sinkLifetimeStartMarkers(F);
1529   buildCoroutineFrame(F, Shape);
1530   replaceFrameSize(Shape);
1531 
1532   // If there are no suspend points, no split required, just remove
1533   // the allocation and deallocation blocks, they are not needed.
1534   if (Shape.CoroSuspends.empty()) {
1535     handleNoSuspendCoroutine(Shape);
1536   } else {
1537     switch (Shape.ABI) {
1538     case coro::ABI::Switch:
1539       splitSwitchCoroutine(F, Shape, Clones);
1540       break;
1541     case coro::ABI::Retcon:
1542     case coro::ABI::RetconOnce:
1543       splitRetconCoroutine(F, Shape, Clones);
1544       break;
1545     }
1546   }
1547 
1548   // Replace all the swifterror operations in the original function.
1549   // This invalidates SwiftErrorOps in the Shape.
1550   replaceSwiftErrorOps(F, Shape, nullptr);
1551 
1552   return Shape;
1553 }
1554 
1555 static void
1556 updateCallGraphAfterCoroutineSplit(Function &F, const coro::Shape &Shape,
1557                                    const SmallVectorImpl<Function *> &Clones,
1558                                    CallGraph &CG, CallGraphSCC &SCC) {
1559   if (!Shape.CoroBegin)
1560     return;
1561 
1562   removeCoroEnds(Shape, &CG);
1563   postSplitCleanup(F);
1564 
1565   // Update call graph and add the functions we created to the SCC.
1566   coro::updateCallGraph(F, Clones, CG, SCC);
1567 }
1568 
1569 static void updateCallGraphAfterCoroutineSplit(
1570     LazyCallGraph::Node &N, const coro::Shape &Shape,
1571     const SmallVectorImpl<Function *> &Clones, LazyCallGraph::SCC &C,
1572     LazyCallGraph &CG, CGSCCAnalysisManager &AM, CGSCCUpdateResult &UR,
1573     FunctionAnalysisManager &FAM) {
1574   if (!Shape.CoroBegin)
1575     return;
1576 
1577   for (llvm::CoroEndInst *End : Shape.CoroEnds) {
1578     auto &Context = End->getContext();
1579     End->replaceAllUsesWith(ConstantInt::getFalse(Context));
1580     End->eraseFromParent();
1581   }
1582 
1583   postSplitCleanup(N.getFunction());
1584 
1585   // To insert the newly created coroutine funclets 'f.resume', 'f.destroy', and
1586   // 'f.cleanup' into the same SCC as the coroutine 'f' they were outlined from,
1587   // we make use of the CallGraphUpdater class, which can modify the internal
1588   // state of the LazyCallGraph.
1589   for (Function *Clone : Clones)
1590     CG.addNewFunctionIntoRefSCC(*Clone, C.getOuterRefSCC());
1591 
1592   // We've inserted instructions into coroutine 'f' that reference the three new
1593   // coroutine funclets. We must now update the call graph so that reference
1594   // edges between 'f' and its funclets are added to it. LazyCallGraph only
1595   // allows CGSCC passes to insert "trivial" reference edges. We've ensured
1596   // above, by inserting the funclets into the same SCC as the corutine, that
1597   // the edges are trivial.
1598   //
1599   // N.B.: If we didn't update the call graph here, a CGSCCToFunctionPassAdaptor
1600   // later in this CGSCC pass pipeline may be run, triggering a call graph
1601   // update of its own. Function passes run by the adaptor are not permitted to
1602   // add new edges of any kind to the graph, and the new edges inserted by this
1603   // pass would be misattributed to that unrelated function pass.
1604   updateCGAndAnalysisManagerForCGSCCPass(CG, C, N, AM, UR, FAM);
1605 }
1606 
1607 // When we see the coroutine the first time, we insert an indirect call to a
1608 // devirt trigger function and mark the coroutine that it is now ready for
1609 // split.
1610 static void prepareForSplit(Function &F, CallGraph &CG) {
1611   Module &M = *F.getParent();
1612   LLVMContext &Context = F.getContext();
1613 #ifndef NDEBUG
1614   Function *DevirtFn = M.getFunction(CORO_DEVIRT_TRIGGER_FN);
1615   assert(DevirtFn && "coro.devirt.trigger function not found");
1616 #endif
1617 
1618   F.addFnAttr(CORO_PRESPLIT_ATTR, PREPARED_FOR_SPLIT);
1619 
1620   // Insert an indirect call sequence that will be devirtualized by CoroElide
1621   // pass:
1622   //    %0 = call i8* @llvm.coro.subfn.addr(i8* null, i8 -1)
1623   //    %1 = bitcast i8* %0 to void(i8*)*
1624   //    call void %1(i8* null)
1625   coro::LowererBase Lowerer(M);
1626   Instruction *InsertPt = F.getEntryBlock().getTerminator();
1627   auto *Null = ConstantPointerNull::get(Type::getInt8PtrTy(Context));
1628   auto *DevirtFnAddr =
1629       Lowerer.makeSubFnCall(Null, CoroSubFnInst::RestartTrigger, InsertPt);
1630   FunctionType *FnTy = FunctionType::get(Type::getVoidTy(Context),
1631                                          {Type::getInt8PtrTy(Context)}, false);
1632   auto *IndirectCall = CallInst::Create(FnTy, DevirtFnAddr, Null, "", InsertPt);
1633 
1634   // Update CG graph with an indirect call we just added.
1635   CG[&F]->addCalledFunction(IndirectCall, CG.getCallsExternalNode());
1636 }
1637 
1638 // Make sure that there is a devirtualization trigger function that the
1639 // coro-split pass uses to force a restart of the CGSCC pipeline. If the devirt
1640 // trigger function is not found, we will create one and add it to the current
1641 // SCC.
1642 static void createDevirtTriggerFunc(CallGraph &CG, CallGraphSCC &SCC) {
1643   Module &M = CG.getModule();
1644   if (M.getFunction(CORO_DEVIRT_TRIGGER_FN))
1645     return;
1646 
1647   LLVMContext &C = M.getContext();
1648   auto *FnTy = FunctionType::get(Type::getVoidTy(C), Type::getInt8PtrTy(C),
1649                                  /*isVarArg=*/false);
1650   Function *DevirtFn =
1651       Function::Create(FnTy, GlobalValue::LinkageTypes::PrivateLinkage,
1652                        CORO_DEVIRT_TRIGGER_FN, &M);
1653   DevirtFn->addFnAttr(Attribute::AlwaysInline);
1654   auto *Entry = BasicBlock::Create(C, "entry", DevirtFn);
1655   ReturnInst::Create(C, Entry);
1656 
1657   auto *Node = CG.getOrInsertFunction(DevirtFn);
1658 
1659   SmallVector<CallGraphNode *, 8> Nodes(SCC.begin(), SCC.end());
1660   Nodes.push_back(Node);
1661   SCC.initialize(Nodes);
1662 }
1663 
1664 /// Replace a call to llvm.coro.prepare.retcon.
1665 static void replacePrepare(CallInst *Prepare, CallGraph &CG) {
1666   auto CastFn = Prepare->getArgOperand(0); // as an i8*
1667   auto Fn = CastFn->stripPointerCasts(); // as its original type
1668 
1669   // Find call graph nodes for the preparation.
1670   CallGraphNode *PrepareUserNode = nullptr, *FnNode = nullptr;
1671   if (auto ConcreteFn = dyn_cast<Function>(Fn)) {
1672     PrepareUserNode = CG[Prepare->getFunction()];
1673     FnNode = CG[ConcreteFn];
1674   }
1675 
1676   // Attempt to peephole this pattern:
1677   //    %0 = bitcast [[TYPE]] @some_function to i8*
1678   //    %1 = call @llvm.coro.prepare.retcon(i8* %0)
1679   //    %2 = bitcast %1 to [[TYPE]]
1680   // ==>
1681   //    %2 = @some_function
1682   for (auto UI = Prepare->use_begin(), UE = Prepare->use_end();
1683          UI != UE; ) {
1684     // Look for bitcasts back to the original function type.
1685     auto *Cast = dyn_cast<BitCastInst>((UI++)->getUser());
1686     if (!Cast || Cast->getType() != Fn->getType()) continue;
1687 
1688     // Check whether the replacement will introduce new direct calls.
1689     // If so, we'll need to update the call graph.
1690     if (PrepareUserNode) {
1691       for (auto &Use : Cast->uses()) {
1692         if (auto *CB = dyn_cast<CallBase>(Use.getUser())) {
1693           if (!CB->isCallee(&Use))
1694             continue;
1695           PrepareUserNode->removeCallEdgeFor(*CB);
1696           PrepareUserNode->addCalledFunction(CB, FnNode);
1697         }
1698       }
1699     }
1700 
1701     // Replace and remove the cast.
1702     Cast->replaceAllUsesWith(Fn);
1703     Cast->eraseFromParent();
1704   }
1705 
1706   // Replace any remaining uses with the function as an i8*.
1707   // This can never directly be a callee, so we don't need to update CG.
1708   Prepare->replaceAllUsesWith(CastFn);
1709   Prepare->eraseFromParent();
1710 
1711   // Kill dead bitcasts.
1712   while (auto *Cast = dyn_cast<BitCastInst>(CastFn)) {
1713     if (!Cast->use_empty()) break;
1714     CastFn = Cast->getOperand(0);
1715     Cast->eraseFromParent();
1716   }
1717 }
1718 
1719 /// Remove calls to llvm.coro.prepare.retcon, a barrier meant to prevent
1720 /// IPO from operating on calls to a retcon coroutine before it's been
1721 /// split.  This is only safe to do after we've split all retcon
1722 /// coroutines in the module.  We can do that this in this pass because
1723 /// this pass does promise to split all retcon coroutines (as opposed to
1724 /// switch coroutines, which are lowered in multiple stages).
1725 static bool replaceAllPrepares(Function *PrepareFn, CallGraph &CG) {
1726   bool Changed = false;
1727   for (auto PI = PrepareFn->use_begin(), PE = PrepareFn->use_end();
1728          PI != PE; ) {
1729     // Intrinsics can only be used in calls.
1730     auto *Prepare = cast<CallInst>((PI++)->getUser());
1731     replacePrepare(Prepare, CG);
1732     Changed = true;
1733   }
1734 
1735   return Changed;
1736 }
1737 
1738 static bool declaresCoroSplitIntrinsics(const Module &M) {
1739   return coro::declaresIntrinsics(
1740       M, {"llvm.coro.begin", "llvm.coro.prepare.retcon"});
1741 }
1742 
1743 PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
1744                                      CGSCCAnalysisManager &AM,
1745                                      LazyCallGraph &CG, CGSCCUpdateResult &UR) {
1746   // NB: One invariant of a valid LazyCallGraph::SCC is that it must contain a
1747   //     non-zero number of nodes, so we assume that here and grab the first
1748   //     node's function's module.
1749   Module &M = *C.begin()->getFunction().getParent();
1750   auto &FAM =
1751       AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
1752 
1753   if (!declaresCoroSplitIntrinsics(M))
1754     return PreservedAnalyses::all();
1755 
1756   // Check for uses of llvm.coro.prepare.retcon.
1757   const auto *PrepareFn = M.getFunction("llvm.coro.prepare.retcon");
1758   if (PrepareFn && PrepareFn->use_empty())
1759     PrepareFn = nullptr;
1760 
1761   // Find coroutines for processing.
1762   SmallVector<LazyCallGraph::Node *, 4> Coroutines;
1763   for (LazyCallGraph::Node &N : C)
1764     if (N.getFunction().hasFnAttribute(CORO_PRESPLIT_ATTR))
1765       Coroutines.push_back(&N);
1766 
1767   if (Coroutines.empty() && !PrepareFn)
1768     return PreservedAnalyses::all();
1769 
1770   if (Coroutines.empty())
1771     llvm_unreachable("new pass manager cannot yet handle "
1772                      "'llvm.coro.prepare.retcon'");
1773 
1774   // Split all the coroutines.
1775   for (LazyCallGraph::Node *N : Coroutines) {
1776     Function &F = N->getFunction();
1777     Attribute Attr = F.getFnAttribute(CORO_PRESPLIT_ATTR);
1778     StringRef Value = Attr.getValueAsString();
1779     LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName()
1780                       << "' state: " << Value << "\n");
1781     if (Value == UNPREPARED_FOR_SPLIT) {
1782       // Enqueue a second iteration of the CGSCC pipeline.
1783       // N.B.:
1784       // The CoroSplitLegacy pass "triggers" a restart of the CGSCC pass
1785       // pipeline by inserting an indirect function call that the
1786       // CoroElideLegacy pass then replaces with a direct function call. The
1787       // legacy CGSCC pipeline's implicit behavior was as if wrapped in the new
1788       // pass manager abstraction DevirtSCCRepeatedPass.
1789       //
1790       // This pass does not need to "trigger" another run of the pipeline.
1791       // Instead, it simply enqueues the same RefSCC onto the pipeline's
1792       // worklist.
1793       UR.CWorklist.insert(&C);
1794       F.addFnAttr(CORO_PRESPLIT_ATTR, PREPARED_FOR_SPLIT);
1795       continue;
1796     }
1797     F.removeFnAttr(CORO_PRESPLIT_ATTR);
1798 
1799     SmallVector<Function *, 4> Clones;
1800     const coro::Shape Shape = splitCoroutine(F, Clones);
1801     updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, C, CG, AM, UR, FAM);
1802   }
1803 
1804   if (PrepareFn)
1805     llvm_unreachable("new pass manager cannot yet handle "
1806                      "'llvm.coro.prepare.retcon'");
1807 
1808   return PreservedAnalyses::none();
1809 }
1810 
1811 namespace {
1812 
1813 // We present a coroutine to LLVM as an ordinary function with suspension
1814 // points marked up with intrinsics. We let the optimizer party on the coroutine
1815 // as a single function for as long as possible. Shortly before the coroutine is
1816 // eligible to be inlined into its callers, we split up the coroutine into parts
1817 // corresponding to initial, resume and destroy invocations of the coroutine,
1818 // add them to the current SCC and restart the IPO pipeline to optimize the
1819 // coroutine subfunctions we extracted before proceeding to the caller of the
1820 // coroutine.
1821 struct CoroSplitLegacy : public CallGraphSCCPass {
1822   static char ID; // Pass identification, replacement for typeid
1823 
1824   CoroSplitLegacy() : CallGraphSCCPass(ID) {
1825     initializeCoroSplitLegacyPass(*PassRegistry::getPassRegistry());
1826   }
1827 
1828   bool Run = false;
1829 
1830   // A coroutine is identified by the presence of coro.begin intrinsic, if
1831   // we don't have any, this pass has nothing to do.
1832   bool doInitialization(CallGraph &CG) override {
1833     Run = declaresCoroSplitIntrinsics(CG.getModule());
1834     return CallGraphSCCPass::doInitialization(CG);
1835   }
1836 
1837   bool runOnSCC(CallGraphSCC &SCC) override {
1838     if (!Run)
1839       return false;
1840 
1841     // Check for uses of llvm.coro.prepare.retcon.
1842     auto PrepareFn =
1843       SCC.getCallGraph().getModule().getFunction("llvm.coro.prepare.retcon");
1844     if (PrepareFn && PrepareFn->use_empty())
1845       PrepareFn = nullptr;
1846 
1847     // Find coroutines for processing.
1848     SmallVector<Function *, 4> Coroutines;
1849     for (CallGraphNode *CGN : SCC)
1850       if (auto *F = CGN->getFunction())
1851         if (F->hasFnAttribute(CORO_PRESPLIT_ATTR))
1852           Coroutines.push_back(F);
1853 
1854     if (Coroutines.empty() && !PrepareFn)
1855       return false;
1856 
1857     CallGraph &CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();
1858 
1859     if (Coroutines.empty())
1860       return replaceAllPrepares(PrepareFn, CG);
1861 
1862     createDevirtTriggerFunc(CG, SCC);
1863 
1864     // Split all the coroutines.
1865     for (Function *F : Coroutines) {
1866       Attribute Attr = F->getFnAttribute(CORO_PRESPLIT_ATTR);
1867       StringRef Value = Attr.getValueAsString();
1868       LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F->getName()
1869                         << "' state: " << Value << "\n");
1870       if (Value == UNPREPARED_FOR_SPLIT) {
1871         prepareForSplit(*F, CG);
1872         continue;
1873       }
1874       F->removeFnAttr(CORO_PRESPLIT_ATTR);
1875 
1876       SmallVector<Function *, 4> Clones;
1877       const coro::Shape Shape = splitCoroutine(*F, Clones);
1878       updateCallGraphAfterCoroutineSplit(*F, Shape, Clones, CG, SCC);
1879     }
1880 
1881     if (PrepareFn)
1882       replaceAllPrepares(PrepareFn, CG);
1883 
1884     return true;
1885   }
1886 
1887   void getAnalysisUsage(AnalysisUsage &AU) const override {
1888     CallGraphSCCPass::getAnalysisUsage(AU);
1889   }
1890 
1891   StringRef getPassName() const override { return "Coroutine Splitting"; }
1892 };
1893 
1894 } // end anonymous namespace
1895 
1896 char CoroSplitLegacy::ID = 0;
1897 
1898 INITIALIZE_PASS_BEGIN(
1899     CoroSplitLegacy, "coro-split",
1900     "Split coroutine into a set of functions driving its state machine", false,
1901     false)
1902 INITIALIZE_PASS_DEPENDENCY(CallGraphWrapperPass)
1903 INITIALIZE_PASS_END(
1904     CoroSplitLegacy, "coro-split",
1905     "Split coroutine into a set of functions driving its state machine", false,
1906     false)
1907 
1908 Pass *llvm::createCoroSplitLegacyPass() { return new CoroSplitLegacy(); }
1909