[CoroutineAccessors] Use typed-malloc.

include/swift/ABI/Coro.h

@@ -30,6 +30,8 @@ enum class CoroAllocatorKind : uint8_t {
   Async = 1,
   // malloc/free
   Malloc = 2,
+  // swift_coroFrameAlloc/free
+  TypedMalloc = 3,
 };
 
 class CoroAllocatorFlags : public FlagSet<uint32_t> {

include/swift/Runtime/RuntimeFunctions.def

@@ -2100,12 +2100,6 @@ FUNCTION(CoroFrameAlloc, Swift, swift_coroFrameAlloc, C_CC,
          NO_ATTRS,
          EFFECT(RuntimeEffect::Allocating),
          UNKNOWN_MEMEFFECTS)
-FUNCTION(coroFrameAllocStub, Swift, swift_coroFrameAllocStub, C_CC,
-         AlwaysAvailable, RETURNS(Int8PtrTy),
-         ARGS(SizeTy, Int64Ty),
-         ATTRS(NoUnwind),
-         EFFECT(RuntimeEffect::Allocating),
-         UNKNOWN_MEMEFFECTS)
 
 // void *_Block_copy(void *block);
 FUNCTION(BlockCopy, BlocksRuntime, _Block_copy, C_CC, AlwaysAvailable,

lib/IRGen/GenCall.cpp

@@ -221,7 +221,11 @@ IRGenFunction::getDefaultCoroutineAllocatorKind() {
     return CoroAllocatorKind::Async;
   }
   if (isCoroutine()) {
-    return CoroAllocatorKind::Malloc;
+    if (getOptions().EmitTypeMallocForCoroFrame) {
+      return CoroAllocatorKind::TypedMalloc;
+    } else {
+      return CoroAllocatorKind::Malloc;
+    }
   }
   if (IGM.SwiftCoroCC != llvm::CallingConv::SwiftCoro) {
     // If the swiftcorocc isn't available, fall back to malloc.
@@ -2491,6 +2495,12 @@ SignatureExpansionABIDetails Signature::getUncachedABIDetails(
   return result;
 }
 
+Signature Signature::forFunction(llvm::Function *fn) {
+  auto sig = Signature(fn->getFunctionType(), fn->getAttributes(),
+                       fn->getCallingConv());
+  return sig;
+}
+
 Signature Signature::forCoroutineContinuation(IRGenModule &IGM,
                                               CanSILFunctionType fnType) {
   assert(fnType->isCoroutine());
@@ -2701,17 +2711,14 @@ irgen::getAsyncFunctionAndSize(IRGenFunction &IGF,
   return {fn, size};
 }
 
-std::pair<llvm::Value *, llvm::Value *>
-irgen::getCoroFunctionAndSize(IRGenFunction &IGF,
-                              FunctionPointer functionPointer,
-                              std::pair<bool, bool> values) {
-  assert(values.first || values.second);
+std::tuple<llvm::Value *, llvm::Value *, llvm::Value *>
+irgen::getCoroFunctionValues(IRGenFunction &IGF,
+                             FunctionPointer functionPointer,
+                             std::tuple<bool, bool, bool> values) {
+  auto [emitFunction, emitSize, emitTypeID] = values;
+  assert(emitFunction || emitSize || emitTypeID);
   assert(functionPointer.getKind() != FunctionPointer::Kind::Function);
 
-  bool emitFunction = values.first;
-  bool emitSize = values.second;
-  assert(emitFunction || emitSize);
-
   // Ensure that the CoroFunctionPointer is not auth'd if it is not used and
   // that it is not auth'd more than once if it is needed.
   //
@@ -2764,7 +2771,15 @@ irgen::getCoroFunctionAndSize(IRGenFunction &IGF,
     size = IGF.Builder.CreateLoad(sizePtr, IGF.IGM.Int32Ty,
                                   IGF.IGM.getPointerAlignment());
   }
-  return {fn, size};
+
+  llvm::Value *typeID = nullptr;
+  if (emitTypeID) {
+    auto *sizePtr = IGF.Builder.CreateStructGEP(IGF.IGM.CoroFunctionPointerTy,
+                                                getCoroPtr(), 2);
+    typeID = IGF.Builder.CreateLoad(sizePtr, IGF.IGM.Int64Ty,
+                                    IGF.IGM.getPointerAlignment());
+  }
+  return {fn, size, typeID};
 }
 
 namespace {
@@ -2804,16 +2819,17 @@ class SyncCallEmission final : public CallEmission {
     assert(!coroAllocator);
 
     if (IsCalleeAllocatedCoroutine) {
-      llvm::Value *bufferSize32;
-      std::tie(calleeFunction, bufferSize32) =
-          getCoroFunctionAndSize(IGF, CurCallee.getFunctionPointer());
+      auto kind = IGF.getDefaultCoroutineAllocatorKind();
+      llvm::Value *bufferSize32, *mallocTypeId;
+      std::tie(calleeFunction, bufferSize32, mallocTypeId) =
+          getCoroFunctionValues(IGF, CurCallee.getFunctionPointer());
       auto *bufferSize = IGF.Builder.CreateZExt(bufferSize32, IGF.IGM.SizeTy);
-      coroStaticFrame = emitAllocYieldOnce2CoroutineFrame(IGF, bufferSize);
+      coroStaticFrame =
+          emitAllocYieldOnce2CoroutineFrame(IGF, bufferSize, mallocTypeId);
       // TODO: CoroutineAccessors: Optimize allocator kind (e.g. async callers
       //                           only need to use the TaskAllocator if the
       //                           coroutine is suspended across an await).
-      coroAllocator = emitYieldOnce2CoroutineAllocator(
-          IGF, IGF.getDefaultCoroutineAllocatorKind());
+      coroAllocator = emitYieldOnce2CoroutineAllocator(IGF, kind);
     }
   }
   void end() override { super::end(); }
@@ -5179,10 +5195,12 @@ void irgen::emitYieldOnce2CoroutineEntry(IRGenFunction &IGF,
   auto deallocFn = IGF.IGM.getOpaquePtr(getCoroDeallocFn(IGF.IGM));
   auto allocFrameFn = IGF.IGM.getOpaquePtr(getCoroAllocFrameFn(IGF.IGM));
   auto deallocFrameFn = IGF.IGM.getOpaquePtr(getCoroDeallocFrameFn(IGF.IGM));
+  auto *typeID = IGF.getMallocTypeId();
   emitRetconCoroutineEntry(
       IGF, fnType, buffer, llvm::Intrinsic::coro_id_retcon_once_dynamic,
       Size(-1) /*dynamic-to-IRGen size*/, IGF.IGM.getCoroStaticFrameAlignment(),
-      {cfp, allocator}, allocFn, deallocFn, {allocFrameFn, deallocFrameFn});
+      {cfp, allocator}, allocFn, deallocFn,
+      {allocFrameFn, deallocFrameFn, typeID});
 }
 void irgen::emitYieldOnce2CoroutineEntry(
     IRGenFunction &IGF, LinkEntity coroFunction, CanSILFunctionType fnType,
@@ -5213,9 +5231,10 @@ Address irgen::emitAllocYieldManyCoroutineBuffer(IRGenFunction &IGF) {
                                  getYieldManyCoroutineBufferAlignment(IGF.IGM));
 }
 
-StackAddress irgen::emitAllocYieldOnce2CoroutineFrame(IRGenFunction &IGF,
-                                                      llvm::Value *size) {
-  return emitAllocCoroStaticFrame(IGF, size);
+StackAddress
+irgen::emitAllocYieldOnce2CoroutineFrame(IRGenFunction &IGF, llvm::Value *size,
+                                         llvm::Value *mallocTypeId) {
+  return emitAllocCoroStaticFrame(IGF, size, mallocTypeId);
 }
 
 void irgen::emitDeallocYieldOnceCoroutineBuffer(IRGenFunction &IGF,
@@ -5264,14 +5283,15 @@ void irgen::emitStaticDeallocAsyncContext(IRGenFunction &IGF, Address context,
 }
 
 StackAddress irgen::emitAllocCoroStaticFrame(IRGenFunction &IGF,
-                                             llvm::Value *size) {
+                                             llvm::Value *size,
+                                             llvm::Value *mallocTypeId) {
+  ASSERT(mallocTypeId);
   // TODO: Avoid swift_task_alloc (async) and malloc (yield_once) if the
   //       suspension doesn't span an apply of an async function or a yield
   //       respectively.
-  auto retval = IGF.emitDynamicAlloca(
-      IGF.IGM.Int8Ty, size, Alignment(MaximumAlignment), AllowsTaskAlloc,
-      IsForCalleeCoroutineFrame_t(IGF.isCalleeAllocatedCoroutine()),
-      "callee-coro-frame");
+  auto retval =
+      IGF.emitDynamicAlloca(IGF.IGM.Int8Ty, size, Alignment(MaximumAlignment),
+                            AllowsTaskAlloc, mallocTypeId, "callee-coro-frame");
   IGF.Builder.CreateLifetimeStart(retval.getAddress(),
                                   Size(-1) /*dynamic size*/);
   return retval;

lib/IRGen/GenCall.h

@@ -143,9 +143,9 @@ namespace irgen {
   std::pair<llvm::Value *, llvm::Value *>
   getAsyncFunctionAndSize(IRGenFunction &IGF, FunctionPointer functionPointer,
                           std::pair<bool, bool> values = {true, true});
-  std::pair<llvm::Value *, llvm::Value *>
-  getCoroFunctionAndSize(IRGenFunction &IGF, FunctionPointer functionPointer,
-                         std::pair<bool, bool> values = {true, true});
+  std::tuple<llvm::Value *, llvm::Value *, llvm::Value *> getCoroFunctionValues(
+      IRGenFunction &IGF, FunctionPointer functionPointer,
+      std::tuple<bool, bool, bool> values = {true, true, true});
   llvm::CallingConv::ID
   expandCallingConv(IRGenModule &IGM, SILFunctionTypeRepresentation convention,
                     bool isAsync, bool isCalleeAllocatedCoro);
@@ -221,7 +221,8 @@ namespace irgen {
                               NativeCCEntryPointArgumentEmission &emission);
 
   StackAddress emitAllocYieldOnce2CoroutineFrame(IRGenFunction &IGF,
-                                                 llvm::Value *size);
+                                                 llvm::Value *size,
+                                                 llvm::Value *mallocTypeId);
   void emitDeallocYieldOnce2CoroutineFrame(IRGenFunction &IGF,
                                            StackAddress allocation);
   void
@@ -252,7 +253,8 @@ namespace irgen {
                               LinkEntity asyncFunction,
                               unsigned asyncContextIndex);
 
-  StackAddress emitAllocCoroStaticFrame(IRGenFunction &IGF, llvm::Value *size);
+  StackAddress emitAllocCoroStaticFrame(IRGenFunction &IGF, llvm::Value *size,
+                                        llvm::Value *mallocTypeId);
   void emitDeallocCoroStaticFrame(IRGenFunction &IGF, StackAddress frame);
 
   /// Yield the given values from the current continuation.