[coro] Async coroutines: Make sure we can handle control flow in suspend point dispatch function

[aschwaighofer]

Summary

Create a valid basic block with a terminator before we call
InlineFunction.

[eeckstein]

LGTM!

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

for (size_t Idx = 0, End = Shape.CoroSuspends.size(); Idx != End; ++Idx) {
     // Coerce the arguments, llvm optimizations seem to ignore the types in
     // vaarg functions and throws away casts in optimized mode.
     SmallVector<Value *, 8> CallArgs;
     coerceArguments(Builder, FnTy, FnArgs, CallArgs);
     auto *TailCall = Builder.CreateCall(FnTy, Fn, CallArgs);
     TailCall->setDebugLoc(DbgLoc);
     TailCall->setTailCall();
     TailCall->setCallingConv(Fn->getCallingConv());
+    Builder.CreateRetVoid();
     InlineFunctionInfo FnInfo;
     auto InlineRes = InlineFunction(*TailCall, FnInfo);
     assert(InlineRes.isSuccess() && "Expected inlining to succeed");
     (void)InlineRes;
-    Builder.CreateRetVoid();
 
     // Replace the lvm.coro.async.resume intrisic call.
     replaceAsyncResumeFunction(Suspend, Continuation);
   }
 
   assert(Clones.size() == Shape.CoroSuspends.size());
   for (size_t Idx = 0, End = Shape.CoroSuspends.size(); Idx != End; ++Idx) {
     auto *Suspend = Shape.CoroSuspends[Idx];

llvm/test/Transforms/Coroutines/coro-async.ll

 }
 
 ; CHECK-LABEL: define swiftcc void @top_level_caller(i8* %ctxt, i8* %task, i8* %actor)
 ; CHECK: store i8* bitcast (void (i8*, i8*, i8*)* @my_async_function.resume.0
 ; CHECK: store i8* %ctxt
 ; CHECK: tail call swiftcc void @asyncSuspend
 ; CHECK: ret void
 
+@dont_crash_on_cf_fp = constant <{ i32, i32 }>
+  <{ i32 trunc ( ; Relative pointer to async function
+       i64 sub (
+         i64 ptrtoint (void (i8*, %async.task*, %async.actor*)* @my_async_function to i64),
+         i64 ptrtoint (i32* getelementptr inbounds (<{ i32, i32 }>, <{ i32, i32 }>* @my_async_function_fp, i32 0, i32 1) to i64)
+       )
+     to i32),
+     i32 128    ; Initial async context size without space for frame
+}>
+
+
+define swiftcc void @dont_crash_on_cf_dispatch(i8* %fnPtr, i8* %async.ctxt, %async.task* %task, %async.actor* %actor) {
+  %isNull = icmp eq %async.task* %task, null
+  br i1 %isNull, label %is_null, label %is_not_null
+
+is_null:
+  ret void
+
+is_not_null:
+  %callee = bitcast i8* %fnPtr to void(i8*, %async.task*, %async.actor*)*
+  tail call swiftcc void %callee(i8* %async.ctxt, %async.task* %task, %async.actor* %actor)
+  ret void
+}
+
+define swiftcc void @dont_crash_on_cf(i8* %async.ctxt, %async.task* %task, %async.actor* %actor)  {
+entry:
+  %id = call token @llvm.coro.id.async(i32 128, i32 16, i32 0,
+          i8* bitcast (<{i32, i32}>* @dont_crash_on_cf_fp to i8*))
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
+  %arg0 = bitcast %async.task* %task to i8*
+  %arg1 = bitcast <{ i32, i32}>* @my_other_async_function_fp to i8*
+  %callee_context = call i8* @llvm.coro.async.context.alloc(i8* %arg0, i8* %arg1)
+  %callee_context.0 = bitcast i8* %callee_context to %async.ctxt*
+  %callee_context.return_to_caller.addr = getelementptr inbounds %async.ctxt, %async.ctxt* %callee_context.0, i32 0, i32 1
+  %return_to_caller.addr = bitcast void(i8*, %async.task*, %async.actor*)** %callee_context.return_to_caller.addr to i8**
+  %resume.func_ptr = call i8* @llvm.coro.async.resume()
+  store i8* %resume.func_ptr, i8** %return_to_caller.addr
+  %callee_context.caller_context.addr = getelementptr inbounds %async.ctxt, %async.ctxt* %callee_context.0, i32 0, i32 0
+  store i8* %async.ctxt, i8** %callee_context.caller_context.addr
+  %resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
+  %callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
+  %res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
+                                                  i8* %resume.func_ptr,
+                                                  i8* %resume_proj_fun,
+                                                  void (i8*, i8*, %async.task*, %async.actor*)* @dont_crash_on_cf_dispatch,
+                                                  i8* %callee, i8* %callee_context, %async.task* %task, %async.actor *%actor)
+
+  call void @llvm.coro.async.context.dealloc(i8* %callee_context)
+  %continuation_task_arg = extractvalue {i8*, i8*, i8*} %res, 1
+  %task.2 =  bitcast i8* %continuation_task_arg to %async.task*
+  tail call swiftcc void @asyncReturn(i8* %async.ctxt, %async.task* %task.2, %async.actor* %actor)
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  unreachable
+}
 declare i8* @llvm.coro.prepare.async(i8*)
 declare token @llvm.coro.id.async(i32, i32, i32, i8*)
 declare i8* @llvm.coro.begin(token, i8*)
 declare i1 @llvm.coro.end(i8*, i1)
 declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i8*, i8*, ...)
 declare i8* @llvm.coro.async.context.alloc(i8*, i8*)
 declare void @llvm.coro.async.context.dealloc(i8*)
 declare swiftcc void @asyncReturn(i8*, %async.task*, %async.actor*)
 declare swiftcc void @asyncSuspend(i8*, %async.task*, %async.actor*)
 declare i8* @llvm.coro.async.resume()