brintos

brintos / llvm-project-archived public Read only

0
0
Text · 91.1 KiB · da2b953 Raw
2455 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8// This pass builds the coroutine frame and outlines resume and destroy parts9// of the coroutine into separate functions.10//11// We present a coroutine to an LLVM as an ordinary function with suspension12// points marked up with intrinsics. We let the optimizer party on the coroutine13// as a single function for as long as possible. Shortly before the coroutine is14// eligible to be inlined into its callers, we split up the coroutine into parts15// 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 the17// coroutine subfunctions we extracted before proceeding to the caller of the18// coroutine.19//===----------------------------------------------------------------------===//20 21#include "llvm/Transforms/Coroutines/CoroSplit.h"22#include "CoroCloner.h"23#include "CoroInternal.h"24#include "llvm/ADT/DenseMap.h"25#include "llvm/ADT/PriorityWorklist.h"26#include "llvm/ADT/STLExtras.h"27#include "llvm/ADT/SmallPtrSet.h"28#include "llvm/ADT/SmallVector.h"29#include "llvm/ADT/StringExtras.h"30#include "llvm/ADT/StringRef.h"31#include "llvm/ADT/Twine.h"32#include "llvm/Analysis/CFG.h"33#include "llvm/Analysis/CallGraph.h"34#include "llvm/Analysis/ConstantFolding.h"35#include "llvm/Analysis/LazyCallGraph.h"36#include "llvm/Analysis/OptimizationRemarkEmitter.h"37#include "llvm/Analysis/TargetTransformInfo.h"38#include "llvm/BinaryFormat/Dwarf.h"39#include "llvm/IR/Argument.h"40#include "llvm/IR/Attributes.h"41#include "llvm/IR/BasicBlock.h"42#include "llvm/IR/CFG.h"43#include "llvm/IR/CallingConv.h"44#include "llvm/IR/Constants.h"45#include "llvm/IR/DIBuilder.h"46#include "llvm/IR/DataLayout.h"47#include "llvm/IR/DebugInfo.h"48#include "llvm/IR/DerivedTypes.h"49#include "llvm/IR/Dominators.h"50#include "llvm/IR/GlobalValue.h"51#include "llvm/IR/GlobalVariable.h"52#include "llvm/IR/InstIterator.h"53#include "llvm/IR/InstrTypes.h"54#include "llvm/IR/Instruction.h"55#include "llvm/IR/Instructions.h"56#include "llvm/IR/IntrinsicInst.h"57#include "llvm/IR/LLVMContext.h"58#include "llvm/IR/Module.h"59#include "llvm/IR/Type.h"60#include "llvm/IR/Value.h"61#include "llvm/IR/Verifier.h"62#include "llvm/Support/Casting.h"63#include "llvm/Support/Debug.h"64#include "llvm/Support/PrettyStackTrace.h"65#include "llvm/Support/raw_ostream.h"66#include "llvm/Transforms/Coroutines/MaterializationUtils.h"67#include "llvm/Transforms/Scalar.h"68#include "llvm/Transforms/Utils/BasicBlockUtils.h"69#include "llvm/Transforms/Utils/CallGraphUpdater.h"70#include "llvm/Transforms/Utils/Cloning.h"71#include "llvm/Transforms/Utils/Local.h"72#include <cassert>73#include <cstddef>74#include <cstdint>75#include <initializer_list>76#include <iterator>77 78using namespace llvm;79 80#define DEBUG_TYPE "coro-split"81 82// FIXME:83// Lower the intrinisc in CoroEarly phase if coroutine frame doesn't escape84// and it is known that other transformations, for example, sanitizers85// won't lead to incorrect code.86static void lowerAwaitSuspend(IRBuilder<> &Builder, CoroAwaitSuspendInst *CB,87                              coro::Shape &Shape) {88  auto Wrapper = CB->getWrapperFunction();89  auto Awaiter = CB->getAwaiter();90  auto FramePtr = CB->getFrame();91 92  Builder.SetInsertPoint(CB);93 94  CallBase *NewCall = nullptr;95  // await_suspend has only 2 parameters, awaiter and handle.96  // Copy parameter attributes from the intrinsic call, but remove the last,97  // because the last parameter now becomes the function that is being called.98  AttributeList NewAttributes =99      CB->getAttributes().removeParamAttributes(CB->getContext(), 2);100 101  if (auto Invoke = dyn_cast<InvokeInst>(CB)) {102    auto WrapperInvoke =103        Builder.CreateInvoke(Wrapper, Invoke->getNormalDest(),104                             Invoke->getUnwindDest(), {Awaiter, FramePtr});105 106    WrapperInvoke->setCallingConv(Invoke->getCallingConv());107    std::copy(Invoke->bundle_op_info_begin(), Invoke->bundle_op_info_end(),108              WrapperInvoke->bundle_op_info_begin());109    WrapperInvoke->setAttributes(NewAttributes);110    WrapperInvoke->setDebugLoc(Invoke->getDebugLoc());111    NewCall = WrapperInvoke;112  } else if (auto Call = dyn_cast<CallInst>(CB)) {113    auto WrapperCall = Builder.CreateCall(Wrapper, {Awaiter, FramePtr});114 115    WrapperCall->setAttributes(NewAttributes);116    WrapperCall->setDebugLoc(Call->getDebugLoc());117    NewCall = WrapperCall;118  } else {119    llvm_unreachable("Unexpected coro_await_suspend invocation method");120  }121 122  if (CB->getCalledFunction()->getIntrinsicID() ==123      Intrinsic::coro_await_suspend_handle) {124    // Follow the lowered await_suspend call above with a lowered resume call125    // to the returned coroutine.126    if (auto *Invoke = dyn_cast<InvokeInst>(CB)) {127      // If the await_suspend call is an invoke, we continue in the next block.128      Builder.SetInsertPoint(Invoke->getNormalDest()->getFirstInsertionPt());129    }130 131    coro::LowererBase LB(*Wrapper->getParent());132    auto *ResumeAddr = LB.makeSubFnCall(NewCall, CoroSubFnInst::ResumeIndex,133                                        &*Builder.GetInsertPoint());134 135    LLVMContext &Ctx = Builder.getContext();136    FunctionType *ResumeTy = FunctionType::get(137        Type::getVoidTy(Ctx), PointerType::getUnqual(Ctx), false);138    auto *ResumeCall = Builder.CreateCall(ResumeTy, ResumeAddr, {NewCall});139    ResumeCall->setCallingConv(CallingConv::Fast);140 141    // We can't insert the 'ret' instruction and adjust the cc until the142    // function has been split, so remember this for later.143    Shape.SymmetricTransfers.push_back(ResumeCall);144 145    NewCall = ResumeCall;146  }147 148  CB->replaceAllUsesWith(NewCall);149  CB->eraseFromParent();150}151 152static void lowerAwaitSuspends(Function &F, coro::Shape &Shape) {153  IRBuilder<> Builder(F.getContext());154  for (auto *AWS : Shape.CoroAwaitSuspends)155    lowerAwaitSuspend(Builder, AWS, Shape);156}157 158static void maybeFreeRetconStorage(IRBuilder<> &Builder,159                                   const coro::Shape &Shape, Value *FramePtr,160                                   CallGraph *CG) {161  assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce);162  if (Shape.RetconLowering.IsFrameInlineInStorage)163    return;164 165  Shape.emitDealloc(Builder, FramePtr, CG);166}167 168/// Replace an llvm.coro.end.async.169/// Will inline the must tail call function call if there is one.170/// \returns true if cleanup of the coro.end block is needed, false otherwise.171static bool replaceCoroEndAsync(AnyCoroEndInst *End) {172  IRBuilder<> Builder(End);173 174  auto *EndAsync = dyn_cast<CoroAsyncEndInst>(End);175  if (!EndAsync) {176    Builder.CreateRetVoid();177    return true /*needs cleanup of coro.end block*/;178  }179 180  auto *MustTailCallFunc = EndAsync->getMustTailCallFunction();181  if (!MustTailCallFunc) {182    Builder.CreateRetVoid();183    return true /*needs cleanup of coro.end block*/;184  }185 186  // Move the must tail call from the predecessor block into the end block.187  auto *CoroEndBlock = End->getParent();188  auto *MustTailCallFuncBlock = CoroEndBlock->getSinglePredecessor();189  assert(MustTailCallFuncBlock && "Must have a single predecessor block");190  auto It = MustTailCallFuncBlock->getTerminator()->getIterator();191  auto *MustTailCall = cast<CallInst>(&*std::prev(It));192  CoroEndBlock->splice(End->getIterator(), MustTailCallFuncBlock,193                       MustTailCall->getIterator());194 195  // Insert the return instruction.196  Builder.SetInsertPoint(End);197  Builder.CreateRetVoid();198  InlineFunctionInfo FnInfo;199 200  // Remove the rest of the block, by splitting it into an unreachable block.201  auto *BB = End->getParent();202  BB->splitBasicBlock(End);203  BB->getTerminator()->eraseFromParent();204 205  auto InlineRes = InlineFunction(*MustTailCall, FnInfo);206  assert(InlineRes.isSuccess() && "Expected inlining to succeed");207  (void)InlineRes;208 209  // We have cleaned up the coro.end block above.210  return false;211}212 213/// Replace a non-unwind call to llvm.coro.end.214static void replaceFallthroughCoroEnd(AnyCoroEndInst *End,215                                      const coro::Shape &Shape, Value *FramePtr,216                                      bool InRamp, CallGraph *CG) {217  // Start inserting right before the coro.end.218  IRBuilder<> Builder(End);219 220  // Create the return instruction.221  switch (Shape.ABI) {222  // The cloned functions in switch-lowering always return void.223  case coro::ABI::Switch:224    assert(!cast<CoroEndInst>(End)->hasResults() &&225           "switch coroutine should not return any values");226    // coro.end doesn't immediately end the coroutine in the main function227    // in this lowering, because we need to deallocate the coroutine.228    if (InRamp)229      return;230    Builder.CreateRetVoid();231    break;232 233  // In async lowering this returns.234  case coro::ABI::Async: {235    bool CoroEndBlockNeedsCleanup = replaceCoroEndAsync(End);236    if (!CoroEndBlockNeedsCleanup)237      return;238    break;239  }240 241  // In unique continuation lowering, the continuations always return void.242  // But we may have implicitly allocated storage.243  case coro::ABI::RetconOnce: {244    maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);245    auto *CoroEnd = cast<CoroEndInst>(End);246    auto *RetTy = Shape.getResumeFunctionType()->getReturnType();247 248    if (!CoroEnd->hasResults()) {249      assert(RetTy->isVoidTy());250      Builder.CreateRetVoid();251      break;252    }253 254    auto *CoroResults = CoroEnd->getResults();255    unsigned NumReturns = CoroResults->numReturns();256 257    if (auto *RetStructTy = dyn_cast<StructType>(RetTy)) {258      assert(RetStructTy->getNumElements() == NumReturns &&259             "numbers of returns should match resume function singature");260      Value *ReturnValue = PoisonValue::get(RetStructTy);261      unsigned Idx = 0;262      for (Value *RetValEl : CoroResults->return_values())263        ReturnValue = Builder.CreateInsertValue(ReturnValue, RetValEl, Idx++);264      Builder.CreateRet(ReturnValue);265    } else if (NumReturns == 0) {266      assert(RetTy->isVoidTy());267      Builder.CreateRetVoid();268    } else {269      assert(NumReturns == 1);270      Builder.CreateRet(*CoroResults->retval_begin());271    }272    CoroResults->replaceAllUsesWith(273        ConstantTokenNone::get(CoroResults->getContext()));274    CoroResults->eraseFromParent();275    break;276  }277 278  // In non-unique continuation lowering, we signal completion by returning279  // a null continuation.280  case coro::ABI::Retcon: {281    assert(!cast<CoroEndInst>(End)->hasResults() &&282           "retcon coroutine should not return any values");283    maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);284    auto RetTy = Shape.getResumeFunctionType()->getReturnType();285    auto RetStructTy = dyn_cast<StructType>(RetTy);286    PointerType *ContinuationTy =287        cast<PointerType>(RetStructTy ? RetStructTy->getElementType(0) : RetTy);288 289    Value *ReturnValue = ConstantPointerNull::get(ContinuationTy);290    if (RetStructTy) {291      ReturnValue = Builder.CreateInsertValue(PoisonValue::get(RetStructTy),292                                              ReturnValue, 0);293    }294    Builder.CreateRet(ReturnValue);295    break;296  }297  }298 299  // Remove the rest of the block, by splitting it into an unreachable block.300  auto *BB = End->getParent();301  BB->splitBasicBlock(End);302  BB->getTerminator()->eraseFromParent();303}304 305// Mark a coroutine as done, which implies that the coroutine is finished and306// never gets resumed.307//308// In resume-switched ABI, the done state is represented by storing zero in309// ResumeFnAddr.310//311// NOTE: We couldn't omit the argument `FramePtr`. It is necessary because the312// pointer to the frame in splitted function is not stored in `Shape`.313static void markCoroutineAsDone(IRBuilder<> &Builder, const coro::Shape &Shape,314                                Value *FramePtr) {315  assert(316      Shape.ABI == coro::ABI::Switch &&317      "markCoroutineAsDone is only supported for Switch-Resumed ABI for now.");318  auto *GepIndex = Builder.CreateStructGEP(319      Shape.FrameTy, FramePtr, coro::Shape::SwitchFieldIndex::Resume,320      "ResumeFn.addr");321  auto *NullPtr = ConstantPointerNull::get(cast<PointerType>(322      Shape.FrameTy->getTypeAtIndex(coro::Shape::SwitchFieldIndex::Resume)));323  Builder.CreateStore(NullPtr, GepIndex);324 325  // If the coroutine don't have unwind coro end, we could omit the store to326  // the final suspend point since we could infer the coroutine is suspended327  // at the final suspend point by the nullness of ResumeFnAddr.328  // However, we can't skip it if the coroutine have unwind coro end. Since329  // the coroutine reaches unwind coro end is considered suspended at the330  // final suspend point (the ResumeFnAddr is null) but in fact the coroutine331  // didn't complete yet. We need the IndexVal for the final suspend point332  // to make the states clear.333  if (Shape.SwitchLowering.HasUnwindCoroEnd &&334      Shape.SwitchLowering.HasFinalSuspend) {335    assert(cast<CoroSuspendInst>(Shape.CoroSuspends.back())->isFinal() &&336           "The final suspend should only live in the last position of "337           "CoroSuspends.");338    ConstantInt *IndexVal = Shape.getIndex(Shape.CoroSuspends.size() - 1);339    auto *FinalIndex = Builder.CreateStructGEP(340        Shape.FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");341 342    Builder.CreateStore(IndexVal, FinalIndex);343  }344}345 346/// Replace an unwind call to llvm.coro.end.347static void replaceUnwindCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,348                                 Value *FramePtr, bool InRamp, CallGraph *CG) {349  IRBuilder<> Builder(End);350 351  switch (Shape.ABI) {352  // In switch-lowering, this does nothing in the main function.353  case coro::ABI::Switch: {354    // In C++'s specification, the coroutine should be marked as done355    // if promise.unhandled_exception() throws.  The frontend will356    // call coro.end(true) along this path.357    //358    // FIXME: We should refactor this once there is other language359    // which uses Switch-Resumed style other than C++.360    markCoroutineAsDone(Builder, Shape, FramePtr);361    if (InRamp)362      return;363    break;364  }365  // In async lowering this does nothing.366  case coro::ABI::Async:367    break;368  // In continuation-lowering, this frees the continuation storage.369  case coro::ABI::Retcon:370  case coro::ABI::RetconOnce:371    maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);372    break;373  }374 375  // If coro.end has an associated bundle, add cleanupret instruction.376  if (auto Bundle = End->getOperandBundle(LLVMContext::OB_funclet)) {377    auto *FromPad = cast<CleanupPadInst>(Bundle->Inputs[0]);378    auto *CleanupRet = Builder.CreateCleanupRet(FromPad, nullptr);379    End->getParent()->splitBasicBlock(End);380    CleanupRet->getParent()->getTerminator()->eraseFromParent();381  }382}383 384static void replaceCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,385                           Value *FramePtr, bool InRamp, CallGraph *CG) {386  if (End->isUnwind())387    replaceUnwindCoroEnd(End, Shape, FramePtr, InRamp, CG);388  else389    replaceFallthroughCoroEnd(End, Shape, FramePtr, InRamp, CG);390  End->eraseFromParent();391}392 393// In the resume function, we remove the last case  (when coro::Shape is built,394// the final suspend point (if present) is always the last element of395// CoroSuspends array) since it is an undefined behavior to resume a coroutine396// suspended at the final suspend point.397// In the destroy function, if it isn't possible that the ResumeFnAddr is NULL398// and the coroutine doesn't suspend at the final suspend point actually (this399// is possible since the coroutine is considered suspended at the final suspend400// point if promise.unhandled_exception() exits via an exception), we can401// remove the last case.402void coro::BaseCloner::handleFinalSuspend() {403  assert(Shape.ABI == coro::ABI::Switch &&404         Shape.SwitchLowering.HasFinalSuspend);405 406  if (isSwitchDestroyFunction() && Shape.SwitchLowering.HasUnwindCoroEnd)407    return;408 409  auto *Switch = cast<SwitchInst>(VMap[Shape.SwitchLowering.ResumeSwitch]);410  auto FinalCaseIt = std::prev(Switch->case_end());411  BasicBlock *ResumeBB = FinalCaseIt->getCaseSuccessor();412  Switch->removeCase(FinalCaseIt);413  if (isSwitchDestroyFunction()) {414    BasicBlock *OldSwitchBB = Switch->getParent();415    auto *NewSwitchBB = OldSwitchBB->splitBasicBlock(Switch, "Switch");416    Builder.SetInsertPoint(OldSwitchBB->getTerminator());417 418    if (NewF->isCoroOnlyDestroyWhenComplete()) {419      // When the coroutine can only be destroyed when complete, we don't need420      // to generate code for other cases.421      Builder.CreateBr(ResumeBB);422    } else {423      auto *GepIndex = Builder.CreateStructGEP(424          Shape.FrameTy, NewFramePtr, coro::Shape::SwitchFieldIndex::Resume,425          "ResumeFn.addr");426      auto *Load =427          Builder.CreateLoad(Shape.getSwitchResumePointerType(), GepIndex);428      auto *Cond = Builder.CreateIsNull(Load);429      Builder.CreateCondBr(Cond, ResumeBB, NewSwitchBB);430    }431    OldSwitchBB->getTerminator()->eraseFromParent();432  }433}434 435static FunctionType *436getFunctionTypeFromAsyncSuspend(AnyCoroSuspendInst *Suspend) {437  auto *AsyncSuspend = cast<CoroSuspendAsyncInst>(Suspend);438  auto *StructTy = cast<StructType>(AsyncSuspend->getType());439  auto &Context = Suspend->getParent()->getParent()->getContext();440  auto *VoidTy = Type::getVoidTy(Context);441  return FunctionType::get(VoidTy, StructTy->elements(), false);442}443 444static Function *createCloneDeclaration(Function &OrigF, coro::Shape &Shape,445                                        const Twine &Suffix,446                                        Module::iterator InsertBefore,447                                        AnyCoroSuspendInst *ActiveSuspend) {448  Module *M = OrigF.getParent();449  auto *FnTy = (Shape.ABI != coro::ABI::Async)450                   ? Shape.getResumeFunctionType()451                   : getFunctionTypeFromAsyncSuspend(ActiveSuspend);452 453  Function *NewF =454      Function::Create(FnTy, GlobalValue::LinkageTypes::InternalLinkage,455                       OrigF.getName() + Suffix);456 457  M->getFunctionList().insert(InsertBefore, NewF);458 459  return NewF;460}461 462/// Replace uses of the active llvm.coro.suspend.retcon/async call with the463/// arguments to the continuation function.464///465/// This assumes that the builder has a meaningful insertion point.466void coro::BaseCloner::replaceRetconOrAsyncSuspendUses() {467  assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce ||468         Shape.ABI == coro::ABI::Async);469 470  auto NewS = VMap[ActiveSuspend];471  if (NewS->use_empty())472    return;473 474  // Copy out all the continuation arguments after the buffer pointer into475  // an easily-indexed data structure for convenience.476  SmallVector<Value *, 8> Args;477  // The async ABI includes all arguments -- including the first argument.478  bool IsAsyncABI = Shape.ABI == coro::ABI::Async;479  for (auto I = IsAsyncABI ? NewF->arg_begin() : std::next(NewF->arg_begin()),480            E = NewF->arg_end();481       I != E; ++I)482    Args.push_back(&*I);483 484  // If the suspend returns a single scalar value, we can just do a simple485  // replacement.486  if (!isa<StructType>(NewS->getType())) {487    assert(Args.size() == 1);488    NewS->replaceAllUsesWith(Args.front());489    return;490  }491 492  // Try to peephole extracts of an aggregate return.493  for (Use &U : llvm::make_early_inc_range(NewS->uses())) {494    auto *EVI = dyn_cast<ExtractValueInst>(U.getUser());495    if (!EVI || EVI->getNumIndices() != 1)496      continue;497 498    EVI->replaceAllUsesWith(Args[EVI->getIndices().front()]);499    EVI->eraseFromParent();500  }501 502  // If we have no remaining uses, we're done.503  if (NewS->use_empty())504    return;505 506  // Otherwise, we need to create an aggregate.507  Value *Aggr = PoisonValue::get(NewS->getType());508  for (auto [Idx, Arg] : llvm::enumerate(Args))509    Aggr = Builder.CreateInsertValue(Aggr, Arg, Idx);510 511  NewS->replaceAllUsesWith(Aggr);512}513 514void coro::BaseCloner::replaceCoroSuspends() {515  Value *SuspendResult;516 517  switch (Shape.ABI) {518  // In switch lowering, replace coro.suspend with the appropriate value519  // for the type of function we're extracting.520  // Replacing coro.suspend with (0) will result in control flow proceeding to521  // a resume label associated with a suspend point, replacing it with (1) will522  // result in control flow proceeding to a cleanup label associated with this523  // suspend point.524  case coro::ABI::Switch:525    SuspendResult = Builder.getInt8(isSwitchDestroyFunction() ? 1 : 0);526    break;527 528  // In async lowering there are no uses of the result.529  case coro::ABI::Async:530    return;531 532  // In returned-continuation lowering, the arguments from earlier533  // continuations are theoretically arbitrary, and they should have been534  // spilled.535  case coro::ABI::RetconOnce:536  case coro::ABI::Retcon:537    return;538  }539 540  for (AnyCoroSuspendInst *CS : Shape.CoroSuspends) {541    // The active suspend was handled earlier.542    if (CS == ActiveSuspend)543      continue;544 545    auto *MappedCS = cast<AnyCoroSuspendInst>(VMap[CS]);546    MappedCS->replaceAllUsesWith(SuspendResult);547    MappedCS->eraseFromParent();548  }549}550 551void coro::BaseCloner::replaceCoroEnds() {552  for (AnyCoroEndInst *CE : Shape.CoroEnds) {553    // We use a null call graph because there's no call graph node for554    // the cloned function yet.  We'll just be rebuilding that later.555    auto *NewCE = cast<AnyCoroEndInst>(VMap[CE]);556    replaceCoroEnd(NewCE, Shape, NewFramePtr, /*in ramp*/ false, nullptr);557  }558}559 560void coro::BaseCloner::replaceCoroIsInRamp() {561  auto &Ctx = OrigF.getContext();562  for (auto *II : Shape.CoroIsInRampInsts) {563    auto *NewII = cast<CoroIsInRampInst>(VMap[II]);564    NewII->replaceAllUsesWith(ConstantInt::getFalse(Ctx));565    NewII->eraseFromParent();566  }567}568 569static void replaceSwiftErrorOps(Function &F, coro::Shape &Shape,570                                 ValueToValueMapTy *VMap) {571  if (Shape.ABI == coro::ABI::Async && Shape.CoroSuspends.empty())572    return;573  Value *CachedSlot = nullptr;574  auto getSwiftErrorSlot = [&](Type *ValueTy) -> Value * {575    if (CachedSlot)576      return CachedSlot;577 578    // Check if the function has a swifterror argument.579    for (auto &Arg : F.args()) {580      if (Arg.isSwiftError()) {581        CachedSlot = &Arg;582        return &Arg;583      }584    }585 586    // Create a swifterror alloca.587    IRBuilder<> Builder(&F.getEntryBlock(),588                        F.getEntryBlock().getFirstNonPHIOrDbg());589    auto Alloca = Builder.CreateAlloca(ValueTy);590    Alloca->setSwiftError(true);591 592    CachedSlot = Alloca;593    return Alloca;594  };595 596  for (CallInst *Op : Shape.SwiftErrorOps) {597    auto MappedOp = VMap ? cast<CallInst>((*VMap)[Op]) : Op;598    IRBuilder<> Builder(MappedOp);599 600    // If there are no arguments, this is a 'get' operation.601    Value *MappedResult;602    if (Op->arg_empty()) {603      auto ValueTy = Op->getType();604      auto Slot = getSwiftErrorSlot(ValueTy);605      MappedResult = Builder.CreateLoad(ValueTy, Slot);606    } else {607      assert(Op->arg_size() == 1);608      auto Value = MappedOp->getArgOperand(0);609      auto ValueTy = Value->getType();610      auto Slot = getSwiftErrorSlot(ValueTy);611      Builder.CreateStore(Value, Slot);612      MappedResult = Slot;613    }614 615    MappedOp->replaceAllUsesWith(MappedResult);616    MappedOp->eraseFromParent();617  }618 619  // If we're updating the original function, we've invalidated SwiftErrorOps.620  if (VMap == nullptr) {621    Shape.SwiftErrorOps.clear();622  }623}624 625/// Returns all debug records in F.626static SmallVector<DbgVariableRecord *>627collectDbgVariableRecords(Function &F) {628  SmallVector<DbgVariableRecord *> DbgVariableRecords;629  for (auto &I : instructions(F)) {630    for (DbgVariableRecord &DVR : filterDbgVars(I.getDbgRecordRange()))631      DbgVariableRecords.push_back(&DVR);632  }633  return DbgVariableRecords;634}635 636void coro::BaseCloner::replaceSwiftErrorOps() {637  ::replaceSwiftErrorOps(*NewF, Shape, &VMap);638}639 640void coro::BaseCloner::salvageDebugInfo() {641  auto DbgVariableRecords = collectDbgVariableRecords(*NewF);642  SmallDenseMap<Argument *, AllocaInst *, 4> ArgToAllocaMap;643 644  // Only 64-bit ABIs have a register we can refer to with the entry value.645  bool UseEntryValue = OrigF.getParent()->getTargetTriple().isArch64Bit();646  for (DbgVariableRecord *DVR : DbgVariableRecords)647    coro::salvageDebugInfo(ArgToAllocaMap, *DVR, UseEntryValue);648 649  // Remove all salvaged dbg.declare intrinsics that became650  // either unreachable or stale due to the CoroSplit transformation.651  DominatorTree DomTree(*NewF);652  auto IsUnreachableBlock = [&](BasicBlock *BB) {653    return !isPotentiallyReachable(&NewF->getEntryBlock(), BB, nullptr,654                                   &DomTree);655  };656  auto RemoveOne = [&](DbgVariableRecord *DVI) {657    if (IsUnreachableBlock(DVI->getParent()))658      DVI->eraseFromParent();659    else if (isa_and_nonnull<AllocaInst>(DVI->getVariableLocationOp(0))) {660      // Count all non-debuginfo uses in reachable blocks.661      unsigned Uses = 0;662      for (auto *User : DVI->getVariableLocationOp(0)->users())663        if (auto *I = dyn_cast<Instruction>(User))664          if (!isa<AllocaInst>(I) && !IsUnreachableBlock(I->getParent()))665            ++Uses;666      if (!Uses)667        DVI->eraseFromParent();668    }669  };670  for_each(DbgVariableRecords, RemoveOne);671}672 673void coro::BaseCloner::replaceEntryBlock() {674  // In the original function, the AllocaSpillBlock is a block immediately675  // following the allocation of the frame object which defines GEPs for676  // all the allocas that have been moved into the frame, and it ends by677  // branching to the original beginning of the coroutine.  Make this678  // the entry block of the cloned function.679  auto *Entry = cast<BasicBlock>(VMap[Shape.AllocaSpillBlock]);680  auto *OldEntry = &NewF->getEntryBlock();681  Entry->setName("entry" + Suffix);682  Entry->moveBefore(OldEntry);683  Entry->getTerminator()->eraseFromParent();684 685  // Clear all predecessors of the new entry block.  There should be686  // exactly one predecessor, which we created when splitting out687  // AllocaSpillBlock to begin with.688  assert(Entry->hasOneUse());689  auto BranchToEntry = cast<BranchInst>(Entry->user_back());690  assert(BranchToEntry->isUnconditional());691  Builder.SetInsertPoint(BranchToEntry);692  Builder.CreateUnreachable();693  BranchToEntry->eraseFromParent();694 695  // Branch from the entry to the appropriate place.696  Builder.SetInsertPoint(Entry);697  switch (Shape.ABI) {698  case coro::ABI::Switch: {699    // In switch-lowering, we built a resume-entry block in the original700    // function.  Make the entry block branch to this.701    auto *SwitchBB =702        cast<BasicBlock>(VMap[Shape.SwitchLowering.ResumeEntryBlock]);703    Builder.CreateBr(SwitchBB);704    SwitchBB->moveAfter(Entry);705    break;706  }707  case coro::ABI::Async:708  case coro::ABI::Retcon:709  case coro::ABI::RetconOnce: {710    // In continuation ABIs, we want to branch to immediately after the711    // active suspend point.  Earlier phases will have put the suspend in its712    // own basic block, so just thread our jump directly to its successor.713    assert((Shape.ABI == coro::ABI::Async &&714            isa<CoroSuspendAsyncInst>(ActiveSuspend)) ||715           ((Shape.ABI == coro::ABI::Retcon ||716             Shape.ABI == coro::ABI::RetconOnce) &&717            isa<CoroSuspendRetconInst>(ActiveSuspend)));718    auto *MappedCS = cast<AnyCoroSuspendInst>(VMap[ActiveSuspend]);719    auto Branch = cast<BranchInst>(MappedCS->getNextNode());720    assert(Branch->isUnconditional());721    Builder.CreateBr(Branch->getSuccessor(0));722    break;723  }724  }725 726  // Any static alloca that's still being used but not reachable from the new727  // entry needs to be moved to the new entry.728  Function *F = OldEntry->getParent();729  DominatorTree DT{*F};730  for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {731    auto *Alloca = dyn_cast<AllocaInst>(&I);732    if (!Alloca || I.use_empty())733      continue;734    if (DT.isReachableFromEntry(I.getParent()) ||735        !isa<ConstantInt>(Alloca->getArraySize()))736      continue;737    I.moveBefore(*Entry, Entry->getFirstInsertionPt());738  }739}740 741/// Derive the value of the new frame pointer.742Value *coro::BaseCloner::deriveNewFramePointer() {743  // Builder should be inserting to the front of the new entry block.744 745  switch (Shape.ABI) {746  // In switch-lowering, the argument is the frame pointer.747  case coro::ABI::Switch:748    return &*NewF->arg_begin();749  // In async-lowering, one of the arguments is an async context as determined750  // by the `llvm.coro.id.async` intrinsic. We can retrieve the async context of751  // the resume function from the async context projection function associated752  // with the active suspend. The frame is located as a tail to the async753  // context header.754  case coro::ABI::Async: {755    auto *ActiveAsyncSuspend = cast<CoroSuspendAsyncInst>(ActiveSuspend);756    auto ContextIdx = ActiveAsyncSuspend->getStorageArgumentIndex() & 0xff;757    auto *CalleeContext = NewF->getArg(ContextIdx);758    auto *ProjectionFunc =759        ActiveAsyncSuspend->getAsyncContextProjectionFunction();760    auto DbgLoc =761        cast<CoroSuspendAsyncInst>(VMap[ActiveSuspend])->getDebugLoc();762    // Calling i8* (i8*)763    auto *CallerContext = Builder.CreateCall(ProjectionFunc->getFunctionType(),764                                             ProjectionFunc, CalleeContext);765    CallerContext->setCallingConv(ProjectionFunc->getCallingConv());766    CallerContext->setDebugLoc(DbgLoc);767    // The frame is located after the async_context header.768    auto &Context = Builder.getContext();769    auto *FramePtrAddr = Builder.CreateConstInBoundsGEP1_32(770        Type::getInt8Ty(Context), CallerContext,771        Shape.AsyncLowering.FrameOffset, "async.ctx.frameptr");772    // Inline the projection function.773    InlineFunctionInfo InlineInfo;774    auto InlineRes = InlineFunction(*CallerContext, InlineInfo);775    assert(InlineRes.isSuccess());776    (void)InlineRes;777    return FramePtrAddr;778  }779  // In continuation-lowering, the argument is the opaque storage.780  case coro::ABI::Retcon:781  case coro::ABI::RetconOnce: {782    Argument *NewStorage = &*NewF->arg_begin();783    auto FramePtrTy = PointerType::getUnqual(Shape.FrameTy->getContext());784 785    // If the storage is inline, just bitcast to the storage to the frame type.786    if (Shape.RetconLowering.IsFrameInlineInStorage)787      return NewStorage;788 789    // Otherwise, load the real frame from the opaque storage.790    return Builder.CreateLoad(FramePtrTy, NewStorage);791  }792  }793  llvm_unreachable("bad ABI");794}795 796/// Adjust the scope line of the funclet to the first line number after the797/// suspend point. This avoids a jump in the line table from the function798/// declaration (where prologue instructions are attributed to) to the suspend799/// point.800/// Only adjust the scope line when the files are the same.801/// If no candidate line number is found, fallback to the line of ActiveSuspend.802static void updateScopeLine(Instruction *ActiveSuspend,803                            DISubprogram &SPToUpdate) {804  if (!ActiveSuspend)805    return;806 807  // No subsequent instruction -> fallback to the location of ActiveSuspend.808  if (!ActiveSuspend->getNextNode()) {809    if (auto DL = ActiveSuspend->getDebugLoc())810      if (SPToUpdate.getFile() == DL->getFile())811        SPToUpdate.setScopeLine(DL->getLine());812    return;813  }814 815  BasicBlock::iterator Successor = ActiveSuspend->getNextNode()->getIterator();816  // Corosplit splits the BB around ActiveSuspend, so the meaningful817  // instructions are not in the same BB.818  if (auto *Branch = dyn_cast_or_null<BranchInst>(Successor);819      Branch && Branch->isUnconditional())820    Successor = Branch->getSuccessor(0)->getFirstNonPHIOrDbg();821 822  // Find the first successor of ActiveSuspend with a non-zero line location.823  // If that matches the file of ActiveSuspend, use it.824  BasicBlock *PBB = Successor->getParent();825  for (; Successor != PBB->end(); Successor = std::next(Successor)) {826    Successor = skipDebugIntrinsics(Successor);827    auto DL = Successor->getDebugLoc();828    if (!DL || DL.getLine() == 0)829      continue;830 831    if (SPToUpdate.getFile() == DL->getFile()) {832      SPToUpdate.setScopeLine(DL.getLine());833      return;834    }835 836    break;837  }838 839  // If the search above failed, fallback to the location of ActiveSuspend.840  if (auto DL = ActiveSuspend->getDebugLoc())841    if (SPToUpdate.getFile() == DL->getFile())842      SPToUpdate.setScopeLine(DL->getLine());843}844 845static void addFramePointerAttrs(AttributeList &Attrs, LLVMContext &Context,846                                 unsigned ParamIndex, uint64_t Size,847                                 Align Alignment, bool NoAlias) {848  AttrBuilder ParamAttrs(Context);849  ParamAttrs.addAttribute(Attribute::NonNull);850  ParamAttrs.addAttribute(Attribute::NoUndef);851 852  if (NoAlias)853    ParamAttrs.addAttribute(Attribute::NoAlias);854 855  ParamAttrs.addAlignmentAttr(Alignment);856  ParamAttrs.addDereferenceableAttr(Size);857  Attrs = Attrs.addParamAttributes(Context, ParamIndex, ParamAttrs);858}859 860static void addAsyncContextAttrs(AttributeList &Attrs, LLVMContext &Context,861                                 unsigned ParamIndex) {862  AttrBuilder ParamAttrs(Context);863  ParamAttrs.addAttribute(Attribute::SwiftAsync);864  Attrs = Attrs.addParamAttributes(Context, ParamIndex, ParamAttrs);865}866 867static void addSwiftSelfAttrs(AttributeList &Attrs, LLVMContext &Context,868                              unsigned ParamIndex) {869  AttrBuilder ParamAttrs(Context);870  ParamAttrs.addAttribute(Attribute::SwiftSelf);871  Attrs = Attrs.addParamAttributes(Context, ParamIndex, ParamAttrs);872}873 874/// Rebuild an aggregate/expression constant from already-remapped operands.875static Constant *rebuildConstant(Constant *C, ArrayRef<Constant *> NewOps) {876  if (auto *CE = dyn_cast<ConstantExpr>(C))877    return CE->getWithOperands(NewOps);878  if (auto *CA = dyn_cast<ConstantArray>(C))879    return ConstantArray::get(CA->getType(), NewOps);880  if (auto *CS = dyn_cast<ConstantStruct>(C))881    return ConstantStruct::get(CS->getType(), NewOps);882  if (isa<ConstantVector>(C))883    return ConstantVector::get(NewOps);884  llvm_unreachable("unexpected constant kind in blockaddress-table remap");885}886 887/// Recursively rebuild \p C, replacing every `blockaddress(&OrigF, BB)` with888/// `blockaddress(&NewF, VMap[BB])` for blocks that were cloned into \p NewF.889/// Returns \p C unchanged if it contains no such blockaddress.890static Constant *891remapBlockAddressesInConstant(Constant *C, Function &OrigF, Function &NewF,892                              ValueToValueMapTy &VMap,893                              DenseMap<Constant *, Constant *> &Cache) {894  auto It = Cache.find(C);895  if (It != Cache.end())896    return It->second;897 898  Constant *Result = C;899  if (auto *BA = dyn_cast<BlockAddress>(C)) {900    if (BA->getFunction() == &OrigF) {901      auto *MappedBB =902          dyn_cast_or_null<BasicBlock>(VMap.lookup(BA->getBasicBlock()));903      if (MappedBB && MappedBB->getParent() == &NewF)904        Result = BlockAddress::get(&NewF, MappedBB);905    }906  } else if (isa<ConstantAggregate>(C) || isa<ConstantExpr>(C)) {907    SmallVector<Constant *, 16> NewOps;908    bool Changed = false;909    for (Use &Op : C->operands()) {910      auto *OpC = cast<Constant>(Op.get());911      Constant *NewOp =912          remapBlockAddressesInConstant(OpC, OrigF, NewF, VMap, Cache);913      Changed |= NewOp != OpC;914      NewOps.push_back(NewOp);915    }916    if (Changed)917      Result = rebuildConstant(C, NewOps);918  }919 920  Cache[C] = Result;921  return Result;922}923 924/// Recursively rebuild \p C, replacing any global in \p Map with its925/// replacement. Returns \p C unchanged if it references no mapped global.926static Constant *927replaceGlobalsInConstant(Constant *C,928                         const DenseMap<GlobalVariable *, GlobalVariable *> &Map,929                         DenseMap<Constant *, Constant *> &Cache) {930  if (auto *GV = dyn_cast<GlobalVariable>(C)) {931    auto It = Map.find(GV);932    return It != Map.end() ? It->second : C;933  }934 935  auto CIt = Cache.find(C);936  if (CIt != Cache.end())937    return CIt->second;938 939  Constant *Result = C;940  if (isa<ConstantAggregate>(C) || isa<ConstantExpr>(C)) {941    SmallVector<Constant *, 16> NewOps;942    bool Changed = false;943    for (Use &Op : C->operands()) {944      auto *OpC = cast<Constant>(Op.get());945      Constant *NewOp = replaceGlobalsInConstant(OpC, Map, Cache);946      Changed |= NewOp != OpC;947      NewOps.push_back(NewOp);948    }949    if (Changed)950      Result = rebuildConstant(C, NewOps);951  }952 953  Cache[C] = Result;954  return Result;955}956 957/// Relocate GNU computed-goto jump tables across a coroutine split.958///959/// When a coroutine takes the address of its own basic blocks -- e.g. glibc's960/// printf_positional drives format dispatch with a module-level constant array961/// of `blockaddress` values feeding an `indirectbr` -- CoroSplit clones the962/// address-taken blocks into the resume/destroy/cleanup clone, but the jump963/// table global still refers to the *ramp* function's blocks. When964/// postSplitCleanup (or later CFG simplification) deletes those now-unreachable965/// ramp blocks, BasicBlock::~BasicBlock RAUWs every `blockaddress` of a deleted966/// block to the `inttoptr(i32 1)` sentinel, poisoning the shared table. The967/// clone's `indirectbr` then reads the poisoned table, which the WebAssembly968/// backend lowers to `unreachable` -> a runtime trap on the resume path.969///970/// Fix: give each clone its own copy of any jump table that references this971/// coroutine's blocks, with the `blockaddress` constants remapped to the972/// clone's blocks, and rewrite the clone's references to use it. The ramp keeps973/// the original table (correct for any address-taken blocks it retains;974/// harmlessly dead if it retains none). This must run while OrigF's blocks975/// still exist (i.e. before postSplitCleanup), which holds for all callers.976static void relocateBlockAddressTablesForClone(Function &OrigF, Function &NewF,977                                               ValueToValueMapTy &VMap) {978  // Collect the GlobalVariables transitively referenced by NewF's body.979  SmallPtrSet<GlobalVariable *, 4> Referenced;980  SmallPtrSet<Constant *, 16> Visited;981  SmallVector<Constant *, 16> Worklist;982  for (Instruction &I : instructions(NewF))983    for (Use &Op : I.operands())984      if (auto *C = dyn_cast<Constant>(Op.get()))985        Worklist.push_back(C);986  while (!Worklist.empty()) {987    Constant *C = Worklist.pop_back_val();988    if (!Visited.insert(C).second)989      continue;990    if (auto *GV = dyn_cast<GlobalVariable>(C)) {991      Referenced.insert(GV);992      continue;993    }994    for (Use &Op : C->operands())995      if (auto *OpC = dyn_cast<Constant>(Op.get()))996        Worklist.push_back(OpC);997  }998 999  // Build a per-clone copy of each referenced table whose initializer holds a1000  // blockaddress of one of OrigF's (now cloned) blocks.1001  DenseMap<GlobalVariable *, GlobalVariable *> TableClones;1002  DenseMap<Constant *, Constant *> RemapCache;1003  for (GlobalVariable *GV : Referenced) {1004    if (!GV->hasInitializer())1005      continue;1006    Constant *Init = GV->getInitializer();1007    Constant *NewInit =1008        remapBlockAddressesInConstant(Init, OrigF, NewF, VMap, RemapCache);1009    if (NewInit == Init)1010      continue; // No blockaddress of this coroutine inside.1011    auto *NewGV = new GlobalVariable(1012        *GV->getParent(), GV->getValueType(), GV->isConstant(),1013        GlobalValue::PrivateLinkage, NewInit, GV->getName() + "." + NewF.getName(),1014        /*InsertBefore=*/nullptr, GV->getThreadLocalMode(),1015        GV->getType()->getAddressSpace());1016    NewGV->setAlignment(GV->getAlign());1017    NewGV->setUnnamedAddr(GV->getUnnamedAddr());1018    TableClones[GV] = NewGV;1019  }1020  if (TableClones.empty())1021    return;1022 1023  // Rewrite NewF's references to the original tables to use the clone's copies.1024  DenseMap<Constant *, Constant *> ReplaceCache;1025  for (Instruction &I : instructions(NewF))1026    for (Use &Op : I.operands())1027      if (auto *C = dyn_cast<Constant>(Op.get())) {1028        Constant *NewC = replaceGlobalsInConstant(C, TableClones, ReplaceCache);1029        if (NewC != C)1030          Op.set(NewC);1031      }1032}1033 1034/// Clone the body of the original function into a resume function of1035/// some sort.1036void coro::BaseCloner::create() {1037  assert(NewF);1038 1039  // Replace all args with dummy instructions. If an argument is the old frame1040  // pointer, the dummy will be replaced by the new frame pointer once it is1041  // computed below. Uses of all other arguments should have already been1042  // rewritten by buildCoroutineFrame() to use loads/stores on the coroutine1043  // frame.1044  SmallVector<Instruction *> DummyArgs;1045  for (Argument &A : OrigF.args()) {1046    DummyArgs.push_back(new FreezeInst(PoisonValue::get(A.getType())));1047    VMap[&A] = DummyArgs.back();1048  }1049 1050  SmallVector<ReturnInst *, 4> Returns;1051 1052  // Ignore attempts to change certain attributes of the function.1053  // TODO: maybe there should be a way to suppress this during cloning?1054  auto savedVisibility = NewF->getVisibility();1055  auto savedUnnamedAddr = NewF->getUnnamedAddr();1056  auto savedDLLStorageClass = NewF->getDLLStorageClass();1057 1058  // NewF's linkage (which CloneFunctionInto does *not* change) might not1059  // be compatible with the visibility of OrigF (which it *does* change),1060  // so protect against that.1061  auto savedLinkage = NewF->getLinkage();1062  NewF->setLinkage(llvm::GlobalValue::ExternalLinkage);1063 1064  CloneFunctionInto(NewF, &OrigF, VMap,1065                    CloneFunctionChangeType::LocalChangesOnly, Returns);1066 1067  // If this coroutine takes the address of its own blocks (a GNU computed-goto1068  // jump table driving an indirectbr), relocate those tables into per-clone1069  // copies so deleting the ramp's address-taken blocks does not poison the1070  // clone's table with the inttoptr(i32 1) sentinel. Must run while OrigF's1071  // blocks still exist, which holds here (postSplitCleanup runs later).1072  relocateBlockAddressTablesForClone(OrigF, *NewF, VMap);1073 1074  auto &Context = NewF->getContext();1075 1076  if (DISubprogram *SP = NewF->getSubprogram()) {1077    assert(SP != OrigF.getSubprogram() && SP->isDistinct());1078    updateScopeLine(ActiveSuspend, *SP);1079 1080    // Update the linkage name and the function name to reflect the modified1081    // name.1082    MDString *NewLinkageName = MDString::get(Context, NewF->getName());1083    SP->replaceLinkageName(NewLinkageName);1084    if (DISubprogram *Decl = SP->getDeclaration()) {1085      TempDISubprogram NewDecl = Decl->clone();1086      NewDecl->replaceLinkageName(NewLinkageName);1087      SP->replaceDeclaration(MDNode::replaceWithUniqued(std::move(NewDecl)));1088    }1089  }1090 1091  NewF->setLinkage(savedLinkage);1092  NewF->setVisibility(savedVisibility);1093  NewF->setUnnamedAddr(savedUnnamedAddr);1094  NewF->setDLLStorageClass(savedDLLStorageClass);1095  // The function sanitizer metadata needs to match the signature of the1096  // function it is being attached to. However this does not hold for split1097  // functions here. Thus remove the metadata for split functions.1098  if (Shape.ABI == coro::ABI::Switch &&1099      NewF->hasMetadata(LLVMContext::MD_func_sanitize))1100    NewF->eraseMetadata(LLVMContext::MD_func_sanitize);1101 1102  // Replace the attributes of the new function:1103  auto OrigAttrs = NewF->getAttributes();1104  auto NewAttrs = AttributeList();1105 1106  switch (Shape.ABI) {1107  case coro::ABI::Switch:1108    // Bootstrap attributes by copying function attributes from the1109    // original function.  This should include optimization settings and so on.1110    NewAttrs = NewAttrs.addFnAttributes(1111        Context, AttrBuilder(Context, OrigAttrs.getFnAttrs()));1112 1113    addFramePointerAttrs(NewAttrs, Context, 0, Shape.FrameSize,1114                         Shape.FrameAlign, /*NoAlias=*/false);1115    break;1116  case coro::ABI::Async: {1117    auto *ActiveAsyncSuspend = cast<CoroSuspendAsyncInst>(ActiveSuspend);1118    if (OrigF.hasParamAttribute(Shape.AsyncLowering.ContextArgNo,1119                                Attribute::SwiftAsync)) {1120      uint32_t ArgAttributeIndices =1121          ActiveAsyncSuspend->getStorageArgumentIndex();1122      auto ContextArgIndex = ArgAttributeIndices & 0xff;1123      addAsyncContextAttrs(NewAttrs, Context, ContextArgIndex);1124 1125      // `swiftasync` must preceed `swiftself` so 0 is not a valid index for1126      // `swiftself`.1127      auto SwiftSelfIndex = ArgAttributeIndices >> 8;1128      if (SwiftSelfIndex)1129        addSwiftSelfAttrs(NewAttrs, Context, SwiftSelfIndex);1130    }1131 1132    // Transfer the original function's attributes.1133    auto FnAttrs = OrigF.getAttributes().getFnAttrs();1134    NewAttrs = NewAttrs.addFnAttributes(Context, AttrBuilder(Context, FnAttrs));1135    break;1136  }1137  case coro::ABI::Retcon:1138  case coro::ABI::RetconOnce:1139    // If we have a continuation prototype, just use its attributes,1140    // full-stop.1141    NewAttrs = Shape.RetconLowering.ResumePrototype->getAttributes();1142 1143    /// FIXME: Is it really good to add the NoAlias attribute?1144    addFramePointerAttrs(NewAttrs, Context, 0,1145                         Shape.getRetconCoroId()->getStorageSize(),1146                         Shape.getRetconCoroId()->getStorageAlignment(),1147                         /*NoAlias=*/true);1148 1149    break;1150  }1151 1152  switch (Shape.ABI) {1153  // In these ABIs, the cloned functions always return 'void', and the1154  // existing return sites are meaningless.  Note that for unique1155  // continuations, this includes the returns associated with suspends;1156  // this is fine because we can't suspend twice.1157  case coro::ABI::Switch:1158  case coro::ABI::RetconOnce:1159    // Remove old returns.1160    for (ReturnInst *Return : Returns)1161      changeToUnreachable(Return);1162    break;1163 1164  // With multi-suspend continuations, we'll already have eliminated the1165  // original returns and inserted returns before all the suspend points,1166  // so we want to leave any returns in place.1167  case coro::ABI::Retcon:1168    break;1169  // Async lowering will insert musttail call functions at all suspend points1170  // followed by a return.1171  // Don't change returns to unreachable because that will trip up the verifier.1172  // These returns should be unreachable from the clone.1173  case coro::ABI::Async:1174    break;1175  }1176 1177  NewF->setAttributes(NewAttrs);1178  NewF->setCallingConv(Shape.getResumeFunctionCC());1179 1180  // Set up the new entry block.1181  replaceEntryBlock();1182 1183  // Turn symmetric transfers into musttail calls.1184  for (CallInst *ResumeCall : Shape.SymmetricTransfers) {1185    ResumeCall = cast<CallInst>(VMap[ResumeCall]);1186    if (TTI.supportsTailCallFor(ResumeCall)) {1187      // FIXME: Could we support symmetric transfer effectively without1188      // musttail?1189      ResumeCall->setTailCallKind(CallInst::TCK_MustTail);1190    }1191 1192    // Put a 'ret void' after the call, and split any remaining instructions to1193    // an unreachable block.1194    BasicBlock *BB = ResumeCall->getParent();1195    BB->splitBasicBlock(ResumeCall->getNextNode());1196    Builder.SetInsertPoint(BB->getTerminator());1197    Builder.CreateRetVoid();1198    BB->getTerminator()->eraseFromParent();1199  }1200 1201  Builder.SetInsertPoint(&NewF->getEntryBlock().front());1202  NewFramePtr = deriveNewFramePointer();1203 1204  // Remap frame pointer.1205  Value *OldFramePtr = VMap[Shape.FramePtr];1206  NewFramePtr->takeName(OldFramePtr);1207  OldFramePtr->replaceAllUsesWith(NewFramePtr);1208 1209  // Remap vFrame pointer.1210  auto *NewVFrame = Builder.CreateBitCast(1211      NewFramePtr, PointerType::getUnqual(Builder.getContext()), "vFrame");1212  Value *OldVFrame = cast<Value>(VMap[Shape.CoroBegin]);1213  if (OldVFrame != NewVFrame)1214    OldVFrame->replaceAllUsesWith(NewVFrame);1215 1216  // All uses of the arguments should have been resolved by this point,1217  // so we can safely remove the dummy values.1218  for (Instruction *DummyArg : DummyArgs) {1219    DummyArg->replaceAllUsesWith(PoisonValue::get(DummyArg->getType()));1220    DummyArg->deleteValue();1221  }1222 1223  switch (Shape.ABI) {1224  case coro::ABI::Switch:1225    // Rewrite final suspend handling as it is not done via switch (allows to1226    // remove final case from the switch, since it is undefined behavior to1227    // resume the coroutine suspended at the final suspend point.1228    if (Shape.SwitchLowering.HasFinalSuspend)1229      handleFinalSuspend();1230    break;1231  case coro::ABI::Async:1232  case coro::ABI::Retcon:1233  case coro::ABI::RetconOnce:1234    // Replace uses of the active suspend with the corresponding1235    // continuation-function arguments.1236    assert(ActiveSuspend != nullptr &&1237           "no active suspend when lowering a continuation-style coroutine");1238    replaceRetconOrAsyncSuspendUses();1239    break;1240  }1241 1242  // Handle suspends.1243  replaceCoroSuspends();1244 1245  // Handle swifterror.1246  replaceSwiftErrorOps();1247 1248  // Remove coro.end intrinsics.1249  replaceCoroEnds();1250 1251  replaceCoroIsInRamp();1252 1253  // Salvage debug info that points into the coroutine frame.1254  salvageDebugInfo();1255}1256 1257void coro::SwitchCloner::create() {1258  // Create a new function matching the original type1259  NewF = createCloneDeclaration(OrigF, Shape, Suffix, OrigF.getParent()->end(),1260                                ActiveSuspend);1261 1262  // Clone the function1263  coro::BaseCloner::create();1264 1265  // Eliminate coro.free from the clones, replacing it with 'null' in cleanup,1266  // to suppress deallocation code.1267  coro::replaceCoroFree(cast<CoroIdInst>(VMap[Shape.CoroBegin->getId()]),1268                        /*Elide=*/FKind == coro::CloneKind::SwitchCleanup);1269}1270 1271static void updateAsyncFuncPointerContextSize(coro::Shape &Shape) {1272  assert(Shape.ABI == coro::ABI::Async);1273 1274  auto *FuncPtrStruct = cast<ConstantStruct>(1275      Shape.AsyncLowering.AsyncFuncPointer->getInitializer());1276  auto *OrigRelativeFunOffset = FuncPtrStruct->getOperand(0);1277  auto *OrigContextSize = FuncPtrStruct->getOperand(1);1278  auto *NewContextSize = ConstantInt::get(OrigContextSize->getType(),1279                                          Shape.AsyncLowering.ContextSize);1280  auto *NewFuncPtrStruct = ConstantStruct::get(1281      FuncPtrStruct->getType(), OrigRelativeFunOffset, NewContextSize);1282 1283  Shape.AsyncLowering.AsyncFuncPointer->setInitializer(NewFuncPtrStruct);1284}1285 1286static TypeSize getFrameSizeForShape(coro::Shape &Shape) {1287  // In the same function all coro.sizes should have the same result type.1288  auto *SizeIntrin = Shape.CoroSizes.back();1289  Module *M = SizeIntrin->getModule();1290  const DataLayout &DL = M->getDataLayout();1291  return DL.getTypeAllocSize(Shape.FrameTy);1292}1293 1294static void replaceFrameSizeAndAlignment(coro::Shape &Shape) {1295  if (Shape.ABI == coro::ABI::Async)1296    updateAsyncFuncPointerContextSize(Shape);1297 1298  for (CoroAlignInst *CA : Shape.CoroAligns) {1299    CA->replaceAllUsesWith(1300        ConstantInt::get(CA->getType(), Shape.FrameAlign.value()));1301    CA->eraseFromParent();1302  }1303 1304  if (Shape.CoroSizes.empty())1305    return;1306 1307  // In the same function all coro.sizes should have the same result type.1308  auto *SizeIntrin = Shape.CoroSizes.back();1309  auto *SizeConstant =1310      ConstantInt::get(SizeIntrin->getType(), getFrameSizeForShape(Shape));1311 1312  for (CoroSizeInst *CS : Shape.CoroSizes) {1313    CS->replaceAllUsesWith(SizeConstant);1314    CS->eraseFromParent();1315  }1316}1317 1318static void postSplitCleanup(Function &F) {1319  removeUnreachableBlocks(F);1320 1321#ifndef NDEBUG1322  // For now, we do a mandatory verification step because we don't1323  // entirely trust this pass.  Note that we don't want to add a verifier1324  // pass to FPM below because it will also verify all the global data.1325  if (verifyFunction(F, &errs()))1326    report_fatal_error("Broken function");1327#endif1328}1329 1330// Coroutine has no suspend points. Remove heap allocation for the coroutine1331// frame if possible.1332static void handleNoSuspendCoroutine(coro::Shape &Shape) {1333  auto *CoroBegin = Shape.CoroBegin;1334  switch (Shape.ABI) {1335  case coro::ABI::Switch: {1336    auto SwitchId = Shape.getSwitchCoroId();1337    auto *AllocInst = SwitchId->getCoroAlloc();1338    coro::replaceCoroFree(SwitchId, /*Elide=*/AllocInst != nullptr);1339    if (AllocInst) {1340      IRBuilder<> Builder(AllocInst);1341      auto *Frame = Builder.CreateAlloca(Shape.FrameTy);1342      Frame->setAlignment(Shape.FrameAlign);1343      AllocInst->replaceAllUsesWith(Builder.getFalse());1344      AllocInst->eraseFromParent();1345      CoroBegin->replaceAllUsesWith(Frame);1346    } else {1347      CoroBegin->replaceAllUsesWith(CoroBegin->getMem());1348    }1349 1350    break;1351  }1352  case coro::ABI::Async:1353  case coro::ABI::Retcon:1354  case coro::ABI::RetconOnce:1355    CoroBegin->replaceAllUsesWith(PoisonValue::get(CoroBegin->getType()));1356    break;1357  }1358 1359  CoroBegin->eraseFromParent();1360  Shape.CoroBegin = nullptr;1361}1362 1363// SimplifySuspendPoint needs to check that there is no calls between1364// coro_save and coro_suspend, since any of the calls may potentially resume1365// the coroutine and if that is the case we cannot eliminate the suspend point.1366static bool hasCallsInBlockBetween(iterator_range<BasicBlock::iterator> R) {1367  for (Instruction &I : R) {1368    // Assume that no intrinsic can resume the coroutine.1369    if (isa<IntrinsicInst>(I))1370      continue;1371 1372    if (isa<CallBase>(I))1373      return true;1374  }1375  return false;1376}1377 1378static bool hasCallsInBlocksBetween(BasicBlock *SaveBB, BasicBlock *ResDesBB) {1379  SmallPtrSet<BasicBlock *, 8> Set;1380  SmallVector<BasicBlock *, 8> Worklist;1381 1382  Set.insert(SaveBB);1383  Worklist.push_back(ResDesBB);1384 1385  // Accumulate all blocks between SaveBB and ResDesBB. Because CoroSaveIntr1386  // returns a token consumed by suspend instruction, all blocks in between1387  // will have to eventually hit SaveBB when going backwards from ResDesBB.1388  while (!Worklist.empty()) {1389    auto *BB = Worklist.pop_back_val();1390    Set.insert(BB);1391    for (auto *Pred : predecessors(BB))1392      if (!Set.contains(Pred))1393        Worklist.push_back(Pred);1394  }1395 1396  // SaveBB and ResDesBB are checked separately in hasCallsBetween.1397  Set.erase(SaveBB);1398  Set.erase(ResDesBB);1399 1400  for (auto *BB : Set)1401    if (hasCallsInBlockBetween({BB->getFirstNonPHIIt(), BB->end()}))1402      return true;1403 1404  return false;1405}1406 1407static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {1408  auto *SaveBB = Save->getParent();1409  auto *ResumeOrDestroyBB = ResumeOrDestroy->getParent();1410  BasicBlock::iterator SaveIt = Save->getIterator();1411  BasicBlock::iterator ResumeOrDestroyIt = ResumeOrDestroy->getIterator();1412 1413  if (SaveBB == ResumeOrDestroyBB)1414    return hasCallsInBlockBetween({std::next(SaveIt), ResumeOrDestroyIt});1415 1416  // Any calls from Save to the end of the block?1417  if (hasCallsInBlockBetween({std::next(SaveIt), SaveBB->end()}))1418    return true;1419 1420  // Any calls from begging of the block up to ResumeOrDestroy?1421  if (hasCallsInBlockBetween(1422          {ResumeOrDestroyBB->getFirstNonPHIIt(), ResumeOrDestroyIt}))1423    return true;1424 1425  // Any calls in all of the blocks between SaveBB and ResumeOrDestroyBB?1426  if (hasCallsInBlocksBetween(SaveBB, ResumeOrDestroyBB))1427    return true;1428 1429  return false;1430}1431 1432// If a SuspendIntrin is preceded by Resume or Destroy, we can eliminate the1433// suspend point and replace it with nornal control flow.1434static bool simplifySuspendPoint(CoroSuspendInst *Suspend,1435                                 CoroBeginInst *CoroBegin) {1436  Instruction *Prev = Suspend->getPrevNode();1437  if (!Prev) {1438    auto *Pred = Suspend->getParent()->getSinglePredecessor();1439    if (!Pred)1440      return false;1441    Prev = Pred->getTerminator();1442  }1443 1444  CallBase *CB = dyn_cast<CallBase>(Prev);1445  if (!CB)1446    return false;1447 1448  auto *Callee = CB->getCalledOperand()->stripPointerCasts();1449 1450  // See if the callsite is for resumption or destruction of the coroutine.1451  auto *SubFn = dyn_cast<CoroSubFnInst>(Callee);1452  if (!SubFn)1453    return false;1454 1455  // Does not refer to the current coroutine, we cannot do anything with it.1456  if (SubFn->getFrame() != CoroBegin)1457    return false;1458 1459  // See if the transformation is safe. Specifically, see if there are any1460  // calls in between Save and CallInstr. They can potenitally resume the1461  // coroutine rendering this optimization unsafe.1462  auto *Save = Suspend->getCoroSave();1463  if (hasCallsBetween(Save, CB))1464    return false;1465 1466  // Replace llvm.coro.suspend with the value that results in resumption over1467  // the resume or cleanup path.1468  Suspend->replaceAllUsesWith(SubFn->getRawIndex());1469  Suspend->eraseFromParent();1470  Save->eraseFromParent();1471 1472  // No longer need a call to coro.resume or coro.destroy.1473  if (auto *Invoke = dyn_cast<InvokeInst>(CB)) {1474    BranchInst::Create(Invoke->getNormalDest(), Invoke->getIterator());1475  }1476 1477  // Grab the CalledValue from CB before erasing the CallInstr.1478  auto *CalledValue = CB->getCalledOperand();1479  CB->eraseFromParent();1480 1481  // If no more users remove it. Usually it is a bitcast of SubFn.1482  if (CalledValue != SubFn && CalledValue->user_empty())1483    if (auto *I = dyn_cast<Instruction>(CalledValue))1484      I->eraseFromParent();1485 1486  // Now we are good to remove SubFn.1487  if (SubFn->user_empty())1488    SubFn->eraseFromParent();1489 1490  return true;1491}1492 1493// Remove suspend points that are simplified.1494static void simplifySuspendPoints(coro::Shape &Shape) {1495  // Currently, the only simplification we do is switch-lowering-specific.1496  if (Shape.ABI != coro::ABI::Switch)1497    return;1498 1499  auto &S = Shape.CoroSuspends;1500  size_t I = 0, N = S.size();1501  if (N == 0)1502    return;1503 1504  size_t ChangedFinalIndex = std::numeric_limits<size_t>::max();1505  while (true) {1506    auto SI = cast<CoroSuspendInst>(S[I]);1507    // Leave final.suspend to handleFinalSuspend since it is undefined behavior1508    // to resume a coroutine suspended at the final suspend point.1509    if (!SI->isFinal() && simplifySuspendPoint(SI, Shape.CoroBegin)) {1510      if (--N == I)1511        break;1512 1513      std::swap(S[I], S[N]);1514 1515      if (cast<CoroSuspendInst>(S[I])->isFinal()) {1516        assert(Shape.SwitchLowering.HasFinalSuspend);1517        ChangedFinalIndex = I;1518      }1519 1520      continue;1521    }1522    if (++I == N)1523      break;1524  }1525  S.resize(N);1526 1527  // Maintain final.suspend in case final suspend was swapped.1528  // Due to we requrie the final suspend to be the last element of CoroSuspends.1529  if (ChangedFinalIndex < N) {1530    assert(cast<CoroSuspendInst>(S[ChangedFinalIndex])->isFinal());1531    std::swap(S[ChangedFinalIndex], S.back());1532  }1533}1534 1535namespace {1536 1537struct SwitchCoroutineSplitter {1538  static void split(Function &F, coro::Shape &Shape,1539                    SmallVectorImpl<Function *> &Clones,1540                    TargetTransformInfo &TTI) {1541    assert(Shape.ABI == coro::ABI::Switch);1542 1543    // Create a resume clone by cloning the body of the original function,1544    // setting new entry block and replacing coro.suspend an appropriate value1545    // to force resume or cleanup pass for every suspend point.1546    createResumeEntryBlock(F, Shape);1547    auto *ResumeClone = coro::SwitchCloner::createClone(1548        F, ".resume", Shape, coro::CloneKind::SwitchResume, TTI);1549    auto *DestroyClone = coro::SwitchCloner::createClone(1550        F, ".destroy", Shape, coro::CloneKind::SwitchUnwind, TTI);1551    auto *CleanupClone = coro::SwitchCloner::createClone(1552        F, ".cleanup", Shape, coro::CloneKind::SwitchCleanup, TTI);1553 1554    postSplitCleanup(*ResumeClone);1555    postSplitCleanup(*DestroyClone);1556    postSplitCleanup(*CleanupClone);1557 1558    // Store addresses resume/destroy/cleanup functions in the coroutine frame.1559    updateCoroFrame(Shape, ResumeClone, DestroyClone, CleanupClone);1560 1561    assert(Clones.empty());1562    Clones.push_back(ResumeClone);1563    Clones.push_back(DestroyClone);1564    Clones.push_back(CleanupClone);1565 1566    // Create a constant array referring to resume/destroy/clone functions1567    // pointed by the last argument of @llvm.coro.info, so that CoroElide pass1568    // can determined correct function to call.1569    setCoroInfo(F, Shape, Clones);1570  }1571 1572  // Create a variant of ramp function that does not perform heap allocation1573  // for a switch ABI coroutine.1574  //1575  // The newly split `.noalloc` ramp function has the following differences:1576  //  - Has one additional frame pointer parameter in lieu of dynamic1577  //  allocation.1578  //  - Suppressed allocations by replacing coro.alloc and coro.free.1579  static Function *createNoAllocVariant(Function &F, coro::Shape &Shape,1580                                        SmallVectorImpl<Function *> &Clones) {1581    assert(Shape.ABI == coro::ABI::Switch);1582    auto *OrigFnTy = F.getFunctionType();1583    auto OldParams = OrigFnTy->params();1584 1585    SmallVector<Type *> NewParams;1586    NewParams.reserve(OldParams.size() + 1);1587    NewParams.append(OldParams.begin(), OldParams.end());1588    NewParams.push_back(PointerType::getUnqual(Shape.FrameTy->getContext()));1589 1590    auto *NewFnTy = FunctionType::get(OrigFnTy->getReturnType(), NewParams,1591                                      OrigFnTy->isVarArg());1592    Function *NoAllocF =1593        Function::Create(NewFnTy, F.getLinkage(), F.getName() + ".noalloc");1594 1595    ValueToValueMapTy VMap;1596    unsigned int Idx = 0;1597    for (const auto &I : F.args()) {1598      VMap[&I] = NoAllocF->getArg(Idx++);1599    }1600    // We just appended the frame pointer as the last argument of the new1601    // function.1602    auto FrameIdx = NoAllocF->arg_size() - 1;1603    SmallVector<ReturnInst *, 4> Returns;1604    CloneFunctionInto(NoAllocF, &F, VMap,1605                      CloneFunctionChangeType::LocalChangesOnly, Returns);1606 1607    if (Shape.CoroBegin) {1608      auto *NewCoroBegin =1609          cast_if_present<CoroBeginInst>(VMap[Shape.CoroBegin]);1610      auto *NewCoroId = cast<CoroIdInst>(NewCoroBegin->getId());1611      coro::replaceCoroFree(NewCoroId, /*Elide=*/true);1612      coro::suppressCoroAllocs(NewCoroId);1613      NewCoroBegin->replaceAllUsesWith(NoAllocF->getArg(FrameIdx));1614      NewCoroBegin->eraseFromParent();1615    }1616 1617    Module *M = F.getParent();1618    M->getFunctionList().insert(M->end(), NoAllocF);1619 1620    removeUnreachableBlocks(*NoAllocF);1621    auto NewAttrs = NoAllocF->getAttributes();1622    // When we elide allocation, we read these attributes to determine the1623    // frame size and alignment.1624    addFramePointerAttrs(NewAttrs, NoAllocF->getContext(), FrameIdx,1625                         Shape.FrameSize, Shape.FrameAlign,1626                         /*NoAlias=*/false);1627 1628    NoAllocF->setAttributes(NewAttrs);1629 1630    Clones.push_back(NoAllocF);1631    // Reset the original function's coro info, make the new noalloc variant1632    // connected to the original ramp function.1633    setCoroInfo(F, Shape, Clones);1634    // After copying, set the linkage to internal linkage. Original function1635    // may have different linkage, but optimization dependent on this function1636    // generally relies on LTO.1637    NoAllocF->setLinkage(llvm::GlobalValue::InternalLinkage);1638    return NoAllocF;1639  }1640 1641private:1642  // Create an entry block for a resume function with a switch that will jump to1643  // suspend points.1644  static void createResumeEntryBlock(Function &F, coro::Shape &Shape) {1645    LLVMContext &C = F.getContext();1646 1647    DIBuilder DBuilder(*F.getParent(), /*AllowUnresolved*/ false);1648    DISubprogram *DIS = F.getSubprogram();1649    // If there is no DISubprogram for F, it implies the function is compiled1650    // without debug info. So we also don't generate debug info for the1651    // suspension points.1652    bool AddDebugLabels = DIS && DIS->getUnit() &&1653                          (DIS->getUnit()->getEmissionKind() ==1654                           DICompileUnit::DebugEmissionKind::FullDebug);1655 1656    // resume.entry:1657    //  %index.addr = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i321658    //  0, i32 2 % index = load i32, i32* %index.addr switch i32 %index, label1659    //  %unreachable [1660    //    i32 0, label %resume.01661    //    i32 1, label %resume.11662    //    ...1663    //  ]1664 1665    auto *NewEntry = BasicBlock::Create(C, "resume.entry", &F);1666    auto *UnreachBB = BasicBlock::Create(C, "unreachable", &F);1667 1668    IRBuilder<> Builder(NewEntry);1669    auto *FramePtr = Shape.FramePtr;1670    auto *FrameTy = Shape.FrameTy;1671    auto *GepIndex = Builder.CreateStructGEP(1672        FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");1673    auto *Index = Builder.CreateLoad(Shape.getIndexType(), GepIndex, "index");1674    auto *Switch =1675        Builder.CreateSwitch(Index, UnreachBB, Shape.CoroSuspends.size());1676    Shape.SwitchLowering.ResumeSwitch = Switch;1677 1678    // Split all coro.suspend calls1679    size_t SuspendIndex = 0;1680    for (auto *AnyS : Shape.CoroSuspends) {1681      auto *S = cast<CoroSuspendInst>(AnyS);1682      ConstantInt *IndexVal = Shape.getIndex(SuspendIndex);1683 1684      // Replace CoroSave with a store to Index:1685      //    %index.addr = getelementptr %f.frame... (index field number)1686      //    store i32 %IndexVal, i32* %index.addr11687      auto *Save = S->getCoroSave();1688      Builder.SetInsertPoint(Save);1689      if (S->isFinal()) {1690        // The coroutine should be marked done if it reaches the final suspend1691        // point.1692        markCoroutineAsDone(Builder, Shape, FramePtr);1693      } else {1694        auto *GepIndex = Builder.CreateStructGEP(1695            FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");1696        Builder.CreateStore(IndexVal, GepIndex);1697      }1698 1699      Save->replaceAllUsesWith(ConstantTokenNone::get(C));1700      Save->eraseFromParent();1701 1702      // Split block before and after coro.suspend and add a jump from an entry1703      // switch:1704      //1705      //  whateverBB:1706      //    whatever1707      //    %0 = call i8 @llvm.coro.suspend(token none, i1 false)1708      //    switch i8 %0, label %suspend[i8 0, label %resume1709      //                                 i8 1, label %cleanup]1710      // becomes:1711      //1712      //  whateverBB:1713      //     whatever1714      //     br label %resume.0.landing1715      //1716      //  resume.0: ; <--- jump from the switch in the resume.entry1717      //        #dbg_label(...)  ; <--- artificial label for debuggers1718      //     %0 = tail call i8 @llvm.coro.suspend(token none, i1 false)1719      //     br label %resume.0.landing1720      //1721      //  resume.0.landing:1722      //     %1 = phi i8[-1, %whateverBB], [%0, %resume.0]1723      //     switch i8 % 1, label %suspend [i8 0, label %resume1724      //                                    i8 1, label %cleanup]1725 1726      auto *SuspendBB = S->getParent();1727      auto *ResumeBB =1728          SuspendBB->splitBasicBlock(S, "resume." + Twine(SuspendIndex));1729      auto *LandingBB = ResumeBB->splitBasicBlock(1730          S->getNextNode(), ResumeBB->getName() + Twine(".landing"));1731      Switch->addCase(IndexVal, ResumeBB);1732 1733      cast<BranchInst>(SuspendBB->getTerminator())->setSuccessor(0, LandingBB);1734      auto *PN = PHINode::Create(Builder.getInt8Ty(), 2, "");1735      PN->insertBefore(LandingBB->begin());1736      S->replaceAllUsesWith(PN);1737      PN->addIncoming(Builder.getInt8(-1), SuspendBB);1738      PN->addIncoming(S, ResumeBB);1739 1740      if (AddDebugLabels) {1741        if (DebugLoc SuspendLoc = S->getDebugLoc()) {1742          std::string LabelName =1743              ("__coro_resume_" + Twine(SuspendIndex)).str();1744          // Take the "inlined at" location recursively, if present. This is1745          // mandatory as the DILabel insertion checks that the scopes of label1746          // and the attached location match. This is not the case when the1747          // suspend location has been inlined due to pointing to the original1748          // scope.1749          DILocation *DILoc = SuspendLoc;1750          while (DILocation *InlinedAt = DILoc->getInlinedAt())1751            DILoc = InlinedAt;1752 1753          DILabel *ResumeLabel =1754              DBuilder.createLabel(DIS, LabelName, DILoc->getFile(),1755                                   SuspendLoc.getLine(), SuspendLoc.getCol(),1756                                   /*IsArtificial=*/true,1757                                   /*CoroSuspendIdx=*/SuspendIndex,1758                                   /*AlwaysPreserve=*/false);1759          DBuilder.insertLabel(ResumeLabel, DILoc, ResumeBB->begin());1760        }1761      }1762 1763      ++SuspendIndex;1764    }1765 1766    Builder.SetInsertPoint(UnreachBB);1767    Builder.CreateUnreachable();1768    DBuilder.finalize();1769 1770    Shape.SwitchLowering.ResumeEntryBlock = NewEntry;1771  }1772 1773  // Store addresses of Resume/Destroy/Cleanup functions in the coroutine frame.1774  static void updateCoroFrame(coro::Shape &Shape, Function *ResumeFn,1775                              Function *DestroyFn, Function *CleanupFn) {1776    IRBuilder<> Builder(&*Shape.getInsertPtAfterFramePtr());1777 1778    auto *ResumeAddr = Builder.CreateStructGEP(1779        Shape.FrameTy, Shape.FramePtr, coro::Shape::SwitchFieldIndex::Resume,1780        "resume.addr");1781    Builder.CreateStore(ResumeFn, ResumeAddr);1782 1783    Value *DestroyOrCleanupFn = DestroyFn;1784 1785    CoroIdInst *CoroId = Shape.getSwitchCoroId();1786    if (CoroAllocInst *CA = CoroId->getCoroAlloc()) {1787      // If there is a CoroAlloc and it returns false (meaning we elide the1788      // allocation, use CleanupFn instead of DestroyFn).1789      DestroyOrCleanupFn = Builder.CreateSelect(CA, DestroyFn, CleanupFn);1790    }1791 1792    auto *DestroyAddr = Builder.CreateStructGEP(1793        Shape.FrameTy, Shape.FramePtr, coro::Shape::SwitchFieldIndex::Destroy,1794        "destroy.addr");1795    Builder.CreateStore(DestroyOrCleanupFn, DestroyAddr);1796  }1797 1798  // Create a global constant array containing pointers to functions provided1799  // and set Info parameter of CoroBegin to point at this constant. Example:1800  //1801  //   @f.resumers = internal constant [2 x void(%f.frame*)*]1802  //                    [void(%f.frame*)* @f.resume, void(%f.frame*)*1803  //                    @f.destroy]1804  //   define void @f() {1805  //     ...1806  //     call i8* @llvm.coro.begin(i8* null, i32 0, i8* null,1807  //                    i8* bitcast([2 x void(%f.frame*)*] * @f.resumers to1808  //                    i8*))1809  //1810  // Assumes that all the functions have the same signature.1811  static void setCoroInfo(Function &F, coro::Shape &Shape,1812                          ArrayRef<Function *> Fns) {1813    // This only works under the switch-lowering ABI because coro elision1814    // only works on the switch-lowering ABI.1815    SmallVector<Constant *, 4> Args(Fns);1816    assert(!Args.empty());1817    Function *Part = *Fns.begin();1818    Module *M = Part->getParent();1819    auto *ArrTy = ArrayType::get(Part->getType(), Args.size());1820 1821    auto *ConstVal = ConstantArray::get(ArrTy, Args);1822    auto *GV = new GlobalVariable(*M, ConstVal->getType(), /*isConstant=*/true,1823                                  GlobalVariable::PrivateLinkage, ConstVal,1824                                  F.getName() + Twine(".resumers"));1825 1826    // Update coro.begin instruction to refer to this constant.1827    LLVMContext &C = F.getContext();1828    auto *BC = ConstantExpr::getPointerCast(GV, PointerType::getUnqual(C));1829    Shape.getSwitchCoroId()->setInfo(BC);1830  }1831};1832 1833} // namespace1834 1835static void replaceAsyncResumeFunction(CoroSuspendAsyncInst *Suspend,1836                                       Value *Continuation) {1837  auto *ResumeIntrinsic = Suspend->getResumeFunction();1838  auto &Context = Suspend->getParent()->getParent()->getContext();1839  auto *Int8PtrTy = PointerType::getUnqual(Context);1840 1841  IRBuilder<> Builder(ResumeIntrinsic);1842  auto *Val = Builder.CreateBitOrPointerCast(Continuation, Int8PtrTy);1843  ResumeIntrinsic->replaceAllUsesWith(Val);1844  ResumeIntrinsic->eraseFromParent();1845  Suspend->setOperand(CoroSuspendAsyncInst::ResumeFunctionArg,1846                      PoisonValue::get(Int8PtrTy));1847}1848 1849/// Coerce the arguments in \p FnArgs according to \p FnTy in \p CallArgs.1850static void coerceArguments(IRBuilder<> &Builder, FunctionType *FnTy,1851                            ArrayRef<Value *> FnArgs,1852                            SmallVectorImpl<Value *> &CallArgs) {1853  size_t ArgIdx = 0;1854  for (auto *paramTy : FnTy->params()) {1855    assert(ArgIdx < FnArgs.size());1856    if (paramTy != FnArgs[ArgIdx]->getType())1857      CallArgs.push_back(1858          Builder.CreateBitOrPointerCast(FnArgs[ArgIdx], paramTy));1859    else1860      CallArgs.push_back(FnArgs[ArgIdx]);1861    ++ArgIdx;1862  }1863}1864 1865CallInst *coro::createMustTailCall(DebugLoc Loc, Function *MustTailCallFn,1866                                   TargetTransformInfo &TTI,1867                                   ArrayRef<Value *> Arguments,1868                                   IRBuilder<> &Builder) {1869  auto *FnTy = MustTailCallFn->getFunctionType();1870  // Coerce the arguments, llvm optimizations seem to ignore the types in1871  // vaarg functions and throws away casts in optimized mode.1872  SmallVector<Value *, 8> CallArgs;1873  coerceArguments(Builder, FnTy, Arguments, CallArgs);1874 1875  auto *TailCall = Builder.CreateCall(FnTy, MustTailCallFn, CallArgs);1876  // Skip targets which don't support tail call.1877  if (TTI.supportsTailCallFor(TailCall)) {1878    TailCall->setTailCallKind(CallInst::TCK_MustTail);1879  }1880  TailCall->setDebugLoc(Loc);1881  TailCall->setCallingConv(MustTailCallFn->getCallingConv());1882  return TailCall;1883}1884 1885void coro::AsyncABI::splitCoroutine(Function &F, coro::Shape &Shape,1886                                    SmallVectorImpl<Function *> &Clones,1887                                    TargetTransformInfo &TTI) {1888  assert(Shape.ABI == coro::ABI::Async);1889  assert(Clones.empty());1890  // Reset various things that the optimizer might have decided it1891  // "knows" about the coroutine function due to not seeing a return.1892  F.removeFnAttr(Attribute::NoReturn);1893  F.removeRetAttr(Attribute::NoAlias);1894  F.removeRetAttr(Attribute::NonNull);1895 1896  auto &Context = F.getContext();1897  auto *Int8PtrTy = PointerType::getUnqual(Context);1898 1899  auto *Id = Shape.getAsyncCoroId();1900  IRBuilder<> Builder(Id);1901 1902  auto *FramePtr = Id->getStorage();1903  FramePtr = Builder.CreateBitOrPointerCast(FramePtr, Int8PtrTy);1904  FramePtr = Builder.CreateConstInBoundsGEP1_32(1905      Type::getInt8Ty(Context), FramePtr, Shape.AsyncLowering.FrameOffset,1906      "async.ctx.frameptr");1907 1908  // Map all uses of llvm.coro.begin to the allocated frame pointer.1909  {1910    // Make sure we don't invalidate Shape.FramePtr.1911    TrackingVH<Value> Handle(Shape.FramePtr);1912    Shape.CoroBegin->replaceAllUsesWith(FramePtr);1913    Shape.FramePtr = Handle.getValPtr();1914  }1915 1916  // Create all the functions in order after the main function.1917  auto NextF = std::next(F.getIterator());1918 1919  // Create a continuation function for each of the suspend points.1920  Clones.reserve(Shape.CoroSuspends.size());1921  for (auto [Idx, CS] : llvm::enumerate(Shape.CoroSuspends)) {1922    auto *Suspend = cast<CoroSuspendAsyncInst>(CS);1923 1924    // Create the clone declaration.1925    auto ResumeNameSuffix = ".resume.";1926    auto ProjectionFunctionName =1927        Suspend->getAsyncContextProjectionFunction()->getName();1928    bool UseSwiftMangling = false;1929    if (ProjectionFunctionName == "__swift_async_resume_project_context") {1930      ResumeNameSuffix = "TQ";1931      UseSwiftMangling = true;1932    } else if (ProjectionFunctionName == "__swift_async_resume_get_context") {1933      ResumeNameSuffix = "TY";1934      UseSwiftMangling = true;1935    }1936    auto *Continuation = createCloneDeclaration(1937        F, Shape,1938        UseSwiftMangling ? ResumeNameSuffix + Twine(Idx) + "_"1939                         : ResumeNameSuffix + Twine(Idx),1940        NextF, Suspend);1941    Clones.push_back(Continuation);1942 1943    // Insert a branch to a new return block immediately before the suspend1944    // point.1945    auto *SuspendBB = Suspend->getParent();1946    auto *NewSuspendBB = SuspendBB->splitBasicBlock(Suspend);1947    auto *Branch = cast<BranchInst>(SuspendBB->getTerminator());1948 1949    // Place it before the first suspend.1950    auto *ReturnBB =1951        BasicBlock::Create(F.getContext(), "coro.return", &F, NewSuspendBB);1952    Branch->setSuccessor(0, ReturnBB);1953 1954    IRBuilder<> Builder(ReturnBB);1955 1956    // Insert the call to the tail call function and inline it.1957    auto *Fn = Suspend->getMustTailCallFunction();1958    SmallVector<Value *, 8> Args(Suspend->args());1959    auto FnArgs = ArrayRef<Value *>(Args).drop_front(1960        CoroSuspendAsyncInst::MustTailCallFuncArg + 1);1961    auto *TailCall = coro::createMustTailCall(Suspend->getDebugLoc(), Fn, TTI,1962                                              FnArgs, Builder);1963    Builder.CreateRetVoid();1964    InlineFunctionInfo FnInfo;1965    (void)InlineFunction(*TailCall, FnInfo);1966 1967    // Replace the lvm.coro.async.resume intrisic call.1968    replaceAsyncResumeFunction(Suspend, Continuation);1969  }1970 1971  assert(Clones.size() == Shape.CoroSuspends.size());1972 1973  for (auto [Idx, CS] : llvm::enumerate(Shape.CoroSuspends)) {1974    auto *Suspend = CS;1975    auto *Clone = Clones[Idx];1976 1977    coro::BaseCloner::createClone(F, "resume." + Twine(Idx), Shape, Clone,1978                                  Suspend, TTI);1979  }1980}1981 1982void coro::AnyRetconABI::splitCoroutine(Function &F, coro::Shape &Shape,1983                                        SmallVectorImpl<Function *> &Clones,1984                                        TargetTransformInfo &TTI) {1985  assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce);1986  assert(Clones.empty());1987 1988  // Reset various things that the optimizer might have decided it1989  // "knows" about the coroutine function due to not seeing a return.1990  F.removeFnAttr(Attribute::NoReturn);1991  F.removeRetAttr(Attribute::NoAlias);1992  F.removeRetAttr(Attribute::NonNull);1993 1994  // Allocate the frame.1995  auto *Id = Shape.getRetconCoroId();1996  Value *RawFramePtr;1997  if (Shape.RetconLowering.IsFrameInlineInStorage) {1998    RawFramePtr = Id->getStorage();1999  } else {2000    IRBuilder<> Builder(Id);2001 2002    // Determine the size of the frame.2003    const DataLayout &DL = F.getDataLayout();2004    auto Size = DL.getTypeAllocSize(Shape.FrameTy);2005 2006    // Allocate.  We don't need to update the call graph node because we're2007    // going to recompute it from scratch after splitting.2008    // FIXME: pass the required alignment2009    RawFramePtr = Shape.emitAlloc(Builder, Builder.getInt64(Size), nullptr);2010    RawFramePtr =2011        Builder.CreateBitCast(RawFramePtr, Shape.CoroBegin->getType());2012 2013    // Stash the allocated frame pointer in the continuation storage.2014    Builder.CreateStore(RawFramePtr, Id->getStorage());2015  }2016 2017  // Map all uses of llvm.coro.begin to the allocated frame pointer.2018  {2019    // Make sure we don't invalidate Shape.FramePtr.2020    TrackingVH<Value> Handle(Shape.FramePtr);2021    Shape.CoroBegin->replaceAllUsesWith(RawFramePtr);2022    Shape.FramePtr = Handle.getValPtr();2023  }2024 2025  // Create a unique return block.2026  BasicBlock *ReturnBB = nullptr;2027  PHINode *ContinuationPhi = nullptr;2028  SmallVector<PHINode *, 4> ReturnPHIs;2029 2030  // Create all the functions in order after the main function.2031  auto NextF = std::next(F.getIterator());2032 2033  // Create a continuation function for each of the suspend points.2034  Clones.reserve(Shape.CoroSuspends.size());2035  for (auto [Idx, CS] : llvm::enumerate(Shape.CoroSuspends)) {2036    auto Suspend = cast<CoroSuspendRetconInst>(CS);2037 2038    // Create the clone declaration.2039    auto Continuation = createCloneDeclaration(2040        F, Shape, ".resume." + Twine(Idx), NextF, nullptr);2041    Clones.push_back(Continuation);2042 2043    // Insert a branch to the unified return block immediately before2044    // the suspend point.2045    auto SuspendBB = Suspend->getParent();2046    auto NewSuspendBB = SuspendBB->splitBasicBlock(Suspend);2047    auto Branch = cast<BranchInst>(SuspendBB->getTerminator());2048 2049    // Create the unified return block.2050    if (!ReturnBB) {2051      // Place it before the first suspend.2052      ReturnBB =2053          BasicBlock::Create(F.getContext(), "coro.return", &F, NewSuspendBB);2054      Shape.RetconLowering.ReturnBlock = ReturnBB;2055 2056      IRBuilder<> Builder(ReturnBB);2057 2058      // First, the continuation.2059      ContinuationPhi =2060          Builder.CreatePHI(Continuation->getType(), Shape.CoroSuspends.size());2061 2062      // Create PHIs for all other return values.2063      assert(ReturnPHIs.empty());2064 2065      // Next, all the directly-yielded values.2066      for (auto *ResultTy : Shape.getRetconResultTypes())2067        ReturnPHIs.push_back(2068            Builder.CreatePHI(ResultTy, Shape.CoroSuspends.size()));2069 2070      // Build the return value.2071      auto RetTy = F.getReturnType();2072 2073      // Cast the continuation value if necessary.2074      // We can't rely on the types matching up because that type would2075      // have to be infinite.2076      auto CastedContinuationTy =2077          (ReturnPHIs.empty() ? RetTy : RetTy->getStructElementType(0));2078      auto *CastedContinuation =2079          Builder.CreateBitCast(ContinuationPhi, CastedContinuationTy);2080 2081      Value *RetV = CastedContinuation;2082      if (!ReturnPHIs.empty()) {2083        auto ValueIdx = 0;2084        RetV = PoisonValue::get(RetTy);2085        RetV = Builder.CreateInsertValue(RetV, CastedContinuation, ValueIdx++);2086 2087        for (auto Phi : ReturnPHIs)2088          RetV = Builder.CreateInsertValue(RetV, Phi, ValueIdx++);2089      }2090 2091      Builder.CreateRet(RetV);2092    }2093 2094    // Branch to the return block.2095    Branch->setSuccessor(0, ReturnBB);2096    assert(ContinuationPhi);2097    ContinuationPhi->addIncoming(Continuation, SuspendBB);2098    for (auto [Phi, VUse] :2099         llvm::zip_equal(ReturnPHIs, Suspend->value_operands()))2100      Phi->addIncoming(VUse, SuspendBB);2101  }2102 2103  assert(Clones.size() == Shape.CoroSuspends.size());2104 2105  for (auto [Idx, CS] : llvm::enumerate(Shape.CoroSuspends)) {2106    auto Suspend = CS;2107    auto Clone = Clones[Idx];2108 2109    coro::BaseCloner::createClone(F, "resume." + Twine(Idx), Shape, Clone,2110                                  Suspend, TTI);2111  }2112}2113 2114namespace {2115class PrettyStackTraceFunction : public PrettyStackTraceEntry {2116  Function &F;2117 2118public:2119  PrettyStackTraceFunction(Function &F) : F(F) {}2120  void print(raw_ostream &OS) const override {2121    OS << "While splitting coroutine ";2122    F.printAsOperand(OS, /*print type*/ false, F.getParent());2123    OS << "\n";2124  }2125};2126} // namespace2127 2128/// Remove calls to llvm.coro.end in the original function.2129static void removeCoroEndsFromRampFunction(const coro::Shape &Shape) {2130  if (Shape.ABI != coro::ABI::Switch) {2131    for (auto *End : Shape.CoroEnds) {2132      replaceCoroEnd(End, Shape, Shape.FramePtr, /*in ramp*/ true, nullptr);2133    }2134  } else {2135    for (llvm::AnyCoroEndInst *End : Shape.CoroEnds)2136      End->eraseFromParent();2137  }2138}2139 2140static void removeCoroIsInRampFromRampFunction(const coro::Shape &Shape) {2141  for (auto *II : Shape.CoroIsInRampInsts) {2142    auto &Ctx = II->getContext();2143    II->replaceAllUsesWith(ConstantInt::getTrue(Ctx));2144    II->eraseFromParent();2145  }2146}2147 2148static bool hasSafeElideCaller(Function &F) {2149  for (auto *U : F.users()) {2150    if (auto *CB = dyn_cast<CallBase>(U)) {2151      auto *Caller = CB->getFunction();2152      if (Caller && Caller->isPresplitCoroutine() &&2153          CB->hasFnAttr(llvm::Attribute::CoroElideSafe))2154        return true;2155    }2156  }2157  return false;2158}2159 2160void coro::SwitchABI::splitCoroutine(Function &F, coro::Shape &Shape,2161                                     SmallVectorImpl<Function *> &Clones,2162                                     TargetTransformInfo &TTI) {2163  SwitchCoroutineSplitter::split(F, Shape, Clones, TTI);2164}2165 2166static void doSplitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,2167                             coro::BaseABI &ABI, TargetTransformInfo &TTI,2168                             bool OptimizeFrame) {2169  PrettyStackTraceFunction prettyStackTrace(F);2170 2171  auto &Shape = ABI.Shape;2172  assert(Shape.CoroBegin);2173 2174  lowerAwaitSuspends(F, Shape);2175 2176  simplifySuspendPoints(Shape);2177 2178  normalizeCoroutine(F, Shape, TTI);2179  ABI.buildCoroutineFrame(OptimizeFrame);2180  replaceFrameSizeAndAlignment(Shape);2181 2182  bool isNoSuspendCoroutine = Shape.CoroSuspends.empty();2183 2184  bool shouldCreateNoAllocVariant =2185      !isNoSuspendCoroutine && Shape.ABI == coro::ABI::Switch &&2186      hasSafeElideCaller(F) && !F.hasFnAttribute(llvm::Attribute::NoInline);2187 2188  // If there are no suspend points, no split required, just remove2189  // the allocation and deallocation blocks, they are not needed.2190  if (isNoSuspendCoroutine) {2191    handleNoSuspendCoroutine(Shape);2192  } else {2193    ABI.splitCoroutine(F, Shape, Clones, TTI);2194  }2195 2196  // Replace all the swifterror operations in the original function.2197  // This invalidates SwiftErrorOps in the Shape.2198  replaceSwiftErrorOps(F, Shape, nullptr);2199 2200  // Salvage debug intrinsics that point into the coroutine frame in the2201  // original function. The Cloner has already salvaged debug info in the new2202  // coroutine funclets.2203  SmallDenseMap<Argument *, AllocaInst *, 4> ArgToAllocaMap;2204  auto DbgVariableRecords = collectDbgVariableRecords(F);2205  for (DbgVariableRecord *DVR : DbgVariableRecords)2206    coro::salvageDebugInfo(ArgToAllocaMap, *DVR, false /*UseEntryValue*/);2207 2208  removeCoroEndsFromRampFunction(Shape);2209  removeCoroIsInRampFromRampFunction(Shape);2210 2211  if (shouldCreateNoAllocVariant)2212    SwitchCoroutineSplitter::createNoAllocVariant(F, Shape, Clones);2213}2214 2215static LazyCallGraph::SCC &updateCallGraphAfterCoroutineSplit(2216    LazyCallGraph::Node &N, const coro::Shape &Shape,2217    const SmallVectorImpl<Function *> &Clones, LazyCallGraph::SCC &C,2218    LazyCallGraph &CG, CGSCCAnalysisManager &AM, CGSCCUpdateResult &UR,2219    FunctionAnalysisManager &FAM) {2220 2221  auto *CurrentSCC = &C;2222  if (!Clones.empty()) {2223    switch (Shape.ABI) {2224    case coro::ABI::Switch:2225      // Each clone in the Switch lowering is independent of the other clones.2226      // Let the LazyCallGraph know about each one separately.2227      for (Function *Clone : Clones)2228        CG.addSplitFunction(N.getFunction(), *Clone);2229      break;2230    case coro::ABI::Async:2231    case coro::ABI::Retcon:2232    case coro::ABI::RetconOnce:2233      // Each clone in the Async/Retcon lowering references of the other clones.2234      // Let the LazyCallGraph know about all of them at once.2235      if (!Clones.empty())2236        CG.addSplitRefRecursiveFunctions(N.getFunction(), Clones);2237      break;2238    }2239 2240    // Let the CGSCC infra handle the changes to the original function.2241    CurrentSCC = &updateCGAndAnalysisManagerForCGSCCPass(CG, *CurrentSCC, N, AM,2242                                                         UR, FAM);2243  }2244 2245  // Do some cleanup and let the CGSCC infra see if we've cleaned up any edges2246  // to the split functions.2247  postSplitCleanup(N.getFunction());2248  CurrentSCC = &updateCGAndAnalysisManagerForFunctionPass(CG, *CurrentSCC, N,2249                                                          AM, UR, FAM);2250  return *CurrentSCC;2251}2252 2253/// Replace a call to llvm.coro.prepare.retcon.2254static void replacePrepare(CallInst *Prepare, LazyCallGraph &CG,2255                           LazyCallGraph::SCC &C) {2256  auto CastFn = Prepare->getArgOperand(0); // as an i8*2257  auto Fn = CastFn->stripPointerCasts();   // as its original type2258 2259  // Attempt to peephole this pattern:2260  //    %0 = bitcast [[TYPE]] @some_function to i8*2261  //    %1 = call @llvm.coro.prepare.retcon(i8* %0)2262  //    %2 = bitcast %1 to [[TYPE]]2263  // ==>2264  //    %2 = @some_function2265  for (Use &U : llvm::make_early_inc_range(Prepare->uses())) {2266    // Look for bitcasts back to the original function type.2267    auto *Cast = dyn_cast<BitCastInst>(U.getUser());2268    if (!Cast || Cast->getType() != Fn->getType())2269      continue;2270 2271    // Replace and remove the cast.2272    Cast->replaceAllUsesWith(Fn);2273    Cast->eraseFromParent();2274  }2275 2276  // Replace any remaining uses with the function as an i8*.2277  // This can never directly be a callee, so we don't need to update CG.2278  Prepare->replaceAllUsesWith(CastFn);2279  Prepare->eraseFromParent();2280 2281  // Kill dead bitcasts.2282  while (auto *Cast = dyn_cast<BitCastInst>(CastFn)) {2283    if (!Cast->use_empty())2284      break;2285    CastFn = Cast->getOperand(0);2286    Cast->eraseFromParent();2287  }2288}2289 2290static bool replaceAllPrepares(Function *PrepareFn, LazyCallGraph &CG,2291                               LazyCallGraph::SCC &C) {2292  bool Changed = false;2293  for (Use &P : llvm::make_early_inc_range(PrepareFn->uses())) {2294    // Intrinsics can only be used in calls.2295    auto *Prepare = cast<CallInst>(P.getUser());2296    replacePrepare(Prepare, CG, C);2297    Changed = true;2298  }2299 2300  return Changed;2301}2302 2303static void addPrepareFunction(const Module &M,2304                               SmallVectorImpl<Function *> &Fns,2305                               StringRef Name) {2306  auto *PrepareFn = M.getFunction(Name);2307  if (PrepareFn && !PrepareFn->use_empty())2308    Fns.push_back(PrepareFn);2309}2310 2311static std::unique_ptr<coro::BaseABI>2312CreateNewABI(Function &F, coro::Shape &S,2313             std::function<bool(Instruction &)> IsMatCallback,2314             const SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs) {2315  if (S.CoroBegin->hasCustomABI()) {2316    unsigned CustomABI = S.CoroBegin->getCustomABI();2317    if (CustomABI >= GenCustomABIs.size())2318      llvm_unreachable("Custom ABI not found amoung those specified");2319    return GenCustomABIs[CustomABI](F, S);2320  }2321 2322  switch (S.ABI) {2323  case coro::ABI::Switch:2324    return std::make_unique<coro::SwitchABI>(F, S, IsMatCallback);2325  case coro::ABI::Async:2326    return std::make_unique<coro::AsyncABI>(F, S, IsMatCallback);2327  case coro::ABI::Retcon:2328    return std::make_unique<coro::AnyRetconABI>(F, S, IsMatCallback);2329  case coro::ABI::RetconOnce:2330    return std::make_unique<coro::AnyRetconABI>(F, S, IsMatCallback);2331  }2332  llvm_unreachable("Unknown ABI");2333}2334 2335CoroSplitPass::CoroSplitPass(bool OptimizeFrame)2336    : CreateAndInitABI([](Function &F, coro::Shape &S) {2337        std::unique_ptr<coro::BaseABI> ABI =2338            CreateNewABI(F, S, coro::isTriviallyMaterializable, {});2339        ABI->init();2340        return ABI;2341      }),2342      OptimizeFrame(OptimizeFrame) {}2343 2344CoroSplitPass::CoroSplitPass(2345    SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs, bool OptimizeFrame)2346    : CreateAndInitABI([=](Function &F, coro::Shape &S) {2347        std::unique_ptr<coro::BaseABI> ABI =2348            CreateNewABI(F, S, coro::isTriviallyMaterializable, GenCustomABIs);2349        ABI->init();2350        return ABI;2351      }),2352      OptimizeFrame(OptimizeFrame) {}2353 2354// For back compatibility, constructor takes a materializable callback and2355// creates a generator for an ABI with a modified materializable callback.2356CoroSplitPass::CoroSplitPass(std::function<bool(Instruction &)> IsMatCallback,2357                             bool OptimizeFrame)2358    : CreateAndInitABI([=](Function &F, coro::Shape &S) {2359        std::unique_ptr<coro::BaseABI> ABI =2360            CreateNewABI(F, S, IsMatCallback, {});2361        ABI->init();2362        return ABI;2363      }),2364      OptimizeFrame(OptimizeFrame) {}2365 2366// For back compatibility, constructor takes a materializable callback and2367// creates a generator for an ABI with a modified materializable callback.2368CoroSplitPass::CoroSplitPass(2369    std::function<bool(Instruction &)> IsMatCallback,2370    SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs, bool OptimizeFrame)2371    : CreateAndInitABI([=](Function &F, coro::Shape &S) {2372        std::unique_ptr<coro::BaseABI> ABI =2373            CreateNewABI(F, S, IsMatCallback, GenCustomABIs);2374        ABI->init();2375        return ABI;2376      }),2377      OptimizeFrame(OptimizeFrame) {}2378 2379PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,2380                                     CGSCCAnalysisManager &AM,2381                                     LazyCallGraph &CG, CGSCCUpdateResult &UR) {2382  // NB: One invariant of a valid LazyCallGraph::SCC is that it must contain a2383  //     non-zero number of nodes, so we assume that here and grab the first2384  //     node's function's module.2385  Module &M = *C.begin()->getFunction().getParent();2386  auto &FAM =2387      AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();2388 2389  // Check for uses of llvm.coro.prepare.retcon/async.2390  SmallVector<Function *, 2> PrepareFns;2391  addPrepareFunction(M, PrepareFns, "llvm.coro.prepare.retcon");2392  addPrepareFunction(M, PrepareFns, "llvm.coro.prepare.async");2393 2394  // Find coroutines for processing.2395  SmallVector<LazyCallGraph::Node *> Coroutines;2396  for (LazyCallGraph::Node &N : C)2397    if (N.getFunction().isPresplitCoroutine())2398      Coroutines.push_back(&N);2399 2400  if (Coroutines.empty() && PrepareFns.empty())2401    return PreservedAnalyses::all();2402 2403  auto *CurrentSCC = &C;2404  // Split all the coroutines.2405  for (LazyCallGraph::Node *N : Coroutines) {2406    Function &F = N->getFunction();2407    LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName()2408                      << "\n");2409 2410    // The suspend-crossing algorithm in buildCoroutineFrame gets tripped up2411    // by unreachable blocks, so remove them as a first pass. Remove the2412    // unreachable blocks before collecting intrinsics into Shape.2413    removeUnreachableBlocks(F);2414 2415    coro::Shape Shape(F);2416    if (!Shape.CoroBegin)2417      continue;2418 2419    F.setSplittedCoroutine();2420 2421    std::unique_ptr<coro::BaseABI> ABI = CreateAndInitABI(F, Shape);2422 2423    SmallVector<Function *, 4> Clones;2424    auto &TTI = FAM.getResult<TargetIRAnalysis>(F);2425    doSplitCoroutine(F, Clones, *ABI, TTI, OptimizeFrame);2426    CurrentSCC = &updateCallGraphAfterCoroutineSplit(2427        *N, Shape, Clones, *CurrentSCC, CG, AM, UR, FAM);2428 2429    auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);2430    ORE.emit([&]() {2431      return OptimizationRemark(DEBUG_TYPE, "CoroSplit", &F)2432             << "Split '" << ore::NV("function", F.getName())2433             << "' (frame_size=" << ore::NV("frame_size", Shape.FrameSize)2434             << ", align=" << ore::NV("align", Shape.FrameAlign.value()) << ")";2435    });2436 2437    if (!Shape.CoroSuspends.empty()) {2438      // Run the CGSCC pipeline on the original and newly split functions.2439      UR.CWorklist.insert(CurrentSCC);2440      for (Function *Clone : Clones)2441        UR.CWorklist.insert(CG.lookupSCC(CG.get(*Clone)));2442    } else if (Shape.ABI == coro::ABI::Async) {2443      // Reprocess the function to inline the tail called return function of2444      // coro.async.end.2445      UR.CWorklist.insert(&C);2446    }2447  }2448 2449  for (auto *PrepareFn : PrepareFns) {2450    replaceAllPrepares(PrepareFn, CG, *CurrentSCC);2451  }2452 2453  return PreservedAnalyses::none();2454}2455