xref: /llvm-project/llvm/lib/Transforms/Coroutines/Coroutines.cpp (revision 6786928c4fe1f9daf720d3b604987de2b013e70b)
1 //===- Coroutines.cpp -----------------------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements the common infrastructure for Coroutine Passes.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "CoroInstr.h"
14 #include "CoroInternal.h"
15 #include "CoroShape.h"
16 #include "llvm/ADT/SmallVector.h"
17 #include "llvm/ADT/StringRef.h"
18 #include "llvm/Analysis/CallGraph.h"
19 #include "llvm/IR/Attributes.h"
20 #include "llvm/IR/Constants.h"
21 #include "llvm/IR/DerivedTypes.h"
22 #include "llvm/IR/Function.h"
23 #include "llvm/IR/InstIterator.h"
24 #include "llvm/IR/Instructions.h"
25 #include "llvm/IR/IntrinsicInst.h"
26 #include "llvm/IR/Intrinsics.h"
27 #include "llvm/IR/Module.h"
28 #include "llvm/IR/Type.h"
29 #include "llvm/Support/Casting.h"
30 #include "llvm/Support/ErrorHandling.h"
31 #include "llvm/Transforms/Utils/Local.h"
32 #include <cassert>
33 #include <cstddef>
34 #include <utility>
35 
36 using namespace llvm;
37 
38 // Construct the lowerer base class and initialize its members.
39 coro::LowererBase::LowererBase(Module &M)
40     : TheModule(M), Context(M.getContext()),
41       Int8Ptr(PointerType::get(Context, 0)),
42       ResumeFnType(FunctionType::get(Type::getVoidTy(Context), Int8Ptr,
43                                      /*isVarArg=*/false)),
44       NullPtr(ConstantPointerNull::get(Int8Ptr)) {}
45 
46 // Creates a call to llvm.coro.subfn.addr to obtain a resume function address.
47 // It generates the following:
48 //
49 //    call ptr @llvm.coro.subfn.addr(ptr %Arg, i8 %index)
50 
51 CallInst *coro::LowererBase::makeSubFnCall(Value *Arg, int Index,
52                                            Instruction *InsertPt) {
53   auto *IndexVal = ConstantInt::get(Type::getInt8Ty(Context), Index);
54   auto *Fn = Intrinsic::getDeclaration(&TheModule, Intrinsic::coro_subfn_addr);
55 
56   assert(Index >= CoroSubFnInst::IndexFirst &&
57          Index < CoroSubFnInst::IndexLast &&
58          "makeSubFnCall: Index value out of range");
59   return CallInst::Create(Fn, {Arg, IndexVal}, "", InsertPt->getIterator());
60 }
61 
62 // NOTE: Must be sorted!
63 static const char *const CoroIntrinsics[] = {
64     "llvm.coro.align",
65     "llvm.coro.alloc",
66     "llvm.coro.async.context.alloc",
67     "llvm.coro.async.context.dealloc",
68     "llvm.coro.async.resume",
69     "llvm.coro.async.size.replace",
70     "llvm.coro.async.store_resume",
71     "llvm.coro.await.suspend.bool",
72     "llvm.coro.await.suspend.handle",
73     "llvm.coro.await.suspend.void",
74     "llvm.coro.begin",
75     "llvm.coro.destroy",
76     "llvm.coro.done",
77     "llvm.coro.end",
78     "llvm.coro.end.async",
79     "llvm.coro.frame",
80     "llvm.coro.free",
81     "llvm.coro.id",
82     "llvm.coro.id.async",
83     "llvm.coro.id.retcon",
84     "llvm.coro.id.retcon.once",
85     "llvm.coro.noop",
86     "llvm.coro.prepare.async",
87     "llvm.coro.prepare.retcon",
88     "llvm.coro.promise",
89     "llvm.coro.resume",
90     "llvm.coro.save",
91     "llvm.coro.size",
92     "llvm.coro.subfn.addr",
93     "llvm.coro.suspend",
94     "llvm.coro.suspend.async",
95     "llvm.coro.suspend.retcon",
96 };
97 
98 #ifndef NDEBUG
99 static bool isCoroutineIntrinsicName(StringRef Name) {
100   return Intrinsic::lookupLLVMIntrinsicByName(CoroIntrinsics, Name, "coro") !=
101          -1;
102 }
103 #endif
104 
105 bool coro::isSuspendBlock(BasicBlock *BB) {
106   return isa<AnyCoroSuspendInst>(BB->front());
107 }
108 
109 bool coro::declaresAnyIntrinsic(const Module &M) {
110   for (StringRef Name : CoroIntrinsics) {
111     assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");
112     if (M.getNamedValue(Name))
113       return true;
114   }
115 
116   return false;
117 }
118 
119 // Verifies if a module has named values listed. Also, in debug mode verifies
120 // that names are intrinsic names.
121 bool coro::declaresIntrinsics(const Module &M,
122                               const std::initializer_list<StringRef> List) {
123   for (StringRef Name : List) {
124     assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");
125     if (M.getNamedValue(Name))
126       return true;
127   }
128 
129   return false;
130 }
131 
132 // Replace all coro.frees associated with the provided CoroId either with 'null'
133 // if Elide is true and with its frame parameter otherwise.
134 void coro::replaceCoroFree(CoroIdInst *CoroId, bool Elide) {
135   SmallVector<CoroFreeInst *, 4> CoroFrees;
136   for (User *U : CoroId->users())
137     if (auto CF = dyn_cast<CoroFreeInst>(U))
138       CoroFrees.push_back(CF);
139 
140   if (CoroFrees.empty())
141     return;
142 
143   Value *Replacement =
144       Elide
145           ? ConstantPointerNull::get(PointerType::get(CoroId->getContext(), 0))
146           : CoroFrees.front()->getFrame();
147 
148   for (CoroFreeInst *CF : CoroFrees) {
149     CF->replaceAllUsesWith(Replacement);
150     CF->eraseFromParent();
151   }
152 }
153 
154 void coro::suppressCoroAllocs(CoroIdInst *CoroId) {
155   SmallVector<CoroAllocInst *, 4> CoroAllocs;
156   for (User *U : CoroId->users())
157     if (auto *CA = dyn_cast<CoroAllocInst>(U))
158       CoroAllocs.push_back(CA);
159 
160   if (CoroAllocs.empty())
161     return;
162 
163   coro::suppressCoroAllocs(CoroId->getContext(), CoroAllocs);
164 }
165 
166 // Replacing llvm.coro.alloc with false will suppress dynamic
167 // allocation as it is expected for the frontend to generate the code that
168 // looks like:
169 //   id = coro.id(...)
170 //   mem = coro.alloc(id) ? malloc(coro.size()) : 0;
171 //   coro.begin(id, mem)
172 void coro::suppressCoroAllocs(LLVMContext &Context,
173                               ArrayRef<CoroAllocInst *> CoroAllocs) {
174   auto *False = ConstantInt::getFalse(Context);
175   for (auto *CA : CoroAllocs) {
176     CA->replaceAllUsesWith(False);
177     CA->eraseFromParent();
178   }
179 }
180 
181 static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
182                                     CoroSuspendInst *SuspendInst) {
183   Module *M = SuspendInst->getModule();
184   auto *Fn = Intrinsic::getDeclaration(M, Intrinsic::coro_save);
185   auto *SaveInst = cast<CoroSaveInst>(
186       CallInst::Create(Fn, CoroBegin, "", SuspendInst->getIterator()));
187   assert(!SuspendInst->getCoroSave());
188   SuspendInst->setArgOperand(0, SaveInst);
189   return SaveInst;
190 }
191 
192 // Collect "interesting" coroutine intrinsics.
193 void coro::Shape::analyze(Function &F,
194                           SmallVectorImpl<CoroFrameInst *> &CoroFrames,
195                           SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves) {
196   clear();
197 
198   bool HasFinalSuspend = false;
199   bool HasUnwindCoroEnd = false;
200   size_t FinalSuspendIndex = 0;
201 
202   for (Instruction &I : instructions(F)) {
203     // FIXME: coro_await_suspend_* are not proper `IntrinisicInst`s
204     // because they might be invoked
205     if (auto AWS = dyn_cast<CoroAwaitSuspendInst>(&I)) {
206       CoroAwaitSuspends.push_back(AWS);
207     } else if (auto II = dyn_cast<IntrinsicInst>(&I)) {
208       switch (II->getIntrinsicID()) {
209       default:
210         continue;
211       case Intrinsic::coro_size:
212         CoroSizes.push_back(cast<CoroSizeInst>(II));
213         break;
214       case Intrinsic::coro_align:
215         CoroAligns.push_back(cast<CoroAlignInst>(II));
216         break;
217       case Intrinsic::coro_frame:
218         CoroFrames.push_back(cast<CoroFrameInst>(II));
219         break;
220       case Intrinsic::coro_save:
221         // After optimizations, coro_suspends using this coro_save might have
222         // been removed, remember orphaned coro_saves to remove them later.
223         if (II->use_empty())
224           UnusedCoroSaves.push_back(cast<CoroSaveInst>(II));
225         break;
226       case Intrinsic::coro_suspend_async: {
227         auto *Suspend = cast<CoroSuspendAsyncInst>(II);
228         Suspend->checkWellFormed();
229         CoroSuspends.push_back(Suspend);
230         break;
231       }
232       case Intrinsic::coro_suspend_retcon: {
233         auto Suspend = cast<CoroSuspendRetconInst>(II);
234         CoroSuspends.push_back(Suspend);
235         break;
236       }
237       case Intrinsic::coro_suspend: {
238         auto Suspend = cast<CoroSuspendInst>(II);
239         CoroSuspends.push_back(Suspend);
240         if (Suspend->isFinal()) {
241           if (HasFinalSuspend)
242             report_fatal_error(
243               "Only one suspend point can be marked as final");
244           HasFinalSuspend = true;
245           FinalSuspendIndex = CoroSuspends.size() - 1;
246         }
247         break;
248       }
249       case Intrinsic::coro_begin: {
250         auto CB = cast<CoroBeginInst>(II);
251 
252         // Ignore coro id's that aren't pre-split.
253         auto Id = dyn_cast<CoroIdInst>(CB->getId());
254         if (Id && !Id->getInfo().isPreSplit())
255           break;
256 
257         if (CoroBegin)
258           report_fatal_error(
259                 "coroutine should have exactly one defining @llvm.coro.begin");
260         CB->addRetAttr(Attribute::NonNull);
261         CB->addRetAttr(Attribute::NoAlias);
262         CB->removeFnAttr(Attribute::NoDuplicate);
263         CoroBegin = CB;
264         break;
265       }
266       case Intrinsic::coro_end_async:
267       case Intrinsic::coro_end:
268         CoroEnds.push_back(cast<AnyCoroEndInst>(II));
269         if (auto *AsyncEnd = dyn_cast<CoroAsyncEndInst>(II)) {
270           AsyncEnd->checkWellFormed();
271         }
272 
273         if (CoroEnds.back()->isUnwind())
274           HasUnwindCoroEnd = true;
275 
276         if (CoroEnds.back()->isFallthrough() && isa<CoroEndInst>(II)) {
277           // Make sure that the fallthrough coro.end is the first element in the
278           // CoroEnds vector.
279           // Note: I don't think this is neccessary anymore.
280           if (CoroEnds.size() > 1) {
281             if (CoroEnds.front()->isFallthrough())
282               report_fatal_error(
283                   "Only one coro.end can be marked as fallthrough");
284             std::swap(CoroEnds.front(), CoroEnds.back());
285           }
286         }
287         break;
288       }
289     }
290   }
291 
292   // If there is no CoroBegin then this is not a coroutine.
293   if (!CoroBegin)
294     return;
295 
296   // Determination of ABI and initializing lowering info
297   auto Id = CoroBegin->getId();
298   switch (auto IntrID = Id->getIntrinsicID()) {
299   case Intrinsic::coro_id: {
300     ABI = coro::ABI::Switch;
301     SwitchLowering.HasFinalSuspend = HasFinalSuspend;
302     SwitchLowering.HasUnwindCoroEnd = HasUnwindCoroEnd;
303 
304     auto SwitchId = getSwitchCoroId();
305     SwitchLowering.ResumeSwitch = nullptr;
306     SwitchLowering.PromiseAlloca = SwitchId->getPromise();
307     SwitchLowering.ResumeEntryBlock = nullptr;
308 
309     // Move final suspend to the last element in the CoroSuspends vector.
310     if (SwitchLowering.HasFinalSuspend &&
311         FinalSuspendIndex != CoroSuspends.size() - 1)
312       std::swap(CoroSuspends[FinalSuspendIndex], CoroSuspends.back());
313     break;
314   }
315   case Intrinsic::coro_id_async: {
316     ABI = coro::ABI::Async;
317     auto *AsyncId = getAsyncCoroId();
318     AsyncId->checkWellFormed();
319     AsyncLowering.Context = AsyncId->getStorage();
320     AsyncLowering.ContextArgNo = AsyncId->getStorageArgumentIndex();
321     AsyncLowering.ContextHeaderSize = AsyncId->getStorageSize();
322     AsyncLowering.ContextAlignment = AsyncId->getStorageAlignment().value();
323     AsyncLowering.AsyncFuncPointer = AsyncId->getAsyncFunctionPointer();
324     AsyncLowering.AsyncCC = F.getCallingConv();
325     break;
326   }
327   case Intrinsic::coro_id_retcon:
328   case Intrinsic::coro_id_retcon_once: {
329     ABI = IntrID == Intrinsic::coro_id_retcon ? coro::ABI::Retcon
330                                               : coro::ABI::RetconOnce;
331     auto ContinuationId = getRetconCoroId();
332     ContinuationId->checkWellFormed();
333     auto Prototype = ContinuationId->getPrototype();
334     RetconLowering.ResumePrototype = Prototype;
335     RetconLowering.Alloc = ContinuationId->getAllocFunction();
336     RetconLowering.Dealloc = ContinuationId->getDeallocFunction();
337     RetconLowering.ReturnBlock = nullptr;
338     RetconLowering.IsFrameInlineInStorage = false;
339     break;
340   }
341   default:
342     llvm_unreachable("coro.begin is not dependent on a coro.id call");
343   }
344 }
345 
346 // If for some reason, we were not able to find coro.begin, bailout.
347 void coro::Shape::invalidateCoroutine(
348     Function &F, SmallVectorImpl<CoroFrameInst *> &CoroFrames) {
349   assert(!CoroBegin);
350   {
351     // Replace coro.frame which are supposed to be lowered to the result of
352     // coro.begin with undef.
353     auto *Undef = UndefValue::get(PointerType::get(F.getContext(), 0));
354     for (CoroFrameInst *CF : CoroFrames) {
355       CF->replaceAllUsesWith(Undef);
356       CF->eraseFromParent();
357     }
358     CoroFrames.clear();
359 
360     // Replace all coro.suspend with undef and remove related coro.saves if
361     // present.
362     for (AnyCoroSuspendInst *CS : CoroSuspends) {
363       CS->replaceAllUsesWith(UndefValue::get(CS->getType()));
364       CS->eraseFromParent();
365       if (auto *CoroSave = CS->getCoroSave())
366         CoroSave->eraseFromParent();
367     }
368     CoroSuspends.clear();
369 
370     // Replace all coro.ends with unreachable instruction.
371     for (AnyCoroEndInst *CE : CoroEnds)
372       changeToUnreachable(CE);
373   }
374 }
375 
376 // Perform semantic checking and initialization of the ABI
377 void coro::Shape::initABI() {
378   switch (ABI) {
379   case coro::ABI::Switch: {
380     for (auto *AnySuspend : CoroSuspends) {
381       auto Suspend = dyn_cast<CoroSuspendInst>(AnySuspend);
382       if (!Suspend) {
383 #ifndef NDEBUG
384         AnySuspend->dump();
385 #endif
386         report_fatal_error("coro.id must be paired with coro.suspend");
387       }
388 
389       if (!Suspend->getCoroSave())
390         createCoroSave(CoroBegin, Suspend);
391     }
392     break;
393   }
394   case coro::ABI::Async: {
395     break;
396   };
397   case coro::ABI::Retcon:
398   case coro::ABI::RetconOnce: {
399     // Determine the result value types, and make sure they match up with
400     // the values passed to the suspends.
401     auto ResultTys = getRetconResultTypes();
402     auto ResumeTys = getRetconResumeTypes();
403 
404     for (auto *AnySuspend : CoroSuspends) {
405       auto Suspend = dyn_cast<CoroSuspendRetconInst>(AnySuspend);
406       if (!Suspend) {
407 #ifndef NDEBUG
408         AnySuspend->dump();
409 #endif
410         report_fatal_error("coro.id.retcon.* must be paired with "
411                            "coro.suspend.retcon");
412       }
413 
414       // Check that the argument types of the suspend match the results.
415       auto SI = Suspend->value_begin(), SE = Suspend->value_end();
416       auto RI = ResultTys.begin(), RE = ResultTys.end();
417       for (; SI != SE && RI != RE; ++SI, ++RI) {
418         auto SrcTy = (*SI)->getType();
419         if (SrcTy != *RI) {
420           // The optimizer likes to eliminate bitcasts leading into variadic
421           // calls, but that messes with our invariants.  Re-insert the
422           // bitcast and ignore this type mismatch.
423           if (CastInst::isBitCastable(SrcTy, *RI)) {
424             auto BCI = new BitCastInst(*SI, *RI, "", Suspend->getIterator());
425             SI->set(BCI);
426             continue;
427           }
428 
429 #ifndef NDEBUG
430           Suspend->dump();
431           RetconLowering.ResumePrototype->getFunctionType()->dump();
432 #endif
433           report_fatal_error("argument to coro.suspend.retcon does not "
434                              "match corresponding prototype function result");
435         }
436       }
437       if (SI != SE || RI != RE) {
438 #ifndef NDEBUG
439         Suspend->dump();
440         RetconLowering.ResumePrototype->getFunctionType()->dump();
441 #endif
442         report_fatal_error("wrong number of arguments to coro.suspend.retcon");
443       }
444 
445       // Check that the result type of the suspend matches the resume types.
446       Type *SResultTy = Suspend->getType();
447       ArrayRef<Type *> SuspendResultTys;
448       if (SResultTy->isVoidTy()) {
449         // leave as empty array
450       } else if (auto SResultStructTy = dyn_cast<StructType>(SResultTy)) {
451         SuspendResultTys = SResultStructTy->elements();
452       } else {
453         // forms an ArrayRef using SResultTy, be careful
454         SuspendResultTys = SResultTy;
455       }
456       if (SuspendResultTys.size() != ResumeTys.size()) {
457 #ifndef NDEBUG
458         Suspend->dump();
459         RetconLowering.ResumePrototype->getFunctionType()->dump();
460 #endif
461         report_fatal_error("wrong number of results from coro.suspend.retcon");
462       }
463       for (size_t I = 0, E = ResumeTys.size(); I != E; ++I) {
464         if (SuspendResultTys[I] != ResumeTys[I]) {
465 #ifndef NDEBUG
466           Suspend->dump();
467           RetconLowering.ResumePrototype->getFunctionType()->dump();
468 #endif
469           report_fatal_error("result from coro.suspend.retcon does not "
470                              "match corresponding prototype function param");
471         }
472       }
473     }
474     break;
475   }
476   }
477 }
478 
479 void coro::Shape::cleanCoroutine(
480     SmallVectorImpl<CoroFrameInst *> &CoroFrames,
481     SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves) {
482   // The coro.frame intrinsic is always lowered to the result of coro.begin.
483   for (CoroFrameInst *CF : CoroFrames) {
484     CF->replaceAllUsesWith(CoroBegin);
485     CF->eraseFromParent();
486   }
487   CoroFrames.clear();
488 
489   // Remove orphaned coro.saves.
490   for (CoroSaveInst *CoroSave : UnusedCoroSaves)
491     CoroSave->eraseFromParent();
492   UnusedCoroSaves.clear();
493 }
494 
495 static void propagateCallAttrsFromCallee(CallInst *Call, Function *Callee) {
496   Call->setCallingConv(Callee->getCallingConv());
497   // TODO: attributes?
498 }
499 
500 static void addCallToCallGraph(CallGraph *CG, CallInst *Call, Function *Callee){
501   if (CG)
502     (*CG)[Call->getFunction()]->addCalledFunction(Call, (*CG)[Callee]);
503 }
504 
505 Value *coro::Shape::emitAlloc(IRBuilder<> &Builder, Value *Size,
506                               CallGraph *CG) const {
507   switch (ABI) {
508   case coro::ABI::Switch:
509     llvm_unreachable("can't allocate memory in coro switch-lowering");
510 
511   case coro::ABI::Retcon:
512   case coro::ABI::RetconOnce: {
513     auto Alloc = RetconLowering.Alloc;
514     Size = Builder.CreateIntCast(Size,
515                                  Alloc->getFunctionType()->getParamType(0),
516                                  /*is signed*/ false);
517     auto *Call = Builder.CreateCall(Alloc, Size);
518     propagateCallAttrsFromCallee(Call, Alloc);
519     addCallToCallGraph(CG, Call, Alloc);
520     return Call;
521   }
522   case coro::ABI::Async:
523     llvm_unreachable("can't allocate memory in coro async-lowering");
524   }
525   llvm_unreachable("Unknown coro::ABI enum");
526 }
527 
528 void coro::Shape::emitDealloc(IRBuilder<> &Builder, Value *Ptr,
529                               CallGraph *CG) const {
530   switch (ABI) {
531   case coro::ABI::Switch:
532     llvm_unreachable("can't allocate memory in coro switch-lowering");
533 
534   case coro::ABI::Retcon:
535   case coro::ABI::RetconOnce: {
536     auto Dealloc = RetconLowering.Dealloc;
537     Ptr = Builder.CreateBitCast(Ptr,
538                                 Dealloc->getFunctionType()->getParamType(0));
539     auto *Call = Builder.CreateCall(Dealloc, Ptr);
540     propagateCallAttrsFromCallee(Call, Dealloc);
541     addCallToCallGraph(CG, Call, Dealloc);
542     return;
543   }
544   case coro::ABI::Async:
545     llvm_unreachable("can't allocate memory in coro async-lowering");
546   }
547   llvm_unreachable("Unknown coro::ABI enum");
548 }
549 
550 [[noreturn]] static void fail(const Instruction *I, const char *Reason,
551                               Value *V) {
552 #ifndef NDEBUG
553   I->dump();
554   if (V) {
555     errs() << "  Value: ";
556     V->printAsOperand(llvm::errs());
557     errs() << '\n';
558   }
559 #endif
560   report_fatal_error(Reason);
561 }
562 
563 /// Check that the given value is a well-formed prototype for the
564 /// llvm.coro.id.retcon.* intrinsics.
565 static void checkWFRetconPrototype(const AnyCoroIdRetconInst *I, Value *V) {
566   auto F = dyn_cast<Function>(V->stripPointerCasts());
567   if (!F)
568     fail(I, "llvm.coro.id.retcon.* prototype not a Function", V);
569 
570   auto FT = F->getFunctionType();
571 
572   if (isa<CoroIdRetconInst>(I)) {
573     bool ResultOkay;
574     if (FT->getReturnType()->isPointerTy()) {
575       ResultOkay = true;
576     } else if (auto SRetTy = dyn_cast<StructType>(FT->getReturnType())) {
577       ResultOkay = (!SRetTy->isOpaque() &&
578                     SRetTy->getNumElements() > 0 &&
579                     SRetTy->getElementType(0)->isPointerTy());
580     } else {
581       ResultOkay = false;
582     }
583     if (!ResultOkay)
584       fail(I, "llvm.coro.id.retcon prototype must return pointer as first "
585               "result", F);
586 
587     if (FT->getReturnType() !=
588           I->getFunction()->getFunctionType()->getReturnType())
589       fail(I, "llvm.coro.id.retcon prototype return type must be same as"
590               "current function return type", F);
591   } else {
592     // No meaningful validation to do here for llvm.coro.id.unique.once.
593   }
594 
595   if (FT->getNumParams() == 0 || !FT->getParamType(0)->isPointerTy())
596     fail(I, "llvm.coro.id.retcon.* prototype must take pointer as "
597             "its first parameter", F);
598 }
599 
600 /// Check that the given value is a well-formed allocator.
601 static void checkWFAlloc(const Instruction *I, Value *V) {
602   auto F = dyn_cast<Function>(V->stripPointerCasts());
603   if (!F)
604     fail(I, "llvm.coro.* allocator not a Function", V);
605 
606   auto FT = F->getFunctionType();
607   if (!FT->getReturnType()->isPointerTy())
608     fail(I, "llvm.coro.* allocator must return a pointer", F);
609 
610   if (FT->getNumParams() != 1 ||
611       !FT->getParamType(0)->isIntegerTy())
612     fail(I, "llvm.coro.* allocator must take integer as only param", F);
613 }
614 
615 /// Check that the given value is a well-formed deallocator.
616 static void checkWFDealloc(const Instruction *I, Value *V) {
617   auto F = dyn_cast<Function>(V->stripPointerCasts());
618   if (!F)
619     fail(I, "llvm.coro.* deallocator not a Function", V);
620 
621   auto FT = F->getFunctionType();
622   if (!FT->getReturnType()->isVoidTy())
623     fail(I, "llvm.coro.* deallocator must return void", F);
624 
625   if (FT->getNumParams() != 1 ||
626       !FT->getParamType(0)->isPointerTy())
627     fail(I, "llvm.coro.* deallocator must take pointer as only param", F);
628 }
629 
630 static void checkConstantInt(const Instruction *I, Value *V,
631                              const char *Reason) {
632   if (!isa<ConstantInt>(V)) {
633     fail(I, Reason, V);
634   }
635 }
636 
637 void AnyCoroIdRetconInst::checkWellFormed() const {
638   checkConstantInt(this, getArgOperand(SizeArg),
639                    "size argument to coro.id.retcon.* must be constant");
640   checkConstantInt(this, getArgOperand(AlignArg),
641                    "alignment argument to coro.id.retcon.* must be constant");
642   checkWFRetconPrototype(this, getArgOperand(PrototypeArg));
643   checkWFAlloc(this, getArgOperand(AllocArg));
644   checkWFDealloc(this, getArgOperand(DeallocArg));
645 }
646 
647 static void checkAsyncFuncPointer(const Instruction *I, Value *V) {
648   auto *AsyncFuncPtrAddr = dyn_cast<GlobalVariable>(V->stripPointerCasts());
649   if (!AsyncFuncPtrAddr)
650     fail(I, "llvm.coro.id.async async function pointer not a global", V);
651 }
652 
653 void CoroIdAsyncInst::checkWellFormed() const {
654   checkConstantInt(this, getArgOperand(SizeArg),
655                    "size argument to coro.id.async must be constant");
656   checkConstantInt(this, getArgOperand(AlignArg),
657                    "alignment argument to coro.id.async must be constant");
658   checkConstantInt(this, getArgOperand(StorageArg),
659                    "storage argument offset to coro.id.async must be constant");
660   checkAsyncFuncPointer(this, getArgOperand(AsyncFuncPtrArg));
661 }
662 
663 static void checkAsyncContextProjectFunction(const Instruction *I,
664                                              Function *F) {
665   auto *FunTy = cast<FunctionType>(F->getValueType());
666   if (!FunTy->getReturnType()->isPointerTy())
667     fail(I,
668          "llvm.coro.suspend.async resume function projection function must "
669          "return a ptr type",
670          F);
671   if (FunTy->getNumParams() != 1 || !FunTy->getParamType(0)->isPointerTy())
672     fail(I,
673          "llvm.coro.suspend.async resume function projection function must "
674          "take one ptr type as parameter",
675          F);
676 }
677 
678 void CoroSuspendAsyncInst::checkWellFormed() const {
679   checkAsyncContextProjectFunction(this, getAsyncContextProjectionFunction());
680 }
681 
682 void CoroAsyncEndInst::checkWellFormed() const {
683   auto *MustTailCallFunc = getMustTailCallFunction();
684   if (!MustTailCallFunc)
685     return;
686   auto *FnTy = MustTailCallFunc->getFunctionType();
687   if (FnTy->getNumParams() != (arg_size() - 3))
688     fail(this,
689          "llvm.coro.end.async must tail call function argument type must "
690          "match the tail arguments",
691          MustTailCallFunc);
692 }
693