xref: /llvm-project/llvm/lib/Transforms/Coroutines/CoroSplit.cpp (revision 285009f202ca8bfcc6b607eba0e919867559e725)

//===- CoroSplit.cpp - Converts a coroutine into a state machine ----------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
// This pass builds the coroutine frame and outlines resume and destroy parts
// of the coroutine into separate functions.
//
// We present a coroutine to an LLVM as an ordinary function with suspension
// points marked up with intrinsics. We let the optimizer party on the coroutine
// as a single function for as long as possible. Shortly before the coroutine is
// eligible to be inlined into its callers, we split up the coroutine into parts
// corresponding to an initial, resume and destroy invocations of the coroutine,
// add them to the current SCC and restart the IPO pipeline to optimize the
// coroutine subfunctions we extracted before proceeding to the caller of the
// coroutine.
//===----------------------------------------------------------------------===//

#include "llvm/Transforms/Coroutines/CoroSplit.h"
#include "CoroCloner.h"
#include "CoroInternal.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/PriorityWorklist.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm/Analysis/CFG.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Analysis/ConstantFolding.h"
#include "llvm/Analysis/LazyCallGraph.h"
#include "llvm/Analysis/OptimizationRemarkEmitter.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/BinaryFormat/Dwarf.h"
#include "llvm/IR/Argument.h"
#include "llvm/IR/Attributes.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/CFG.h"
#include "llvm/IR/CallingConv.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/DebugInfo.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/GlobalValue.h"
#include "llvm/IR/GlobalVariable.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/Value.h"
#include "llvm/IR/Verifier.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/PrettyStackTrace.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Coroutines/MaterializationUtils.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/CallGraphUpdater.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/Local.h"
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <initializer_list>
#include <iterator>

using namespace llvm;

#define DEBUG_TYPE "coro-split"

namespace {
/// Collect (a known) subset of global debug info metadata potentially used by
/// the function \p F.
///
/// This metadata set can be used to avoid cloning debug info not owned by \p F
/// and is shared among all potential clones \p F.
MetadataSetTy collectCommonDebugInfo(Function &F) {
  TimeTraceScope FunctionScope("CollectCommonDebugInfo");

  DebugInfoFinder DIFinder;
  DISubprogram *SPClonedWithinModule = CollectDebugInfoForCloning(
      F, CloneFunctionChangeType::LocalChangesOnly, DIFinder);

  return FindDebugInfoToIdentityMap(CloneFunctionChangeType::LocalChangesOnly,
                                    DIFinder, SPClonedWithinModule);
}
} // end anonymous namespace

// FIXME:
// Lower the intrinisc in CoroEarly phase if coroutine frame doesn't escape
// and it is known that other transformations, for example, sanitizers
// won't lead to incorrect code.
static void lowerAwaitSuspend(IRBuilder<> &Builder, CoroAwaitSuspendInst *CB,
                              coro::Shape &Shape) {
  auto Wrapper = CB->getWrapperFunction();
  auto Awaiter = CB->getAwaiter();
  auto FramePtr = CB->getFrame();

  Builder.SetInsertPoint(CB);

  CallBase *NewCall = nullptr;
  // await_suspend has only 2 parameters, awaiter and handle.
  // Copy parameter attributes from the intrinsic call, but remove the last,
  // because the last parameter now becomes the function that is being called.
  AttributeList NewAttributes =
      CB->getAttributes().removeParamAttributes(CB->getContext(), 2);

  if (auto Invoke = dyn_cast<InvokeInst>(CB)) {
    auto WrapperInvoke =
        Builder.CreateInvoke(Wrapper, Invoke->getNormalDest(),
                             Invoke->getUnwindDest(), {Awaiter, FramePtr});

    WrapperInvoke->setCallingConv(Invoke->getCallingConv());
    std::copy(Invoke->bundle_op_info_begin(), Invoke->bundle_op_info_end(),
              WrapperInvoke->bundle_op_info_begin());
    WrapperInvoke->setAttributes(NewAttributes);
    WrapperInvoke->setDebugLoc(Invoke->getDebugLoc());
    NewCall = WrapperInvoke;
  } else if (auto Call = dyn_cast<CallInst>(CB)) {
    auto WrapperCall = Builder.CreateCall(Wrapper, {Awaiter, FramePtr});

    WrapperCall->setAttributes(NewAttributes);
    WrapperCall->setDebugLoc(Call->getDebugLoc());
    NewCall = WrapperCall;
  } else {
    llvm_unreachable("Unexpected coro_await_suspend invocation method");
  }

  if (CB->getCalledFunction()->getIntrinsicID() ==
      Intrinsic::coro_await_suspend_handle) {
    // Follow the lowered await_suspend call above with a lowered resume call
    // to the returned coroutine.
    if (auto *Invoke = dyn_cast<InvokeInst>(CB)) {
      // If the await_suspend call is an invoke, we continue in the next block.
      Builder.SetInsertPoint(Invoke->getNormalDest()->getFirstInsertionPt());
    }

    coro::LowererBase LB(*Wrapper->getParent());
    auto *ResumeAddr = LB.makeSubFnCall(NewCall, CoroSubFnInst::ResumeIndex,
                                        &*Builder.GetInsertPoint());

    LLVMContext &Ctx = Builder.getContext();
    FunctionType *ResumeTy = FunctionType::get(
        Type::getVoidTy(Ctx), PointerType::getUnqual(Ctx), false);
    auto *ResumeCall = Builder.CreateCall(ResumeTy, ResumeAddr, {NewCall});
    ResumeCall->setCallingConv(CallingConv::Fast);

    // We can't insert the 'ret' instruction and adjust the cc until the
    // function has been split, so remember this for later.
    Shape.SymmetricTransfers.push_back(ResumeCall);

    NewCall = ResumeCall;
  }

  CB->replaceAllUsesWith(NewCall);
  CB->eraseFromParent();
}

static void lowerAwaitSuspends(Function &F, coro::Shape &Shape) {
  IRBuilder<> Builder(F.getContext());
  for (auto *AWS : Shape.CoroAwaitSuspends)
    lowerAwaitSuspend(Builder, AWS, Shape);
}

static void maybeFreeRetconStorage(IRBuilder<> &Builder,
                                   const coro::Shape &Shape, Value *FramePtr,
                                   CallGraph *CG) {
  assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce);
  if (Shape.RetconLowering.IsFrameInlineInStorage)
    return;

  Shape.emitDealloc(Builder, FramePtr, CG);
}

/// Replace an llvm.coro.end.async.
/// Will inline the must tail call function call if there is one.
/// \returns true if cleanup of the coro.end block is needed, false otherwise.
static bool replaceCoroEndAsync(AnyCoroEndInst *End) {
  IRBuilder<> Builder(End);

  auto *EndAsync = dyn_cast<CoroAsyncEndInst>(End);
  if (!EndAsync) {
    Builder.CreateRetVoid();
    return true /*needs cleanup of coro.end block*/;
  }

  auto *MustTailCallFunc = EndAsync->getMustTailCallFunction();
  if (!MustTailCallFunc) {
    Builder.CreateRetVoid();
    return true /*needs cleanup of coro.end block*/;
  }

  // Move the must tail call from the predecessor block into the end block.
  auto *CoroEndBlock = End->getParent();
  auto *MustTailCallFuncBlock = CoroEndBlock->getSinglePredecessor();
  assert(MustTailCallFuncBlock && "Must have a single predecessor block");
  auto It = MustTailCallFuncBlock->getTerminator()->getIterator();
  auto *MustTailCall = cast<CallInst>(&*std::prev(It));
  CoroEndBlock->splice(End->getIterator(), MustTailCallFuncBlock,
                       MustTailCall->getIterator());

  // Insert the return instruction.
  Builder.SetInsertPoint(End);
  Builder.CreateRetVoid();
  InlineFunctionInfo FnInfo;

  // Remove the rest of the block, by splitting it into an unreachable block.
  auto *BB = End->getParent();
  BB->splitBasicBlock(End);
  BB->getTerminator()->eraseFromParent();

  auto InlineRes = InlineFunction(*MustTailCall, FnInfo);
  assert(InlineRes.isSuccess() && "Expected inlining to succeed");
  (void)InlineRes;

  // We have cleaned up the coro.end block above.
  return false;
}

/// Replace a non-unwind call to llvm.coro.end.
static void replaceFallthroughCoroEnd(AnyCoroEndInst *End,
                                      const coro::Shape &Shape, Value *FramePtr,
                                      bool InResume, CallGraph *CG) {
  // Start inserting right before the coro.end.
  IRBuilder<> Builder(End);

  // Create the return instruction.
  switch (Shape.ABI) {
  // The cloned functions in switch-lowering always return void.
  case coro::ABI::Switch:
    assert(!cast<CoroEndInst>(End)->hasResults() &&
           "switch coroutine should not return any values");
    // coro.end doesn't immediately end the coroutine in the main function
    // in this lowering, because we need to deallocate the coroutine.
    if (!InResume)
      return;
    Builder.CreateRetVoid();
    break;

  // In async lowering this returns.
  case coro::ABI::Async: {
    bool CoroEndBlockNeedsCleanup = replaceCoroEndAsync(End);
    if (!CoroEndBlockNeedsCleanup)
      return;
    break;
  }

  // In unique continuation lowering, the continuations always return void.
  // But we may have implicitly allocated storage.
  case coro::ABI::RetconOnce: {
    maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
    auto *CoroEnd = cast<CoroEndInst>(End);
    auto *RetTy = Shape.getResumeFunctionType()->getReturnType();

    if (!CoroEnd->hasResults()) {
      assert(RetTy->isVoidTy());
      Builder.CreateRetVoid();
      break;
    }

    auto *CoroResults = CoroEnd->getResults();
    unsigned NumReturns = CoroResults->numReturns();

    if (auto *RetStructTy = dyn_cast<StructType>(RetTy)) {
      assert(RetStructTy->getNumElements() == NumReturns &&
             "numbers of returns should match resume function singature");
      Value *ReturnValue = PoisonValue::get(RetStructTy);
      unsigned Idx = 0;
      for (Value *RetValEl : CoroResults->return_values())
        ReturnValue = Builder.CreateInsertValue(ReturnValue, RetValEl, Idx++);
      Builder.CreateRet(ReturnValue);
    } else if (NumReturns == 0) {
      assert(RetTy->isVoidTy());
      Builder.CreateRetVoid();
    } else {
      assert(NumReturns == 1);
      Builder.CreateRet(*CoroResults->retval_begin());
    }
    CoroResults->replaceAllUsesWith(
        ConstantTokenNone::get(CoroResults->getContext()));
    CoroResults->eraseFromParent();
    break;
  }

  // In non-unique continuation lowering, we signal completion by returning
  // a null continuation.
  case coro::ABI::Retcon: {
    assert(!cast<CoroEndInst>(End)->hasResults() &&
           "retcon coroutine should not return any values");
    maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
    auto RetTy = Shape.getResumeFunctionType()->getReturnType();
    auto RetStructTy = dyn_cast<StructType>(RetTy);
    PointerType *ContinuationTy =
        cast<PointerType>(RetStructTy ? RetStructTy->getElementType(0) : RetTy);

    Value *ReturnValue = ConstantPointerNull::get(ContinuationTy);
    if (RetStructTy) {
      ReturnValue = Builder.CreateInsertValue(PoisonValue::get(RetStructTy),
                                              ReturnValue, 0);
    }
    Builder.CreateRet(ReturnValue);
    break;
  }
  }

  // Remove the rest of the block, by splitting it into an unreachable block.
  auto *BB = End->getParent();
  BB->splitBasicBlock(End);
  BB->getTerminator()->eraseFromParent();
}

// Mark a coroutine as done, which implies that the coroutine is finished and
// never get resumed.
//
// In resume-switched ABI, the done state is represented by storing zero in
// ResumeFnAddr.
//
// NOTE: We couldn't omit the argument `FramePtr`. It is necessary because the
// pointer to the frame in splitted function is not stored in `Shape`.
static void markCoroutineAsDone(IRBuilder<> &Builder, const coro::Shape &Shape,
                                Value *FramePtr) {
  assert(
      Shape.ABI == coro::ABI::Switch &&
      "markCoroutineAsDone is only supported for Switch-Resumed ABI for now.");
  auto *GepIndex = Builder.CreateStructGEP(
      Shape.FrameTy, FramePtr, coro::Shape::SwitchFieldIndex::Resume,
      "ResumeFn.addr");
  auto *NullPtr = ConstantPointerNull::get(cast<PointerType>(
      Shape.FrameTy->getTypeAtIndex(coro::Shape::SwitchFieldIndex::Resume)));
  Builder.CreateStore(NullPtr, GepIndex);

  // If the coroutine don't have unwind coro end, we could omit the store to
  // the final suspend point since we could infer the coroutine is suspended
  // at the final suspend point by the nullness of ResumeFnAddr.
  // However, we can't skip it if the coroutine have unwind coro end. Since
  // the coroutine reaches unwind coro end is considered suspended at the
  // final suspend point (the ResumeFnAddr is null) but in fact the coroutine
  // didn't complete yet. We need the IndexVal for the final suspend point
  // to make the states clear.
  if (Shape.SwitchLowering.HasUnwindCoroEnd &&
      Shape.SwitchLowering.HasFinalSuspend) {
    assert(cast<CoroSuspendInst>(Shape.CoroSuspends.back())->isFinal() &&
           "The final suspend should only live in the last position of "
           "CoroSuspends.");
    ConstantInt *IndexVal = Shape.getIndex(Shape.CoroSuspends.size() - 1);
    auto *FinalIndex = Builder.CreateStructGEP(
        Shape.FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");

    Builder.CreateStore(IndexVal, FinalIndex);
  }
}

/// Replace an unwind call to llvm.coro.end.
static void replaceUnwindCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,
                                 Value *FramePtr, bool InResume,
                                 CallGraph *CG) {
  IRBuilder<> Builder(End);

  switch (Shape.ABI) {
  // In switch-lowering, this does nothing in the main function.
  case coro::ABI::Switch: {
    // In C++'s specification, the coroutine should be marked as done
    // if promise.unhandled_exception() throws.  The frontend will
    // call coro.end(true) along this path.
    //
    // FIXME: We should refactor this once there is other language
    // which uses Switch-Resumed style other than C++.
    markCoroutineAsDone(Builder, Shape, FramePtr);
    if (!InResume)
      return;
    break;
  }
  // In async lowering this does nothing.
  case coro::ABI::Async:
    break;
  // In continuation-lowering, this frees the continuation storage.
  case coro::ABI::Retcon:
  case coro::ABI::RetconOnce:
    maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
    break;
  }

  // If coro.end has an associated bundle, add cleanupret instruction.
  if (auto Bundle = End->getOperandBundle(LLVMContext::OB_funclet)) {
    auto *FromPad = cast<CleanupPadInst>(Bundle->Inputs[0]);
    auto *CleanupRet = Builder.CreateCleanupRet(FromPad, nullptr);
    End->getParent()->splitBasicBlock(End);
    CleanupRet->getParent()->getTerminator()->eraseFromParent();
  }
}

static void replaceCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,
                           Value *FramePtr, bool InResume, CallGraph *CG) {
  if (End->isUnwind())
    replaceUnwindCoroEnd(End, Shape, FramePtr, InResume, CG);
  else
    replaceFallthroughCoroEnd(End, Shape, FramePtr, InResume, CG);

  auto &Context = End->getContext();
  End->replaceAllUsesWith(InResume ? ConstantInt::getTrue(Context)
                                   : ConstantInt::getFalse(Context));
  End->eraseFromParent();
}

// In the resume function, we remove the last case  (when coro::Shape is built,
// the final suspend point (if present) is always the last element of
// CoroSuspends array) since it is an undefined behavior to resume a coroutine
// suspended at the final suspend point.
// In the destroy function, if it isn't possible that the ResumeFnAddr is NULL
// and the coroutine doesn't suspend at the final suspend point actually (this
// is possible since the coroutine is considered suspended at the final suspend
// point if promise.unhandled_exception() exits via an exception), we can
// remove the last case.
void coro::BaseCloner::handleFinalSuspend() {
  assert(Shape.ABI == coro::ABI::Switch &&
         Shape.SwitchLowering.HasFinalSuspend);

  if (isSwitchDestroyFunction() && Shape.SwitchLowering.HasUnwindCoroEnd)
    return;

  auto *Switch = cast<SwitchInst>(VMap[Shape.SwitchLowering.ResumeSwitch]);
  auto FinalCaseIt = std::prev(Switch->case_end());
  BasicBlock *ResumeBB = FinalCaseIt->getCaseSuccessor();
  Switch->removeCase(FinalCaseIt);
  if (isSwitchDestroyFunction()) {
    BasicBlock *OldSwitchBB = Switch->getParent();
    auto *NewSwitchBB = OldSwitchBB->splitBasicBlock(Switch, "Switch");
    Builder.SetInsertPoint(OldSwitchBB->getTerminator());

    if (NewF->isCoroOnlyDestroyWhenComplete()) {
      // When the coroutine can only be destroyed when complete, we don't need
      // to generate code for other cases.
      Builder.CreateBr(ResumeBB);
    } else {
      auto *GepIndex = Builder.CreateStructGEP(
          Shape.FrameTy, NewFramePtr, coro::Shape::SwitchFieldIndex::Resume,
          "ResumeFn.addr");
      auto *Load =
          Builder.CreateLoad(Shape.getSwitchResumePointerType(), GepIndex);
      auto *Cond = Builder.CreateIsNull(Load);
      Builder.CreateCondBr(Cond, ResumeBB, NewSwitchBB);
    }
    OldSwitchBB->getTerminator()->eraseFromParent();
  }
}

static FunctionType *
getFunctionTypeFromAsyncSuspend(AnyCoroSuspendInst *Suspend) {
  auto *AsyncSuspend = cast<CoroSuspendAsyncInst>(Suspend);
  auto *StructTy = cast<StructType>(AsyncSuspend->getType());
  auto &Context = Suspend->getParent()->getParent()->getContext();
  auto *VoidTy = Type::getVoidTy(Context);
  return FunctionType::get(VoidTy, StructTy->elements(), false);
}

static Function *createCloneDeclaration(Function &OrigF, coro::Shape &Shape,
                                        const Twine &Suffix,
                                        Module::iterator InsertBefore,
                                        AnyCoroSuspendInst *ActiveSuspend) {
  Module *M = OrigF.getParent();
  auto *FnTy = (Shape.ABI != coro::ABI::Async)
                   ? Shape.getResumeFunctionType()
                   : getFunctionTypeFromAsyncSuspend(ActiveSuspend);

  Function *NewF =
      Function::Create(FnTy, GlobalValue::LinkageTypes::InternalLinkage,
                       OrigF.getName() + Suffix);

  M->getFunctionList().insert(InsertBefore, NewF);

  return NewF;
}

/// Replace uses of the active llvm.coro.suspend.retcon/async call with the
/// arguments to the continuation function.
///
/// This assumes that the builder has a meaningful insertion point.
void coro::BaseCloner::replaceRetconOrAsyncSuspendUses() {
  assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce ||
         Shape.ABI == coro::ABI::Async);

  auto NewS = VMap[ActiveSuspend];
  if (NewS->use_empty())
    return;

  // Copy out all the continuation arguments after the buffer pointer into
  // an easily-indexed data structure for convenience.
  SmallVector<Value *, 8> Args;
  // The async ABI includes all arguments -- including the first argument.
  bool IsAsyncABI = Shape.ABI == coro::ABI::Async;
  for (auto I = IsAsyncABI ? NewF->arg_begin() : std::next(NewF->arg_begin()),
            E = NewF->arg_end();
       I != E; ++I)
    Args.push_back(&*I);

  // If the suspend returns a single scalar value, we can just do a simple
  // replacement.
  if (!isa<StructType>(NewS->getType())) {
    assert(Args.size() == 1);
    NewS->replaceAllUsesWith(Args.front());
    return;
  }

  // Try to peephole extracts of an aggregate return.
  for (Use &U : llvm::make_early_inc_range(NewS->uses())) {
    auto *EVI = dyn_cast<ExtractValueInst>(U.getUser());
    if (!EVI || EVI->getNumIndices() != 1)
      continue;

    EVI->replaceAllUsesWith(Args[EVI->getIndices().front()]);
    EVI->eraseFromParent();
  }

  // If we have no remaining uses, we're done.
  if (NewS->use_empty())
    return;

  // Otherwise, we need to create an aggregate.
  Value *Aggr = PoisonValue::get(NewS->getType());
  for (auto [Idx, Arg] : llvm::enumerate(Args))
    Aggr = Builder.CreateInsertValue(Aggr, Arg, Idx);

  NewS->replaceAllUsesWith(Aggr);
}

void coro::BaseCloner::replaceCoroSuspends() {
  Value *SuspendResult;

  switch (Shape.ABI) {
  // In switch lowering, replace coro.suspend with the appropriate value
  // for the type of function we're extracting.
  // Replacing coro.suspend with (0) will result in control flow proceeding to
  // a resume label associated with a suspend point, replacing it with (1) will
  // result in control flow proceeding to a cleanup label associated with this
  // suspend point.
  case coro::ABI::Switch:
    SuspendResult = Builder.getInt8(isSwitchDestroyFunction() ? 1 : 0);
    break;

  // In async lowering there are no uses of the result.
  case coro::ABI::Async:
    return;

  // In returned-continuation lowering, the arguments from earlier
  // continuations are theoretically arbitrary, and they should have been
  // spilled.
  case coro::ABI::RetconOnce:
  case coro::ABI::Retcon:
    return;
  }

  for (AnyCoroSuspendInst *CS : Shape.CoroSuspends) {
    // The active suspend was handled earlier.
    if (CS == ActiveSuspend)
      continue;

    auto *MappedCS = cast<AnyCoroSuspendInst>(VMap[CS]);
    MappedCS->replaceAllUsesWith(SuspendResult);
    MappedCS->eraseFromParent();
  }
}

void coro::BaseCloner::replaceCoroEnds() {
  for (AnyCoroEndInst *CE : Shape.CoroEnds) {
    // We use a null call graph because there's no call graph node for
    // the cloned function yet.  We'll just be rebuilding that later.
    auto *NewCE = cast<AnyCoroEndInst>(VMap[CE]);
    replaceCoroEnd(NewCE, Shape, NewFramePtr, /*in resume*/ true, nullptr);
  }
}

static void replaceSwiftErrorOps(Function &F, coro::Shape &Shape,
                                 ValueToValueMapTy *VMap) {
  if (Shape.ABI == coro::ABI::Async && Shape.CoroSuspends.empty())
    return;
  Value *CachedSlot = nullptr;
  auto getSwiftErrorSlot = [&](Type *ValueTy) -> Value * {
    if (CachedSlot)
      return CachedSlot;

    // Check if the function has a swifterror argument.
    for (auto &Arg : F.args()) {
      if (Arg.isSwiftError()) {
        CachedSlot = &Arg;
        return &Arg;
      }
    }

    // Create a swifterror alloca.
    IRBuilder<> Builder(&F.getEntryBlock(),
                        F.getEntryBlock().getFirstNonPHIOrDbg());
    auto Alloca = Builder.CreateAlloca(ValueTy);
    Alloca->setSwiftError(true);

    CachedSlot = Alloca;
    return Alloca;
  };

  for (CallInst *Op : Shape.SwiftErrorOps) {
    auto MappedOp = VMap ? cast<CallInst>((*VMap)[Op]) : Op;
    IRBuilder<> Builder(MappedOp);

    // If there are no arguments, this is a 'get' operation.
    Value *MappedResult;
    if (Op->arg_empty()) {
      auto ValueTy = Op->getType();
      auto Slot = getSwiftErrorSlot(ValueTy);
      MappedResult = Builder.CreateLoad(ValueTy, Slot);
    } else {
      assert(Op->arg_size() == 1);
      auto Value = MappedOp->getArgOperand(0);
      auto ValueTy = Value->getType();
      auto Slot = getSwiftErrorSlot(ValueTy);
      Builder.CreateStore(Value, Slot);
      MappedResult = Slot;
    }

    MappedOp->replaceAllUsesWith(MappedResult);
    MappedOp->eraseFromParent();
  }

  // If we're updating the original function, we've invalidated SwiftErrorOps.
  if (VMap == nullptr) {
    Shape.SwiftErrorOps.clear();
  }
}

/// Returns all DbgVariableIntrinsic in F.
static std::pair<SmallVector<DbgVariableIntrinsic *, 8>,
                 SmallVector<DbgVariableRecord *>>
collectDbgVariableIntrinsics(Function &F) {
  SmallVector<DbgVariableIntrinsic *, 8> Intrinsics;
  SmallVector<DbgVariableRecord *> DbgVariableRecords;
  for (auto &I : instructions(F)) {
    for (DbgVariableRecord &DVR : filterDbgVars(I.getDbgRecordRange()))
      DbgVariableRecords.push_back(&DVR);
    if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&I))
      Intrinsics.push_back(DVI);
  }
  return {Intrinsics, DbgVariableRecords};
}

void coro::BaseCloner::replaceSwiftErrorOps() {
  ::replaceSwiftErrorOps(*NewF, Shape, &VMap);
}

void coro::BaseCloner::salvageDebugInfo() {
  auto [Worklist, DbgVariableRecords] = collectDbgVariableIntrinsics(*NewF);
  SmallDenseMap<Argument *, AllocaInst *, 4> ArgToAllocaMap;

  // Only 64-bit ABIs have a register we can refer to with the entry value.
  bool UseEntryValue =
      llvm::Triple(OrigF.getParent()->getTargetTriple()).isArch64Bit();
  for (DbgVariableIntrinsic *DVI : Worklist)
    coro::salvageDebugInfo(ArgToAllocaMap, *DVI, UseEntryValue);
  for (DbgVariableRecord *DVR : DbgVariableRecords)
    coro::salvageDebugInfo(ArgToAllocaMap, *DVR, UseEntryValue);

  // Remove all salvaged dbg.declare intrinsics that became
  // either unreachable or stale due to the CoroSplit transformation.
  DominatorTree DomTree(*NewF);
  auto IsUnreachableBlock = [&](BasicBlock *BB) {
    return !isPotentiallyReachable(&NewF->getEntryBlock(), BB, nullptr,
                                   &DomTree);
  };
  auto RemoveOne = [&](auto *DVI) {
    if (IsUnreachableBlock(DVI->getParent()))
      DVI->eraseFromParent();
    else if (isa_and_nonnull<AllocaInst>(DVI->getVariableLocationOp(0))) {
      // Count all non-debuginfo uses in reachable blocks.
      unsigned Uses = 0;
      for (auto *User : DVI->getVariableLocationOp(0)->users())
        if (auto *I = dyn_cast<Instruction>(User))
          if (!isa<AllocaInst>(I) && !IsUnreachableBlock(I->getParent()))
            ++Uses;
      if (!Uses)
        DVI->eraseFromParent();
    }
  };
  for_each(Worklist, RemoveOne);
  for_each(DbgVariableRecords, RemoveOne);
}

void coro::BaseCloner::replaceEntryBlock() {
  // In the original function, the AllocaSpillBlock is a block immediately
  // following the allocation of the frame object which defines GEPs for
  // all the allocas that have been moved into the frame, and it ends by
  // branching to the original beginning of the coroutine.  Make this
  // the entry block of the cloned function.
  auto *Entry = cast<BasicBlock>(VMap[Shape.AllocaSpillBlock]);
  auto *OldEntry = &NewF->getEntryBlock();
  Entry->setName("entry" + Suffix);
  Entry->moveBefore(OldEntry);
  Entry->getTerminator()->eraseFromParent();

  // Clear all predecessors of the new entry block.  There should be
  // exactly one predecessor, which we created when splitting out
  // AllocaSpillBlock to begin with.
  assert(Entry->hasOneUse());
  auto BranchToEntry = cast<BranchInst>(Entry->user_back());
  assert(BranchToEntry->isUnconditional());
  Builder.SetInsertPoint(BranchToEntry);
  Builder.CreateUnreachable();
  BranchToEntry->eraseFromParent();

  // Branch from the entry to the appropriate place.
  Builder.SetInsertPoint(Entry);
  switch (Shape.ABI) {
  case coro::ABI::Switch: {
    // In switch-lowering, we built a resume-entry block in the original
    // function.  Make the entry block branch to this.
    auto *SwitchBB =
        cast<BasicBlock>(VMap[Shape.SwitchLowering.ResumeEntryBlock]);
    Builder.CreateBr(SwitchBB);
    break;
  }
  case coro::ABI::Async:
  case coro::ABI::Retcon:
  case coro::ABI::RetconOnce: {
    // In continuation ABIs, we want to branch to immediately after the
    // active suspend point.  Earlier phases will have put the suspend in its
    // own basic block, so just thread our jump directly to its successor.
    assert((Shape.ABI == coro::ABI::Async &&
            isa<CoroSuspendAsyncInst>(ActiveSuspend)) ||
           ((Shape.ABI == coro::ABI::Retcon ||
             Shape.ABI == coro::ABI::RetconOnce) &&
            isa<CoroSuspendRetconInst>(ActiveSuspend)));
    auto *MappedCS = cast<AnyCoroSuspendInst>(VMap[ActiveSuspend]);
    auto Branch = cast<BranchInst>(MappedCS->getNextNode());
    assert(Branch->isUnconditional());
    Builder.CreateBr(Branch->getSuccessor(0));
    break;
  }
  }

  // Any static alloca that's still being used but not reachable from the new
  // entry needs to be moved to the new entry.
  Function *F = OldEntry->getParent();
  DominatorTree DT{*F};
  for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
    auto *Alloca = dyn_cast<AllocaInst>(&I);
    if (!Alloca || I.use_empty())
      continue;
    if (DT.isReachableFromEntry(I.getParent()) ||
        !isa<ConstantInt>(Alloca->getArraySize()))
      continue;
    I.moveBefore(*Entry, Entry->getFirstInsertionPt());
  }
}

/// Derive the value of the new frame pointer.
Value *coro::BaseCloner::deriveNewFramePointer() {
  // Builder should be inserting to the front of the new entry block.

  switch (Shape.ABI) {
  // In switch-lowering, the argument is the frame pointer.
  case coro::ABI::Switch:
    return &*NewF->arg_begin();
  // In async-lowering, one of the arguments is an async context as determined
  // by the `llvm.coro.id.async` intrinsic. We can retrieve the async context of
  // the resume function from the async context projection function associated
  // with the active suspend. The frame is located as a tail to the async
  // context header.
  case coro::ABI::Async: {
    auto *ActiveAsyncSuspend = cast<CoroSuspendAsyncInst>(ActiveSuspend);
    auto ContextIdx = ActiveAsyncSuspend->getStorageArgumentIndex() & 0xff;
    auto *CalleeContext = NewF->getArg(ContextIdx);
    auto *ProjectionFunc =
        ActiveAsyncSuspend->getAsyncContextProjectionFunction();
    auto DbgLoc =
        cast<CoroSuspendAsyncInst>(VMap[ActiveSuspend])->getDebugLoc();
    // Calling i8* (i8*)
    auto *CallerContext = Builder.CreateCall(ProjectionFunc->getFunctionType(),
                                             ProjectionFunc, CalleeContext);
    CallerContext->setCallingConv(ProjectionFunc->getCallingConv());
    CallerContext->setDebugLoc(DbgLoc);
    // The frame is located after the async_context header.
    auto &Context = Builder.getContext();
    auto *FramePtrAddr = Builder.CreateConstInBoundsGEP1_32(
        Type::getInt8Ty(Context), CallerContext,
        Shape.AsyncLowering.FrameOffset, "async.ctx.frameptr");
    // Inline the projection function.
    InlineFunctionInfo InlineInfo;
    auto InlineRes = InlineFunction(*CallerContext, InlineInfo);
    assert(InlineRes.isSuccess());
    (void)InlineRes;
    return FramePtrAddr;
  }
  // In continuation-lowering, the argument is the opaque storage.
  case coro::ABI::Retcon:
  case coro::ABI::RetconOnce: {
    Argument *NewStorage = &*NewF->arg_begin();
    auto FramePtrTy = PointerType::getUnqual(Shape.FrameTy->getContext());

    // If the storage is inline, just bitcast to the storage to the frame type.
    if (Shape.RetconLowering.IsFrameInlineInStorage)
      return NewStorage;

    // Otherwise, load the real frame from the opaque storage.
    return Builder.CreateLoad(FramePtrTy, NewStorage);
  }
  }
  llvm_unreachable("bad ABI");
}

/// Adjust the scope line of the funclet to the first line number after the
/// suspend point. This avoids a jump in the line table from the function
/// declaration (where prologue instructions are attributed to) to the suspend
/// point.
/// Only adjust the scope line when the files are the same.
/// If no candidate line number is found, fallback to the line of ActiveSuspend.
static void updateScopeLine(Instruction *ActiveSuspend,
                            DISubprogram &SPToUpdate) {
  if (!ActiveSuspend)
    return;

  // No subsequent instruction -> fallback to the location of ActiveSuspend.
  if (!ActiveSuspend->getNextNonDebugInstruction()) {
    if (auto DL = ActiveSuspend->getDebugLoc())
      if (SPToUpdate.getFile() == DL->getFile())
        SPToUpdate.setScopeLine(DL->getLine());
    return;
  }

  BasicBlock::iterator Successor =
      ActiveSuspend->getNextNonDebugInstruction()->getIterator();
  // Corosplit splits the BB around ActiveSuspend, so the meaningful
  // instructions are not in the same BB.
  if (auto *Branch = dyn_cast_or_null<BranchInst>(Successor);
      Branch && Branch->isUnconditional())
    Successor = Branch->getSuccessor(0)->getFirstNonPHIOrDbg();

  // Find the first successor of ActiveSuspend with a non-zero line location.
  // If that matches the file of ActiveSuspend, use it.
  BasicBlock *PBB = Successor->getParent();
  for (; Successor != PBB->end(); Successor = std::next(Successor)) {
    Successor = skipDebugIntrinsics(Successor);
    auto DL = Successor->getDebugLoc();
    if (!DL || DL.getLine() == 0)
      continue;

    if (SPToUpdate.getFile() == DL->getFile()) {
      SPToUpdate.setScopeLine(DL.getLine());
      return;
    }

    break;
  }

  // If the search above failed, fallback to the location of ActiveSuspend.
  if (auto DL = ActiveSuspend->getDebugLoc())
    if (SPToUpdate.getFile() == DL->getFile())
      SPToUpdate.setScopeLine(DL->getLine());
}

static void addFramePointerAttrs(AttributeList &Attrs, LLVMContext &Context,
                                 unsigned ParamIndex, uint64_t Size,
                                 Align Alignment, bool NoAlias) {
  AttrBuilder ParamAttrs(Context);
  ParamAttrs.addAttribute(Attribute::NonNull);
  ParamAttrs.addAttribute(Attribute::NoUndef);

  if (NoAlias)
    ParamAttrs.addAttribute(Attribute::NoAlias);

  ParamAttrs.addAlignmentAttr(Alignment);
  ParamAttrs.addDereferenceableAttr(Size);
  Attrs = Attrs.addParamAttributes(Context, ParamIndex, ParamAttrs);
}

static void addAsyncContextAttrs(AttributeList &Attrs, LLVMContext &Context,
                                 unsigned ParamIndex) {
  AttrBuilder ParamAttrs(Context);
  ParamAttrs.addAttribute(Attribute::SwiftAsync);
  Attrs = Attrs.addParamAttributes(Context, ParamIndex, ParamAttrs);
}

static void addSwiftSelfAttrs(AttributeList &Attrs, LLVMContext &Context,
                              unsigned ParamIndex) {
  AttrBuilder ParamAttrs(Context);
  ParamAttrs.addAttribute(Attribute::SwiftSelf);
  Attrs = Attrs.addParamAttributes(Context, ParamIndex, ParamAttrs);
}

/// Clone the body of the original function into a resume function of
/// some sort.
void coro::BaseCloner::create() {
  assert(NewF);

  // Replace all args with dummy instructions. If an argument is the old frame
  // pointer, the dummy will be replaced by the new frame pointer once it is
  // computed below. Uses of all other arguments should have already been
  // rewritten by buildCoroutineFrame() to use loads/stores on the coroutine
  // frame.
  SmallVector<Instruction *> DummyArgs;
  for (Argument &A : OrigF.args()) {
    DummyArgs.push_back(new FreezeInst(PoisonValue::get(A.getType())));
    VMap[&A] = DummyArgs.back();
  }

  SmallVector<ReturnInst *, 4> Returns;

  // Ignore attempts to change certain attributes of the function.
  // TODO: maybe there should be a way to suppress this during cloning?
  auto savedVisibility = NewF->getVisibility();
  auto savedUnnamedAddr = NewF->getUnnamedAddr();
  auto savedDLLStorageClass = NewF->getDLLStorageClass();

  // NewF's linkage (which CloneFunctionInto does *not* change) might not
  // be compatible with the visibility of OrigF (which it *does* change),
  // so protect against that.
  auto savedLinkage = NewF->getLinkage();
  NewF->setLinkage(llvm::GlobalValue::ExternalLinkage);

  CloneFunctionAttributesInto(NewF, &OrigF, VMap, false);
  CloneFunctionMetadataInto(*NewF, OrigF, VMap, RF_None, nullptr, nullptr,
                            &CommonDebugInfo);
  CloneFunctionBodyInto(*NewF, OrigF, VMap, RF_None, Returns, "", nullptr,
                        nullptr, nullptr, &CommonDebugInfo);

  auto &Context = NewF->getContext();

  if (DISubprogram *SP = NewF->getSubprogram()) {
    assert(SP != OrigF.getSubprogram() && SP->isDistinct());
    updateScopeLine(ActiveSuspend, *SP);

    // Update the linkage name to reflect the modified symbol name. It
    // is necessary to update the linkage name in Swift, since the
    // mangling changes for resume functions. It might also be the
    // right thing to do in C++, but due to a limitation in LLVM's
    // AsmPrinter we can only do this if the function doesn't have an
    // abstract specification, since the DWARF backend expects the
    // abstract specification to contain the linkage name and asserts
    // that they are identical.
    if (SP->getUnit() &&
        SP->getUnit()->getSourceLanguage() == dwarf::DW_LANG_Swift) {
      SP->replaceLinkageName(MDString::get(Context, NewF->getName()));
      if (auto *Decl = SP->getDeclaration()) {
        auto *NewDecl = DISubprogram::get(
            Decl->getContext(), Decl->getScope(), Decl->getName(),
            NewF->getName(), Decl->getFile(), Decl->getLine(), Decl->getType(),
            Decl->getScopeLine(), Decl->getContainingType(),
            Decl->getVirtualIndex(), Decl->getThisAdjustment(),
            Decl->getFlags(), Decl->getSPFlags(), Decl->getUnit(),
            Decl->getTemplateParams(), nullptr, Decl->getRetainedNodes(),
            Decl->getThrownTypes(), Decl->getAnnotations(),
            Decl->getTargetFuncName());
        SP->replaceDeclaration(NewDecl);
      }
    }
  }

  NewF->setLinkage(savedLinkage);
  NewF->setVisibility(savedVisibility);
  NewF->setUnnamedAddr(savedUnnamedAddr);
  NewF->setDLLStorageClass(savedDLLStorageClass);
  // The function sanitizer metadata needs to match the signature of the
  // function it is being attached to. However this does not hold for split
  // functions here. Thus remove the metadata for split functions.
  if (Shape.ABI == coro::ABI::Switch &&
      NewF->hasMetadata(LLVMContext::MD_func_sanitize))
    NewF->eraseMetadata(LLVMContext::MD_func_sanitize);

  // Replace the attributes of the new function:
  auto OrigAttrs = NewF->getAttributes();
  auto NewAttrs = AttributeList();

  switch (Shape.ABI) {
  case coro::ABI::Switch:
    // Bootstrap attributes by copying function attributes from the
    // original function.  This should include optimization settings and so on.
    NewAttrs = NewAttrs.addFnAttributes(
        Context, AttrBuilder(Context, OrigAttrs.getFnAttrs()));

    addFramePointerAttrs(NewAttrs, Context, 0, Shape.FrameSize,
                         Shape.FrameAlign, /*NoAlias=*/false);
    break;
  case coro::ABI::Async: {
    auto *ActiveAsyncSuspend = cast<CoroSuspendAsyncInst>(ActiveSuspend);
    if (OrigF.hasParamAttribute(Shape.AsyncLowering.ContextArgNo,
                                Attribute::SwiftAsync)) {
      uint32_t ArgAttributeIndices =
          ActiveAsyncSuspend->getStorageArgumentIndex();
      auto ContextArgIndex = ArgAttributeIndices & 0xff;
      addAsyncContextAttrs(NewAttrs, Context, ContextArgIndex);

      // `swiftasync` must preceed `swiftself` so 0 is not a valid index for
      // `swiftself`.
      auto SwiftSelfIndex = ArgAttributeIndices >> 8;
      if (SwiftSelfIndex)
        addSwiftSelfAttrs(NewAttrs, Context, SwiftSelfIndex);
    }

    // Transfer the original function's attributes.
    auto FnAttrs = OrigF.getAttributes().getFnAttrs();
    NewAttrs = NewAttrs.addFnAttributes(Context, AttrBuilder(Context, FnAttrs));
    break;
  }
  case coro::ABI::Retcon:
  case coro::ABI::RetconOnce:
    // If we have a continuation prototype, just use its attributes,
    // full-stop.
    NewAttrs = Shape.RetconLowering.ResumePrototype->getAttributes();

    /// FIXME: Is it really good to add the NoAlias attribute?
    addFramePointerAttrs(NewAttrs, Context, 0,
                         Shape.getRetconCoroId()->getStorageSize(),
                         Shape.getRetconCoroId()->getStorageAlignment(),
                         /*NoAlias=*/true);

    break;
  }

  switch (Shape.ABI) {
  // In these ABIs, the cloned functions always return 'void', and the
  // existing return sites are meaningless.  Note that for unique
  // continuations, this includes the returns associated with suspends;
  // this is fine because we can't suspend twice.
  case coro::ABI::Switch:
  case coro::ABI::RetconOnce:
    // Remove old returns.
    for (ReturnInst *Return : Returns)
      changeToUnreachable(Return);
    break;

  // With multi-suspend continuations, we'll already have eliminated the
  // original returns and inserted returns before all the suspend points,
  // so we want to leave any returns in place.
  case coro::ABI::Retcon:
    break;
  // Async lowering will insert musttail call functions at all suspend points
  // followed by a return.
  // Don't change returns to unreachable because that will trip up the verifier.
  // These returns should be unreachable from the clone.
  case coro::ABI::Async:
    break;
  }

  NewF->setAttributes(NewAttrs);
  NewF->setCallingConv(Shape.getResumeFunctionCC());

  // Set up the new entry block.
  replaceEntryBlock();

  // Turn symmetric transfers into musttail calls.
  for (CallInst *ResumeCall : Shape.SymmetricTransfers) {
    ResumeCall = cast<CallInst>(VMap[ResumeCall]);
    if (TTI.supportsTailCallFor(ResumeCall)) {
      // FIXME: Could we support symmetric transfer effectively without
      // musttail?
      ResumeCall->setTailCallKind(CallInst::TCK_MustTail);
    }

    // Put a 'ret void' after the call, and split any remaining instructions to
    // an unreachable block.
    BasicBlock *BB = ResumeCall->getParent();
    BB->splitBasicBlock(ResumeCall->getNextNode());
    Builder.SetInsertPoint(BB->getTerminator());
    Builder.CreateRetVoid();
    BB->getTerminator()->eraseFromParent();
  }

  Builder.SetInsertPoint(&NewF->getEntryBlock().front());
  NewFramePtr = deriveNewFramePointer();

  // Remap frame pointer.
  Value *OldFramePtr = VMap[Shape.FramePtr];
  NewFramePtr->takeName(OldFramePtr);
  OldFramePtr->replaceAllUsesWith(NewFramePtr);

  // Remap vFrame pointer.
  auto *NewVFrame = Builder.CreateBitCast(
      NewFramePtr, PointerType::getUnqual(Builder.getContext()), "vFrame");
  Value *OldVFrame = cast<Value>(VMap[Shape.CoroBegin]);
  if (OldVFrame != NewVFrame)
    OldVFrame->replaceAllUsesWith(NewVFrame);

  // All uses of the arguments should have been resolved by this point,
  // so we can safely remove the dummy values.
  for (Instruction *DummyArg : DummyArgs) {
    DummyArg->replaceAllUsesWith(PoisonValue::get(DummyArg->getType()));
    DummyArg->deleteValue();
  }

  switch (Shape.ABI) {
  case coro::ABI::Switch:
    // Rewrite final suspend handling as it is not done via switch (allows to
    // remove final case from the switch, since it is undefined behavior to
    // resume the coroutine suspended at the final suspend point.
    if (Shape.SwitchLowering.HasFinalSuspend)
      handleFinalSuspend();
    break;
  case coro::ABI::Async:
  case coro::ABI::Retcon:
  case coro::ABI::RetconOnce:
    // Replace uses of the active suspend with the corresponding
    // continuation-function arguments.
    assert(ActiveSuspend != nullptr &&
           "no active suspend when lowering a continuation-style coroutine");
    replaceRetconOrAsyncSuspendUses();
    break;
  }

  // Handle suspends.
  replaceCoroSuspends();

  // Handle swifterror.
  replaceSwiftErrorOps();

  // Remove coro.end intrinsics.
  replaceCoroEnds();

  // Salvage debug info that points into the coroutine frame.
  salvageDebugInfo();
}

void coro::SwitchCloner::create() {
  // Create a new function matching the original type
  NewF = createCloneDeclaration(OrigF, Shape, Suffix, OrigF.getParent()->end(),
                                ActiveSuspend);

  // Clone the function
  coro::BaseCloner::create();

  // Eliminate coro.free from the clones, replacing it with 'null' in cleanup,
  // to suppress deallocation code.
  coro::replaceCoroFree(cast<CoroIdInst>(VMap[Shape.CoroBegin->getId()]),
                        /*Elide=*/FKind == coro::CloneKind::SwitchCleanup);
}

static void updateAsyncFuncPointerContextSize(coro::Shape &Shape) {
  assert(Shape.ABI == coro::ABI::Async);

  auto *FuncPtrStruct = cast<ConstantStruct>(
      Shape.AsyncLowering.AsyncFuncPointer->getInitializer());
  auto *OrigRelativeFunOffset = FuncPtrStruct->getOperand(0);
  auto *OrigContextSize = FuncPtrStruct->getOperand(1);
  auto *NewContextSize = ConstantInt::get(OrigContextSize->getType(),
                                          Shape.AsyncLowering.ContextSize);
  auto *NewFuncPtrStruct = ConstantStruct::get(
      FuncPtrStruct->getType(), OrigRelativeFunOffset, NewContextSize);

  Shape.AsyncLowering.AsyncFuncPointer->setInitializer(NewFuncPtrStruct);
}

static TypeSize getFrameSizeForShape(coro::Shape &Shape) {
  // In the same function all coro.sizes should have the same result type.
  auto *SizeIntrin = Shape.CoroSizes.back();
  Module *M = SizeIntrin->getModule();
  const DataLayout &DL = M->getDataLayout();
  return DL.getTypeAllocSize(Shape.FrameTy);
}

static void replaceFrameSizeAndAlignment(coro::Shape &Shape) {
  if (Shape.ABI == coro::ABI::Async)
    updateAsyncFuncPointerContextSize(Shape);

  for (CoroAlignInst *CA : Shape.CoroAligns) {
    CA->replaceAllUsesWith(
        ConstantInt::get(CA->getType(), Shape.FrameAlign.value()));
    CA->eraseFromParent();
  }

  if (Shape.CoroSizes.empty())
    return;

  // In the same function all coro.sizes should have the same result type.
  auto *SizeIntrin = Shape.CoroSizes.back();
  auto *SizeConstant =
      ConstantInt::get(SizeIntrin->getType(), getFrameSizeForShape(Shape));

  for (CoroSizeInst *CS : Shape.CoroSizes) {
    CS->replaceAllUsesWith(SizeConstant);
    CS->eraseFromParent();
  }
}

static void postSplitCleanup(Function &F) {
  removeUnreachableBlocks(F);

#ifndef NDEBUG
  // For now, we do a mandatory verification step because we don't
  // entirely trust this pass.  Note that we don't want to add a verifier
  // pass to FPM below because it will also verify all the global data.
  if (verifyFunction(F, &errs()))
    report_fatal_error("Broken function");
#endif
}

// Coroutine has no suspend points. Remove heap allocation for the coroutine
// frame if possible.
static void handleNoSuspendCoroutine(coro::Shape &Shape) {
  auto *CoroBegin = Shape.CoroBegin;
  switch (Shape.ABI) {
  case coro::ABI::Switch: {
    auto SwitchId = Shape.getSwitchCoroId();
    auto *AllocInst = SwitchId->getCoroAlloc();
    coro::replaceCoroFree(SwitchId, /*Elide=*/AllocInst != nullptr);
    if (AllocInst) {
      IRBuilder<> Builder(AllocInst);
      auto *Frame = Builder.CreateAlloca(Shape.FrameTy);
      Frame->setAlignment(Shape.FrameAlign);
      AllocInst->replaceAllUsesWith(Builder.getFalse());
      AllocInst->eraseFromParent();
      CoroBegin->replaceAllUsesWith(Frame);
    } else {
      CoroBegin->replaceAllUsesWith(CoroBegin->getMem());
    }

    break;
  }
  case coro::ABI::Async:
  case coro::ABI::Retcon:
  case coro::ABI::RetconOnce:
    CoroBegin->replaceAllUsesWith(PoisonValue::get(CoroBegin->getType()));
    break;
  }

  CoroBegin->eraseFromParent();
  Shape.CoroBegin = nullptr;
}

// SimplifySuspendPoint needs to check that there is no calls between
// coro_save and coro_suspend, since any of the calls may potentially resume
// the coroutine and if that is the case we cannot eliminate the suspend point.
static bool hasCallsInBlockBetween(iterator_range<BasicBlock::iterator> R) {
  for (Instruction &I : R) {
    // Assume that no intrinsic can resume the coroutine.
    if (isa<IntrinsicInst>(I))
      continue;

    if (isa<CallBase>(I))
      return true;
  }
  return false;
}

static bool hasCallsInBlocksBetween(BasicBlock *SaveBB, BasicBlock *ResDesBB) {
  SmallPtrSet<BasicBlock *, 8> Set;
  SmallVector<BasicBlock *, 8> Worklist;

  Set.insert(SaveBB);
  Worklist.push_back(ResDesBB);

  // Accumulate all blocks between SaveBB and ResDesBB. Because CoroSaveIntr
  // returns a token consumed by suspend instruction, all blocks in between
  // will have to eventually hit SaveBB when going backwards from ResDesBB.
  while (!Worklist.empty()) {
    auto *BB = Worklist.pop_back_val();
    Set.insert(BB);
    for (auto *Pred : predecessors(BB))
      if (!Set.contains(Pred))
        Worklist.push_back(Pred);
  }

  // SaveBB and ResDesBB are checked separately in hasCallsBetween.
  Set.erase(SaveBB);
  Set.erase(ResDesBB);

  for (auto *BB : Set)
    if (hasCallsInBlockBetween({BB->getFirstNonPHIIt(), BB->end()}))
      return true;

  return false;
}

static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
  auto *SaveBB = Save->getParent();
  auto *ResumeOrDestroyBB = ResumeOrDestroy->getParent();
  BasicBlock::iterator SaveIt = Save->getIterator();
  BasicBlock::iterator ResumeOrDestroyIt = ResumeOrDestroy->getIterator();

  if (SaveBB == ResumeOrDestroyBB)
    return hasCallsInBlockBetween({std::next(SaveIt), ResumeOrDestroyIt});

  // Any calls from Save to the end of the block?
  if (hasCallsInBlockBetween({std::next(SaveIt), SaveBB->end()}))
    return true;

  // Any calls from begging of the block up to ResumeOrDestroy?
  if (hasCallsInBlockBetween(
          {ResumeOrDestroyBB->getFirstNonPHIIt(), ResumeOrDestroyIt}))
    return true;

  // Any calls in all of the blocks between SaveBB and ResumeOrDestroyBB?
  if (hasCallsInBlocksBetween(SaveBB, ResumeOrDestroyBB))
    return true;

  return false;
}

// If a SuspendIntrin is preceded by Resume or Destroy, we can eliminate the
// suspend point and replace it with nornal control flow.
static bool simplifySuspendPoint(CoroSuspendInst *Suspend,
                                 CoroBeginInst *CoroBegin) {
  Instruction *Prev = Suspend->getPrevNode();
  if (!Prev) {
    auto *Pred = Suspend->getParent()->getSinglePredecessor();
    if (!Pred)
      return false;
    Prev = Pred->getTerminator();
  }

  CallBase *CB = dyn_cast<CallBase>(Prev);
  if (!CB)
    return false;

  auto *Callee = CB->getCalledOperand()->stripPointerCasts();

  // See if the callsite is for resumption or destruction of the coroutine.
  auto *SubFn = dyn_cast<CoroSubFnInst>(Callee);
  if (!SubFn)
    return false;

  // Does not refer to the current coroutine, we cannot do anything with it.
  if (SubFn->getFrame() != CoroBegin)
    return false;

  // See if the transformation is safe. Specifically, see if there are any
  // calls in between Save and CallInstr. They can potenitally resume the
  // coroutine rendering this optimization unsafe.
  auto *Save = Suspend->getCoroSave();
  if (hasCallsBetween(Save, CB))
    return false;

  // Replace llvm.coro.suspend with the value that results in resumption over
  // the resume or cleanup path.
  Suspend->replaceAllUsesWith(SubFn->getRawIndex());
  Suspend->eraseFromParent();
  Save->eraseFromParent();

  // No longer need a call to coro.resume or coro.destroy.
  if (auto *Invoke = dyn_cast<InvokeInst>(CB)) {
    BranchInst::Create(Invoke->getNormalDest(), Invoke->getIterator());
  }

  // Grab the CalledValue from CB before erasing the CallInstr.
  auto *CalledValue = CB->getCalledOperand();
  CB->eraseFromParent();

  // If no more users remove it. Usually it is a bitcast of SubFn.
  if (CalledValue != SubFn && CalledValue->user_empty())
    if (auto *I = dyn_cast<Instruction>(CalledValue))
      I->eraseFromParent();

  // Now we are good to remove SubFn.
  if (SubFn->user_empty())
    SubFn->eraseFromParent();

  return true;
}

// Remove suspend points that are simplified.
static void simplifySuspendPoints(coro::Shape &Shape) {
  // Currently, the only simplification we do is switch-lowering-specific.
  if (Shape.ABI != coro::ABI::Switch)
    return;

  auto &S = Shape.CoroSuspends;
  size_t I = 0, N = S.size();
  if (N == 0)
    return;

  size_t ChangedFinalIndex = std::numeric_limits<size_t>::max();
  while (true) {
    auto SI = cast<CoroSuspendInst>(S[I]);
    // Leave final.suspend to handleFinalSuspend since it is undefined behavior
    // to resume a coroutine suspended at the final suspend point.
    if (!SI->isFinal() && simplifySuspendPoint(SI, Shape.CoroBegin)) {
      if (--N == I)
        break;

      std::swap(S[I], S[N]);

      if (cast<CoroSuspendInst>(S[I])->isFinal()) {
        assert(Shape.SwitchLowering.HasFinalSuspend);
        ChangedFinalIndex = I;
      }

      continue;
    }
    if (++I == N)
      break;
  }
  S.resize(N);

  // Maintain final.suspend in case final suspend was swapped.
  // Due to we requrie the final suspend to be the last element of CoroSuspends.
  if (ChangedFinalIndex < N) {
    assert(cast<CoroSuspendInst>(S[ChangedFinalIndex])->isFinal());
    std::swap(S[ChangedFinalIndex], S.back());
  }
}

namespace {

struct SwitchCoroutineSplitter {
  static void split(Function &F, coro::Shape &Shape,
                    SmallVectorImpl<Function *> &Clones,
                    TargetTransformInfo &TTI) {
    assert(Shape.ABI == coro::ABI::Switch);

    MetadataSetTy CommonDebugInfo{collectCommonDebugInfo(F)};

    // Create a resume clone by cloning the body of the original function,
    // setting new entry block and replacing coro.suspend an appropriate value
    // to force resume or cleanup pass for every suspend point.
    createResumeEntryBlock(F, Shape);
    auto *ResumeClone = coro::SwitchCloner::createClone(
        F, ".resume", Shape, coro::CloneKind::SwitchResume, TTI,
        CommonDebugInfo);
    auto *DestroyClone = coro::SwitchCloner::createClone(
        F, ".destroy", Shape, coro::CloneKind::SwitchUnwind, TTI,
        CommonDebugInfo);
    auto *CleanupClone = coro::SwitchCloner::createClone(
        F, ".cleanup", Shape, coro::CloneKind::SwitchCleanup, TTI,
        CommonDebugInfo);

    postSplitCleanup(*ResumeClone);
    postSplitCleanup(*DestroyClone);
    postSplitCleanup(*CleanupClone);

    // Store addresses resume/destroy/cleanup functions in the coroutine frame.
    updateCoroFrame(Shape, ResumeClone, DestroyClone, CleanupClone);

    assert(Clones.empty());
    Clones.push_back(ResumeClone);
    Clones.push_back(DestroyClone);
    Clones.push_back(CleanupClone);

    // Create a constant array referring to resume/destroy/clone functions
    // pointed by the last argument of @llvm.coro.info, so that CoroElide pass
    // can determined correct function to call.
    setCoroInfo(F, Shape, Clones);
  }

  // Create a variant of ramp function that does not perform heap allocation
  // for a switch ABI coroutine.
  //
  // The newly split `.noalloc` ramp function has the following differences:
  //  - Has one additional frame pointer parameter in lieu of dynamic
  //  allocation.
  //  - Suppressed allocations by replacing coro.alloc and coro.free.
  static Function *createNoAllocVariant(Function &F, coro::Shape &Shape,
                                        SmallVectorImpl<Function *> &Clones) {
    assert(Shape.ABI == coro::ABI::Switch);
    auto *OrigFnTy = F.getFunctionType();
    auto OldParams = OrigFnTy->params();

    SmallVector<Type *> NewParams;
    NewParams.reserve(OldParams.size() + 1);
    NewParams.append(OldParams.begin(), OldParams.end());
    NewParams.push_back(PointerType::getUnqual(Shape.FrameTy->getContext()));

    auto *NewFnTy = FunctionType::get(OrigFnTy->getReturnType(), NewParams,
                                      OrigFnTy->isVarArg());
    Function *NoAllocF =
        Function::Create(NewFnTy, F.getLinkage(), F.getName() + ".noalloc");

    ValueToValueMapTy VMap;
    unsigned int Idx = 0;
    for (const auto &I : F.args()) {
      VMap[&I] = NoAllocF->getArg(Idx++);
    }
    // We just appended the frame pointer as the last argument of the new
    // function.
    auto FrameIdx = NoAllocF->arg_size() - 1;
    SmallVector<ReturnInst *, 4> Returns;
    CloneFunctionInto(NoAllocF, &F, VMap,
                      CloneFunctionChangeType::LocalChangesOnly, Returns);

    if (Shape.CoroBegin) {
      auto *NewCoroBegin =
          cast_if_present<CoroBeginInst>(VMap[Shape.CoroBegin]);
      auto *NewCoroId = cast<CoroIdInst>(NewCoroBegin->getId());
      coro::replaceCoroFree(NewCoroId, /*Elide=*/true);
      coro::suppressCoroAllocs(NewCoroId);
      NewCoroBegin->replaceAllUsesWith(NoAllocF->getArg(FrameIdx));
      NewCoroBegin->eraseFromParent();
    }

    Module *M = F.getParent();
    M->getFunctionList().insert(M->end(), NoAllocF);

    removeUnreachableBlocks(*NoAllocF);
    auto NewAttrs = NoAllocF->getAttributes();
    // When we elide allocation, we read these attributes to determine the
    // frame size and alignment.
    addFramePointerAttrs(NewAttrs, NoAllocF->getContext(), FrameIdx,
                         Shape.FrameSize, Shape.FrameAlign,
                         /*NoAlias=*/false);

    NoAllocF->setAttributes(NewAttrs);

    Clones.push_back(NoAllocF);
    // Reset the original function's coro info, make the new noalloc variant
    // connected to the original ramp function.
    setCoroInfo(F, Shape, Clones);
    // After copying, set the linkage to internal linkage. Original function
    // may have different linkage, but optimization dependent on this function
    // generally relies on LTO.
    NoAllocF->setLinkage(llvm::GlobalValue::InternalLinkage);
    return NoAllocF;
  }

private:
  // Create an entry block for a resume function with a switch that will jump to
  // suspend points.
  static void createResumeEntryBlock(Function &F, coro::Shape &Shape) {
    LLVMContext &C = F.getContext();

    // resume.entry:
    //  %index.addr = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32
    //  0, i32 2 % index = load i32, i32* %index.addr switch i32 %index, label
    //  %unreachable [
    //    i32 0, label %resume.0
    //    i32 1, label %resume.1
    //    ...
    //  ]

    auto *NewEntry = BasicBlock::Create(C, "resume.entry", &F);
    auto *UnreachBB = BasicBlock::Create(C, "unreachable", &F);

    IRBuilder<> Builder(NewEntry);
    auto *FramePtr = Shape.FramePtr;
    auto *FrameTy = Shape.FrameTy;
    auto *GepIndex = Builder.CreateStructGEP(
        FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");
    auto *Index = Builder.CreateLoad(Shape.getIndexType(), GepIndex, "index");
    auto *Switch =
        Builder.CreateSwitch(Index, UnreachBB, Shape.CoroSuspends.size());
    Shape.SwitchLowering.ResumeSwitch = Switch;

    size_t SuspendIndex = 0;
    for (auto *AnyS : Shape.CoroSuspends) {
      auto *S = cast<CoroSuspendInst>(AnyS);
      ConstantInt *IndexVal = Shape.getIndex(SuspendIndex);

      // Replace CoroSave with a store to Index:
      //    %index.addr = getelementptr %f.frame... (index field number)
      //    store i32 %IndexVal, i32* %index.addr1
      auto *Save = S->getCoroSave();
      Builder.SetInsertPoint(Save);
      if (S->isFinal()) {
        // The coroutine should be marked done if it reaches the final suspend
        // point.
        markCoroutineAsDone(Builder, Shape, FramePtr);
      } else {
        auto *GepIndex = Builder.CreateStructGEP(
            FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");
        Builder.CreateStore(IndexVal, GepIndex);
      }

      Save->replaceAllUsesWith(ConstantTokenNone::get(C));
      Save->eraseFromParent();

      // Split block before and after coro.suspend and add a jump from an entry
      // switch:
      //
      //  whateverBB:
      //    whatever
      //    %0 = call i8 @llvm.coro.suspend(token none, i1 false)
      //    switch i8 %0, label %suspend[i8 0, label %resume
      //                                 i8 1, label %cleanup]
      // becomes:
      //
      //  whateverBB:
      //     whatever
      //     br label %resume.0.landing
      //
      //  resume.0: ; <--- jump from the switch in the resume.entry
      //     %0 = tail call i8 @llvm.coro.suspend(token none, i1 false)
      //     br label %resume.0.landing
      //
      //  resume.0.landing:
      //     %1 = phi i8[-1, %whateverBB], [%0, %resume.0]
      //     switch i8 % 1, label %suspend [i8 0, label %resume
      //                                    i8 1, label %cleanup]

      auto *SuspendBB = S->getParent();
      auto *ResumeBB =
          SuspendBB->splitBasicBlock(S, "resume." + Twine(SuspendIndex));
      auto *LandingBB = ResumeBB->splitBasicBlock(
          S->getNextNode(), ResumeBB->getName() + Twine(".landing"));
      Switch->addCase(IndexVal, ResumeBB);

      cast<BranchInst>(SuspendBB->getTerminator())->setSuccessor(0, LandingBB);
      auto *PN = PHINode::Create(Builder.getInt8Ty(), 2, "");
      PN->insertBefore(LandingBB->begin());
      S->replaceAllUsesWith(PN);
      PN->addIncoming(Builder.getInt8(-1), SuspendBB);
      PN->addIncoming(S, ResumeBB);

      ++SuspendIndex;
    }

    Builder.SetInsertPoint(UnreachBB);
    Builder.CreateUnreachable();

    Shape.SwitchLowering.ResumeEntryBlock = NewEntry;
  }

  // Store addresses of Resume/Destroy/Cleanup functions in the coroutine frame.
  static void updateCoroFrame(coro::Shape &Shape, Function *ResumeFn,
                              Function *DestroyFn, Function *CleanupFn) {
    IRBuilder<> Builder(&*Shape.getInsertPtAfterFramePtr());

    auto *ResumeAddr = Builder.CreateStructGEP(
        Shape.FrameTy, Shape.FramePtr, coro::Shape::SwitchFieldIndex::Resume,
        "resume.addr");
    Builder.CreateStore(ResumeFn, ResumeAddr);

    Value *DestroyOrCleanupFn = DestroyFn;

    CoroIdInst *CoroId = Shape.getSwitchCoroId();
    if (CoroAllocInst *CA = CoroId->getCoroAlloc()) {
      // If there is a CoroAlloc and it returns false (meaning we elide the
      // allocation, use CleanupFn instead of DestroyFn).
      DestroyOrCleanupFn = Builder.CreateSelect(CA, DestroyFn, CleanupFn);
    }

    auto *DestroyAddr = Builder.CreateStructGEP(
        Shape.FrameTy, Shape.FramePtr, coro::Shape::SwitchFieldIndex::Destroy,
        "destroy.addr");
    Builder.CreateStore(DestroyOrCleanupFn, DestroyAddr);
  }

  // Create a global constant array containing pointers to functions provided
  // and set Info parameter of CoroBegin to point at this constant. Example:
  //
  //   @f.resumers = internal constant [2 x void(%f.frame*)*]
  //                    [void(%f.frame*)* @f.resume, void(%f.frame*)*
  //                    @f.destroy]
  //   define void @f() {
  //     ...
  //     call i8* @llvm.coro.begin(i8* null, i32 0, i8* null,
  //                    i8* bitcast([2 x void(%f.frame*)*] * @f.resumers to
  //                    i8*))
  //
  // Assumes that all the functions have the same signature.
  static void setCoroInfo(Function &F, coro::Shape &Shape,
                          ArrayRef<Function *> Fns) {
    // This only works under the switch-lowering ABI because coro elision
    // only works on the switch-lowering ABI.
    SmallVector<Constant *, 4> Args(Fns);
    assert(!Args.empty());
    Function *Part = *Fns.begin();
    Module *M = Part->getParent();
    auto *ArrTy = ArrayType::get(Part->getType(), Args.size());

    auto *ConstVal = ConstantArray::get(ArrTy, Args);
    auto *GV = new GlobalVariable(*M, ConstVal->getType(), /*isConstant=*/true,
                                  GlobalVariable::PrivateLinkage, ConstVal,
                                  F.getName() + Twine(".resumers"));

    // Update coro.begin instruction to refer to this constant.
    LLVMContext &C = F.getContext();
    auto *BC = ConstantExpr::getPointerCast(GV, PointerType::getUnqual(C));
    Shape.getSwitchCoroId()->setInfo(BC);
  }
};

} // namespace

static void replaceAsyncResumeFunction(CoroSuspendAsyncInst *Suspend,
                                       Value *Continuation) {
  auto *ResumeIntrinsic = Suspend->getResumeFunction();
  auto &Context = Suspend->getParent()->getParent()->getContext();
  auto *Int8PtrTy = PointerType::getUnqual(Context);

  IRBuilder<> Builder(ResumeIntrinsic);
  auto *Val = Builder.CreateBitOrPointerCast(Continuation, Int8PtrTy);
  ResumeIntrinsic->replaceAllUsesWith(Val);
  ResumeIntrinsic->eraseFromParent();
  Suspend->setOperand(CoroSuspendAsyncInst::ResumeFunctionArg,
                      PoisonValue::get(Int8PtrTy));
}

/// Coerce the arguments in \p FnArgs according to \p FnTy in \p CallArgs.
static void coerceArguments(IRBuilder<> &Builder, FunctionType *FnTy,
                            ArrayRef<Value *> FnArgs,
                            SmallVectorImpl<Value *> &CallArgs) {
  size_t ArgIdx = 0;
  for (auto *paramTy : FnTy->params()) {
    assert(ArgIdx < FnArgs.size());
    if (paramTy != FnArgs[ArgIdx]->getType())
      CallArgs.push_back(
          Builder.CreateBitOrPointerCast(FnArgs[ArgIdx], paramTy));
    else
      CallArgs.push_back(FnArgs[ArgIdx]);
    ++ArgIdx;
  }
}

CallInst *coro::createMustTailCall(DebugLoc Loc, Function *MustTailCallFn,
                                   TargetTransformInfo &TTI,
                                   ArrayRef<Value *> Arguments,
                                   IRBuilder<> &Builder) {
  auto *FnTy = MustTailCallFn->getFunctionType();
  // Coerce the arguments, llvm optimizations seem to ignore the types in
  // vaarg functions and throws away casts in optimized mode.
  SmallVector<Value *, 8> CallArgs;
  coerceArguments(Builder, FnTy, Arguments, CallArgs);

  auto *TailCall = Builder.CreateCall(FnTy, MustTailCallFn, CallArgs);
  // Skip targets which don't support tail call.
  if (TTI.supportsTailCallFor(TailCall)) {
    TailCall->setTailCallKind(CallInst::TCK_MustTail);
  }
  TailCall->setDebugLoc(Loc);
  TailCall->setCallingConv(MustTailCallFn->getCallingConv());
  return TailCall;
}

void coro::AsyncABI::splitCoroutine(Function &F, coro::Shape &Shape,
                                    SmallVectorImpl<Function *> &Clones,
                                    TargetTransformInfo &TTI) {
  assert(Shape.ABI == coro::ABI::Async);
  assert(Clones.empty());
  // Reset various things that the optimizer might have decided it
  // "knows" about the coroutine function due to not seeing a return.
  F.removeFnAttr(Attribute::NoReturn);
  F.removeRetAttr(Attribute::NoAlias);
  F.removeRetAttr(Attribute::NonNull);

  auto &Context = F.getContext();
  auto *Int8PtrTy = PointerType::getUnqual(Context);

  auto *Id = Shape.getAsyncCoroId();
  IRBuilder<> Builder(Id);

  auto *FramePtr = Id->getStorage();
  FramePtr = Builder.CreateBitOrPointerCast(FramePtr, Int8PtrTy);
  FramePtr = Builder.CreateConstInBoundsGEP1_32(
      Type::getInt8Ty(Context), FramePtr, Shape.AsyncLowering.FrameOffset,
      "async.ctx.frameptr");

  // Map all uses of llvm.coro.begin to the allocated frame pointer.
  {
    // Make sure we don't invalidate Shape.FramePtr.
    TrackingVH<Value> Handle(Shape.FramePtr);
    Shape.CoroBegin->replaceAllUsesWith(FramePtr);
    Shape.FramePtr = Handle.getValPtr();
  }

  // Create all the functions in order after the main function.
  auto NextF = std::next(F.getIterator());

  // Create a continuation function for each of the suspend points.
  Clones.reserve(Shape.CoroSuspends.size());
  for (auto [Idx, CS] : llvm::enumerate(Shape.CoroSuspends)) {
    auto *Suspend = cast<CoroSuspendAsyncInst>(CS);

    // Create the clone declaration.
    auto ResumeNameSuffix = ".resume.";
    auto ProjectionFunctionName =
        Suspend->getAsyncContextProjectionFunction()->getName();
    bool UseSwiftMangling = false;
    if (ProjectionFunctionName == "__swift_async_resume_project_context") {
      ResumeNameSuffix = "TQ";
      UseSwiftMangling = true;
    } else if (ProjectionFunctionName == "__swift_async_resume_get_context") {
      ResumeNameSuffix = "TY";
      UseSwiftMangling = true;
    }
    auto *Continuation = createCloneDeclaration(
        F, Shape,
        UseSwiftMangling ? ResumeNameSuffix + Twine(Idx) + "_"
                         : ResumeNameSuffix + Twine(Idx),
        NextF, Suspend);
    Clones.push_back(Continuation);

    // Insert a branch to a new return block immediately before the suspend
    // point.
    auto *SuspendBB = Suspend->getParent();
    auto *NewSuspendBB = SuspendBB->splitBasicBlock(Suspend);
    auto *Branch = cast<BranchInst>(SuspendBB->getTerminator());

    // Place it before the first suspend.
    auto *ReturnBB =
        BasicBlock::Create(F.getContext(), "coro.return", &F, NewSuspendBB);
    Branch->setSuccessor(0, ReturnBB);

    IRBuilder<> Builder(ReturnBB);

    // Insert the call to the tail call function and inline it.
    auto *Fn = Suspend->getMustTailCallFunction();
    SmallVector<Value *, 8> Args(Suspend->args());
    auto FnArgs = ArrayRef<Value *>(Args).drop_front(
        CoroSuspendAsyncInst::MustTailCallFuncArg + 1);
    auto *TailCall = coro::createMustTailCall(Suspend->getDebugLoc(), Fn, TTI,
                                              FnArgs, Builder);
    Builder.CreateRetVoid();
    InlineFunctionInfo FnInfo;
    (void)InlineFunction(*TailCall, FnInfo);

    // Replace the lvm.coro.async.resume intrisic call.
    replaceAsyncResumeFunction(Suspend, Continuation);
  }

  assert(Clones.size() == Shape.CoroSuspends.size());

  MetadataSetTy CommonDebugInfo{collectCommonDebugInfo(F)};

  for (auto [Idx, CS] : llvm::enumerate(Shape.CoroSuspends)) {
    auto *Suspend = CS;
    auto *Clone = Clones[Idx];

    coro::BaseCloner::createClone(F, "resume." + Twine(Idx), Shape, Clone,
                                  Suspend, TTI, CommonDebugInfo);
  }
}

void coro::AnyRetconABI::splitCoroutine(Function &F, coro::Shape &Shape,
                                        SmallVectorImpl<Function *> &Clones,
                                        TargetTransformInfo &TTI) {
  assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce);
  assert(Clones.empty());

  // Reset various things that the optimizer might have decided it
  // "knows" about the coroutine function due to not seeing a return.
  F.removeFnAttr(Attribute::NoReturn);
  F.removeRetAttr(Attribute::NoAlias);
  F.removeRetAttr(Attribute::NonNull);

  // Allocate the frame.
  auto *Id = Shape.getRetconCoroId();
  Value *RawFramePtr;
  if (Shape.RetconLowering.IsFrameInlineInStorage) {
    RawFramePtr = Id->getStorage();
  } else {
    IRBuilder<> Builder(Id);

    // Determine the size of the frame.
    const DataLayout &DL = F.getDataLayout();
    auto Size = DL.getTypeAllocSize(Shape.FrameTy);

    // Allocate.  We don't need to update the call graph node because we're
    // going to recompute it from scratch after splitting.
    // FIXME: pass the required alignment
    RawFramePtr = Shape.emitAlloc(Builder, Builder.getInt64(Size), nullptr);
    RawFramePtr =
        Builder.CreateBitCast(RawFramePtr, Shape.CoroBegin->getType());

    // Stash the allocated frame pointer in the continuation storage.
    Builder.CreateStore(RawFramePtr, Id->getStorage());
  }

  // Map all uses of llvm.coro.begin to the allocated frame pointer.
  {
    // Make sure we don't invalidate Shape.FramePtr.
    TrackingVH<Value> Handle(Shape.FramePtr);
    Shape.CoroBegin->replaceAllUsesWith(RawFramePtr);
    Shape.FramePtr = Handle.getValPtr();
  }

  // Create a unique return block.
  BasicBlock *ReturnBB = nullptr;
  PHINode *ContinuationPhi = nullptr;
  SmallVector<PHINode *, 4> ReturnPHIs;

  // Create all the functions in order after the main function.
  auto NextF = std::next(F.getIterator());

  // Create a continuation function for each of the suspend points.
  Clones.reserve(Shape.CoroSuspends.size());
  for (auto [Idx, CS] : llvm::enumerate(Shape.CoroSuspends)) {
    auto Suspend = cast<CoroSuspendRetconInst>(CS);

    // Create the clone declaration.
    auto Continuation = createCloneDeclaration(
        F, Shape, ".resume." + Twine(Idx), NextF, nullptr);
    Clones.push_back(Continuation);

    // Insert a branch to the unified return block immediately before
    // the suspend point.
    auto SuspendBB = Suspend->getParent();
    auto NewSuspendBB = SuspendBB->splitBasicBlock(Suspend);
    auto Branch = cast<BranchInst>(SuspendBB->getTerminator());

    // Create the unified return block.
    if (!ReturnBB) {
      // Place it before the first suspend.
      ReturnBB =
          BasicBlock::Create(F.getContext(), "coro.return", &F, NewSuspendBB);
      Shape.RetconLowering.ReturnBlock = ReturnBB;

      IRBuilder<> Builder(ReturnBB);

      // First, the continuation.
      ContinuationPhi =
          Builder.CreatePHI(Continuation->getType(), Shape.CoroSuspends.size());

      // Create PHIs for all other return values.
      assert(ReturnPHIs.empty());

      // Next, all the directly-yielded values.
      for (auto *ResultTy : Shape.getRetconResultTypes())
        ReturnPHIs.push_back(
            Builder.CreatePHI(ResultTy, Shape.CoroSuspends.size()));

      // Build the return value.
      auto RetTy = F.getReturnType();

      // Cast the continuation value if necessary.
      // We can't rely on the types matching up because that type would
      // have to be infinite.
      auto CastedContinuationTy =
          (ReturnPHIs.empty() ? RetTy : RetTy->getStructElementType(0));
      auto *CastedContinuation =
          Builder.CreateBitCast(ContinuationPhi, CastedContinuationTy);

      Value *RetV = CastedContinuation;
      if (!ReturnPHIs.empty()) {
        auto ValueIdx = 0;
        RetV = PoisonValue::get(RetTy);
        RetV = Builder.CreateInsertValue(RetV, CastedContinuation, ValueIdx++);

        for (auto Phi : ReturnPHIs)
          RetV = Builder.CreateInsertValue(RetV, Phi, ValueIdx++);
      }

      Builder.CreateRet(RetV);
    }

    // Branch to the return block.
    Branch->setSuccessor(0, ReturnBB);
    assert(ContinuationPhi);
    ContinuationPhi->addIncoming(Continuation, SuspendBB);
    for (auto [Phi, VUse] :
         llvm::zip_equal(ReturnPHIs, Suspend->value_operands()))
      Phi->addIncoming(VUse, SuspendBB);
  }

  assert(Clones.size() == Shape.CoroSuspends.size());

  MetadataSetTy CommonDebugInfo{collectCommonDebugInfo(F)};

  for (auto [Idx, CS] : llvm::enumerate(Shape.CoroSuspends)) {
    auto Suspend = CS;
    auto Clone = Clones[Idx];

    coro::BaseCloner::createClone(F, "resume." + Twine(Idx), Shape, Clone,
                                  Suspend, TTI, CommonDebugInfo);
  }
}

namespace {
class PrettyStackTraceFunction : public PrettyStackTraceEntry {
  Function &F;

public:
  PrettyStackTraceFunction(Function &F) : F(F) {}
  void print(raw_ostream &OS) const override {
    OS << "While splitting coroutine ";
    F.printAsOperand(OS, /*print type*/ false, F.getParent());
    OS << "\n";
  }
};
} // namespace

/// Remove calls to llvm.coro.end in the original function.
static void removeCoroEndsFromRampFunction(const coro::Shape &Shape) {
  if (Shape.ABI != coro::ABI::Switch) {
    for (auto *End : Shape.CoroEnds) {
      replaceCoroEnd(End, Shape, Shape.FramePtr, /*in resume*/ false, nullptr);
    }
  } else {
    for (llvm::AnyCoroEndInst *End : Shape.CoroEnds) {
      auto &Context = End->getContext();
      End->replaceAllUsesWith(ConstantInt::getFalse(Context));
      End->eraseFromParent();
    }
  }
}

static bool hasSafeElideCaller(Function &F) {
  for (auto *U : F.users()) {
    if (auto *CB = dyn_cast<CallBase>(U)) {
      auto *Caller = CB->getFunction();
      if (Caller && Caller->isPresplitCoroutine() &&
          CB->hasFnAttr(llvm::Attribute::CoroElideSafe))
        return true;
    }
  }
  return false;
}

void coro::SwitchABI::splitCoroutine(Function &F, coro::Shape &Shape,
                                     SmallVectorImpl<Function *> &Clones,
                                     TargetTransformInfo &TTI) {
  SwitchCoroutineSplitter::split(F, Shape, Clones, TTI);
}

static void doSplitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
                             coro::BaseABI &ABI, TargetTransformInfo &TTI,
                             bool OptimizeFrame) {
  PrettyStackTraceFunction prettyStackTrace(F);

  auto &Shape = ABI.Shape;
  assert(Shape.CoroBegin);

  lowerAwaitSuspends(F, Shape);

  simplifySuspendPoints(Shape);

  normalizeCoroutine(F, Shape, TTI);
  ABI.buildCoroutineFrame(OptimizeFrame);
  replaceFrameSizeAndAlignment(Shape);

  bool isNoSuspendCoroutine = Shape.CoroSuspends.empty();

  bool shouldCreateNoAllocVariant =
      !isNoSuspendCoroutine && Shape.ABI == coro::ABI::Switch &&
      hasSafeElideCaller(F) && !F.hasFnAttribute(llvm::Attribute::NoInline);

  // If there are no suspend points, no split required, just remove
  // the allocation and deallocation blocks, they are not needed.
  if (isNoSuspendCoroutine) {
    handleNoSuspendCoroutine(Shape);
  } else {
    ABI.splitCoroutine(F, Shape, Clones, TTI);
  }

  // Replace all the swifterror operations in the original function.
  // This invalidates SwiftErrorOps in the Shape.
  replaceSwiftErrorOps(F, Shape, nullptr);

  // Salvage debug intrinsics that point into the coroutine frame in the
  // original function. The Cloner has already salvaged debug info in the new
  // coroutine funclets.
  SmallDenseMap<Argument *, AllocaInst *, 4> ArgToAllocaMap;
  auto [DbgInsts, DbgVariableRecords] = collectDbgVariableIntrinsics(F);
  for (auto *DDI : DbgInsts)
    coro::salvageDebugInfo(ArgToAllocaMap, *DDI, false /*UseEntryValue*/);
  for (DbgVariableRecord *DVR : DbgVariableRecords)
    coro::salvageDebugInfo(ArgToAllocaMap, *DVR, false /*UseEntryValue*/);

  removeCoroEndsFromRampFunction(Shape);

  if (shouldCreateNoAllocVariant)
    SwitchCoroutineSplitter::createNoAllocVariant(F, Shape, Clones);
}

static LazyCallGraph::SCC &updateCallGraphAfterCoroutineSplit(
    LazyCallGraph::Node &N, const coro::Shape &Shape,
    const SmallVectorImpl<Function *> &Clones, LazyCallGraph::SCC &C,
    LazyCallGraph &CG, CGSCCAnalysisManager &AM, CGSCCUpdateResult &UR,
    FunctionAnalysisManager &FAM) {

  auto *CurrentSCC = &C;
  if (!Clones.empty()) {
    switch (Shape.ABI) {
    case coro::ABI::Switch:
      // Each clone in the Switch lowering is independent of the other clones.
      // Let the LazyCallGraph know about each one separately.
      for (Function *Clone : Clones)
        CG.addSplitFunction(N.getFunction(), *Clone);
      break;
    case coro::ABI::Async:
    case coro::ABI::Retcon:
    case coro::ABI::RetconOnce:
      // Each clone in the Async/Retcon lowering references of the other clones.
      // Let the LazyCallGraph know about all of them at once.
      if (!Clones.empty())
        CG.addSplitRefRecursiveFunctions(N.getFunction(), Clones);
      break;
    }

    // Let the CGSCC infra handle the changes to the original function.
    CurrentSCC = &updateCGAndAnalysisManagerForCGSCCPass(CG, *CurrentSCC, N, AM,
                                                         UR, FAM);
  }

  // Do some cleanup and let the CGSCC infra see if we've cleaned up any edges
  // to the split functions.
  postSplitCleanup(N.getFunction());
  CurrentSCC = &updateCGAndAnalysisManagerForFunctionPass(CG, *CurrentSCC, N,
                                                          AM, UR, FAM);
  return *CurrentSCC;
}

/// Replace a call to llvm.coro.prepare.retcon.
static void replacePrepare(CallInst *Prepare, LazyCallGraph &CG,
                           LazyCallGraph::SCC &C) {
  auto CastFn = Prepare->getArgOperand(0); // as an i8*
  auto Fn = CastFn->stripPointerCasts();   // as its original type

  // Attempt to peephole this pattern:
  //    %0 = bitcast [[TYPE]] @some_function to i8*
  //    %1 = call @llvm.coro.prepare.retcon(i8* %0)
  //    %2 = bitcast %1 to [[TYPE]]
  // ==>
  //    %2 = @some_function
  for (Use &U : llvm::make_early_inc_range(Prepare->uses())) {
    // Look for bitcasts back to the original function type.
    auto *Cast = dyn_cast<BitCastInst>(U.getUser());
    if (!Cast || Cast->getType() != Fn->getType())
      continue;

    // Replace and remove the cast.
    Cast->replaceAllUsesWith(Fn);
    Cast->eraseFromParent();
  }

  // Replace any remaining uses with the function as an i8*.
  // This can never directly be a callee, so we don't need to update CG.
  Prepare->replaceAllUsesWith(CastFn);
  Prepare->eraseFromParent();

  // Kill dead bitcasts.
  while (auto *Cast = dyn_cast<BitCastInst>(CastFn)) {
    if (!Cast->use_empty())
      break;
    CastFn = Cast->getOperand(0);
    Cast->eraseFromParent();
  }
}

static bool replaceAllPrepares(Function *PrepareFn, LazyCallGraph &CG,
                               LazyCallGraph::SCC &C) {
  bool Changed = false;
  for (Use &P : llvm::make_early_inc_range(PrepareFn->uses())) {
    // Intrinsics can only be used in calls.
    auto *Prepare = cast<CallInst>(P.getUser());
    replacePrepare(Prepare, CG, C);
    Changed = true;
  }

  return Changed;
}

static void addPrepareFunction(const Module &M,
                               SmallVectorImpl<Function *> &Fns,
                               StringRef Name) {
  auto *PrepareFn = M.getFunction(Name);
  if (PrepareFn && !PrepareFn->use_empty())
    Fns.push_back(PrepareFn);
}

static std::unique_ptr<coro::BaseABI>
CreateNewABI(Function &F, coro::Shape &S,
             std::function<bool(Instruction &)> IsMatCallback,
             const SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs) {
  if (S.CoroBegin->hasCustomABI()) {
    unsigned CustomABI = S.CoroBegin->getCustomABI();
    if (CustomABI >= GenCustomABIs.size())
      llvm_unreachable("Custom ABI not found amoung those specified");
    return GenCustomABIs[CustomABI](F, S);
  }

  switch (S.ABI) {
  case coro::ABI::Switch:
    return std::make_unique<coro::SwitchABI>(F, S, IsMatCallback);
  case coro::ABI::Async:
    return std::make_unique<coro::AsyncABI>(F, S, IsMatCallback);
  case coro::ABI::Retcon:
    return std::make_unique<coro::AnyRetconABI>(F, S, IsMatCallback);
  case coro::ABI::RetconOnce:
    return std::make_unique<coro::AnyRetconABI>(F, S, IsMatCallback);
  }
  llvm_unreachable("Unknown ABI");
}

CoroSplitPass::CoroSplitPass(bool OptimizeFrame)
    : CreateAndInitABI([](Function &F, coro::Shape &S) {
        std::unique_ptr<coro::BaseABI> ABI =
            CreateNewABI(F, S, coro::isTriviallyMaterializable, {});
        ABI->init();
        return ABI;
      }),
      OptimizeFrame(OptimizeFrame) {}

CoroSplitPass::CoroSplitPass(
    SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs, bool OptimizeFrame)
    : CreateAndInitABI([=](Function &F, coro::Shape &S) {
        std::unique_ptr<coro::BaseABI> ABI =
            CreateNewABI(F, S, coro::isTriviallyMaterializable, GenCustomABIs);
        ABI->init();
        return ABI;
      }),
      OptimizeFrame(OptimizeFrame) {}

// For back compatibility, constructor takes a materializable callback and
// creates a generator for an ABI with a modified materializable callback.
CoroSplitPass::CoroSplitPass(std::function<bool(Instruction &)> IsMatCallback,
                             bool OptimizeFrame)
    : CreateAndInitABI([=](Function &F, coro::Shape &S) {
        std::unique_ptr<coro::BaseABI> ABI =
            CreateNewABI(F, S, IsMatCallback, {});
        ABI->init();
        return ABI;
      }),
      OptimizeFrame(OptimizeFrame) {}

// For back compatibility, constructor takes a materializable callback and
// creates a generator for an ABI with a modified materializable callback.
CoroSplitPass::CoroSplitPass(
    std::function<bool(Instruction &)> IsMatCallback,
    SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs, bool OptimizeFrame)
    : CreateAndInitABI([=](Function &F, coro::Shape &S) {
        std::unique_ptr<coro::BaseABI> ABI =
            CreateNewABI(F, S, IsMatCallback, GenCustomABIs);
        ABI->init();
        return ABI;
      }),
      OptimizeFrame(OptimizeFrame) {}

PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
                                     CGSCCAnalysisManager &AM,
                                     LazyCallGraph &CG, CGSCCUpdateResult &UR) {
  // NB: One invariant of a valid LazyCallGraph::SCC is that it must contain a
  //     non-zero number of nodes, so we assume that here and grab the first
  //     node's function's module.
  Module &M = *C.begin()->getFunction().getParent();
  auto &FAM =
      AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();

  // Check for uses of llvm.coro.prepare.retcon/async.
  SmallVector<Function *, 2> PrepareFns;
  addPrepareFunction(M, PrepareFns, "llvm.coro.prepare.retcon");
  addPrepareFunction(M, PrepareFns, "llvm.coro.prepare.async");

  // Find coroutines for processing.
  SmallVector<LazyCallGraph::Node *> Coroutines;
  for (LazyCallGraph::Node &N : C)
    if (N.getFunction().isPresplitCoroutine())
      Coroutines.push_back(&N);

  if (Coroutines.empty() && PrepareFns.empty())
    return PreservedAnalyses::all();

  auto *CurrentSCC = &C;
  // Split all the coroutines.
  for (LazyCallGraph::Node *N : Coroutines) {
    Function &F = N->getFunction();
    LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName()
                      << "\n");

    // The suspend-crossing algorithm in buildCoroutineFrame gets tripped up
    // by unreachable blocks, so remove them as a first pass. Remove the
    // unreachable blocks before collecting intrinsics into Shape.
    removeUnreachableBlocks(F);

    coro::Shape Shape(F);
    if (!Shape.CoroBegin)
      continue;

    F.setSplittedCoroutine();

    std::unique_ptr<coro::BaseABI> ABI = CreateAndInitABI(F, Shape);

    SmallVector<Function *, 4> Clones;
    auto &TTI = FAM.getResult<TargetIRAnalysis>(F);
    doSplitCoroutine(F, Clones, *ABI, TTI, OptimizeFrame);
    CurrentSCC = &updateCallGraphAfterCoroutineSplit(
        *N, Shape, Clones, *CurrentSCC, CG, AM, UR, FAM);

    auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
    ORE.emit([&]() {
      return OptimizationRemark(DEBUG_TYPE, "CoroSplit", &F)
             << "Split '" << ore::NV("function", F.getName())
             << "' (frame_size=" << ore::NV("frame_size", Shape.FrameSize)
             << ", align=" << ore::NV("align", Shape.FrameAlign.value()) << ")";
    });

    if (!Shape.CoroSuspends.empty()) {
      // Run the CGSCC pipeline on the original and newly split functions.
      UR.CWorklist.insert(CurrentSCC);
      for (Function *Clone : Clones)
        UR.CWorklist.insert(CG.lookupSCC(CG.get(*Clone)));
    }
  }

  for (auto *PrepareFn : PrepareFns) {
    replaceAllPrepares(PrepareFn, CG, *CurrentSCC);
  }

  return PreservedAnalyses::none();
}