Get rid of getOrCreateAwaitAsyncSupendFn

Author: aschwaighofer

lib/IRGen/GenCall.cpp

@@ -4581,3 +4581,17 @@ void irgen::emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &asyncLayout,
   auto call = IGF.Builder.CreateCall(fnPtr, Args);
   call->setTailCall();
 }
+
+FunctionPointer
+IRGenFunction::getFunctionPointerForResumeIntrinsic(llvm::Value *resume) {
+  auto *fnTy = llvm::FunctionType::get(
+      IGM.VoidTy, {IGM.Int8PtrTy, IGM.Int8PtrTy, IGM.Int8PtrTy},
+      false /*vaargs*/);
+  auto signature =
+      Signature(fnTy, IGM.constructInitialAttributes(), IGM.SwiftCC);
+  auto fnPtr = FunctionPointer(
+      FunctionPointer::KindTy::Function,
+      Builder.CreateBitOrPointerCast(resume, fnTy->getPointerTo()),
+      PointerAuthInfo(), signature);
+  return fnPtr;
+}

lib/IRGen/GenFunc.cpp

@@ -2395,21 +2395,31 @@ llvm::Function *IRGenFunction::getOrCreateResumePrjFn() {
       },
       false /*isNoInline*/));
 }
-
 llvm::Function *
 IRGenFunction::createAsyncDispatchFn(const FunctionPointer &fnPtr,
                                      ArrayRef<llvm::Value *> args) {
   SmallVector<llvm::Type*, 8> argTys;
-  argTys.push_back(IGM.Int8PtrTy); // Function pointer to be called.
   for (auto arg : args) {
     auto *ty = arg->getType();
     argTys.push_back(ty);
   }
+  return createAsyncDispatchFn(fnPtr, argTys);
+}
+
+llvm::Function *
+IRGenFunction::createAsyncDispatchFn(const FunctionPointer &fnPtr,
+                                     ArrayRef<llvm::Type *> argTypes) {
+  SmallVector<llvm::Type*, 8> argTys;
+  argTys.push_back(IGM.Int8PtrTy); // Function pointer to be called.
+  for (auto ty : argTypes) {
+    argTys.push_back(ty);
+  }
   auto calleeFnPtrType = fnPtr.getRawPointer()->getType();
   auto *dispatchFnTy =
       llvm::FunctionType::get(IGM.VoidTy, argTys, false /*vaargs*/);
   llvm::SmallString<40> name;
-  llvm::raw_svector_ostream(name) << "__swift_suspend_dispatch_" << args.size();
+  llvm::raw_svector_ostream(name)
+      << "__swift_suspend_dispatch_" << argTypes.size();
   llvm::Function *dispatch =
       llvm::Function::Create(dispatchFnTy, llvm::Function::InternalLinkage,
                              llvm::StringRef(name), &IGM.Module);
@@ -2433,31 +2443,3 @@ IRGenFunction::createAsyncDispatchFn(const FunctionPointer &fnPtr,
   Builder.CreateRetVoid();
   return dispatch;
 }
-
-llvm::Function *IRGenFunction::getOrCreateAwaitAsyncSupendFn() {
-  auto name = "__swift_async_await_async_suspend";
-  return cast<llvm::Function>(IGM.getOrCreateHelperFunction(
-      name, IGM.VoidTy, {IGM.Int8PtrTy, IGM.Int8PtrTy, IGM.Int8PtrTy, IGM.Int8PtrTy},
-      [&](IRGenFunction &IGF) {
-        auto it = IGF.CurFn->arg_begin();
-        auto &Builder = IGF.Builder;
-        auto *fnTy = llvm::FunctionType::get(
-            IGM.VoidTy, {IGM.Int8PtrTy, IGM.Int8PtrTy, IGM.Int8PtrTy},
-            false /*vaargs*/);
-        llvm::Value *fn = &*(it++);
-        SmallVector<llvm::Value *, 8> callArgs;
-        for (auto end = IGF.CurFn->arg_end(); it != end; ++it)
-          callArgs.push_back(&*it);
-        auto signature =
-            Signature(fnTy, IGM.constructInitialAttributes(), IGM.SwiftCC);
-        auto fnPtr = FunctionPointer(
-            FunctionPointer::KindTy::Function,
-            Builder.CreateBitOrPointerCast(fn, fnTy->getPointerTo()),
-            PointerAuthInfo(), signature);
-        auto call = Builder.CreateCall(fnPtr, callArgs);
-        call->setTailCall();
-        call->setCallingConv(IGM.SwiftCC);
-        Builder.CreateRetVoid();
-      },
-      false /*isNoInline*/));
-}

lib/IRGen/IRGenFunction.cpp

@@ -23,6 +23,7 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include "Callee.h"
 #include "Explosion.h"
 #include "IRGenDebugInfo.h"
 #include "IRGenFunction.h"
@@ -697,18 +698,22 @@ void IRGenFunction::emitAwaitAsyncContinuation(
     auto resumeProjFn = getOrCreateResumePrjFn();
     arguments.push_back(
         Builder.CreateBitOrPointerCast(resumeProjFn, IGM.Int8PtrTy));
+    // The dispatch function just calls the resume point.
+    auto resumeFnPtr =
+        getFunctionPointerForResumeIntrinsic(AsyncCoroutineCurrentResume);
     arguments.push_back(Builder.CreateBitOrPointerCast(
-        getOrCreateAwaitAsyncSupendFn(), IGM.Int8PtrTy));
+        createAsyncDispatchFn(resumeFnPtr,
+                              {IGM.Int8PtrTy, IGM.Int8PtrTy, IGM.Int8PtrTy}),
+        IGM.Int8PtrTy));
     arguments.push_back(AsyncCoroutineCurrentResume);
     arguments.push_back(
         Builder.CreateBitOrPointerCast(getAsyncTask(), IGM.Int8PtrTy));
     arguments.push_back(
         Builder.CreateBitOrPointerCast(getAsyncExecutor(), IGM.Int8PtrTy));
     arguments.push_back(Builder.CreateBitOrPointerCast(
         AsyncCoroutineCurrentContinuationContext, IGM.Int8PtrTy));
-    auto *id = Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_suspend_async,
+    Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_suspend_async, arguments);
 
-                                           arguments);
     auto results = Builder.CreateAtomicCmpXchg(
         contAwaitSyncAddr, null, one,
         llvm::AtomicOrdering::Release /*success ordering*/,

lib/IRGen/IRGenFunction.h

@@ -135,7 +135,8 @@ class IRGenFunction {
   llvm::Function *getOrCreateResumePrjFn();
   llvm::Function *createAsyncDispatchFn(const FunctionPointer &fnPtr,
                                         ArrayRef<llvm::Value *> args);
-  llvm::Function *getOrCreateAwaitAsyncSupendFn();
+  llvm::Function *createAsyncDispatchFn(const FunctionPointer &fnPtr,
+                                        ArrayRef<llvm::Type *> argTypes);
 
   void emitGetAsyncContinuation(SILType silTy, StackAddress optionalResultAddr,
                                 Explosion &out);
@@ -147,6 +148,9 @@ class IRGenFunction {
                                   llvm::PHINode *&optionalErrorPhi,
                                   llvm::BasicBlock *&optionalErrorBB);
 
+  FunctionPointer
+  getFunctionPointerForResumeIntrinsic(llvm::Value *resumeIntrinsic);
+
 private:
   void emitPrologue();
   void emitEpilogue();