[CoroutineAccessors] Open code dealloc fn.

include/swift/Runtime/Coro.h

@@ -21,9 +21,6 @@
 #include "swift/Runtime/Config.h"
 #include <cstddef>
 
-namespace swift {
+namespace swift {} // end namespace swift
-SWIFT_RUNTIME_EXPORT
-SWIFT_CC(swift) void swift_coro_dealloc(CoroAllocator *allocator, void *ptr);
-} // end namespace swift
 
 #endif

include/swift/Runtime/RuntimeFunctions.def

@@ -3008,16 +3008,6 @@ FUNCTION(MemsetS, c, memset_s, C_CC, AlwaysAvailable,
          EFFECT(RuntimeEffect::NoEffect),
          UNKNOWN_MEMEFFECTS)
 
-// void swift_coro_dealloc(CoroAllocator *, void *ptr);
-FUNCTION(CoroDealloc,
-         Swift, swift_coro_dealloc, SwiftCC,
-         CoroutineAccessorsAvailability,
-         RETURNS(VoidTy),
-         ARGS(CoroAllocatorPtrTy, Int8PtrTy),
-         NO_ATTRS,
-         EFFECT(RuntimeEffect::Deallocating, RuntimeEffect::Concurrency),
-         UNKNOWN_MEMEFFECTS)
-
 #undef RETURNS
 #undef ARGS
 #undef ATTRS

lib/IRGen/GenCall.cpp

@@ -5173,33 +5173,78 @@ static llvm::Constant *getCoroAllocFn(IRGenModule &IGM) {
       /*optionalLinkageOverride=*/nullptr, IGM.SwiftCoroCC);
 }
 
-static llvm::Constant *getCoroDeallocWrapperFn(IRGenModule &IGM) {
+static llvm::Constant *getCoroDeallocFn(IRGenModule &IGM) {
+  auto isSwiftCoroCCAvailable = IGM.SwiftCoroCC == llvm::CallingConv::SwiftCoro;
   return IGM.getOrCreateHelperFunction(
-      "__swift_coro_dealloc_", IGM.VoidTy,
+      "_swift_coro_dealloc", IGM.VoidTy,
       {IGM.CoroAllocatorPtrTy, IGM.Int8PtrTy},
-      [](IRGenFunction &IGF) {
+      [isSwiftCoroCCAvailable](IRGenFunction &IGF) {
         auto parameters = IGF.collectParameters();
         auto *allocator = parameters.claimNext();
         auto *ptr = parameters.claimNext();
-        auto *nullAllocator = IGF.Builder.CreateCmp(
+        if (isSwiftCoroCCAvailable) {
-            llvm::CmpInst::Predicate::ICMP_EQ, allocator,
+          // swiftcorocc is available, so if there's no allocator pointer,
-            llvm::ConstantPointerNull::get(
+          // storage was allocated on the stack which will be naturally cleaned
-                cast<llvm::PointerType>(allocator->getType())));
+          // up when the coroutine's frame is "freed".
-        auto *bailBlock = IGF.createBasicBlock("bail");
+          auto *nullAllocator = IGF.Builder.CreateCmp(
-        auto *forwardBlock = IGF.createBasicBlock("forward");
+              llvm::CmpInst::Predicate::ICMP_EQ, allocator,
-        IGF.Builder.CreateCondBr(nullAllocator, bailBlock, forwardBlock);
+              llvm::ConstantPointerNull::get(
-        IGF.Builder.emitBlock(bailBlock);
+                  cast<llvm::PointerType>(allocator->getType())));
-        // Emit the dynamic alloca.
+          auto *bailBlock = IGF.createBasicBlock("null_allocator");
+          auto *normalBlock = IGF.createBasicBlock("nonnull_allocator");
+          IGF.Builder.CreateCondBr(nullAllocator, bailBlock, normalBlock);
+          IGF.Builder.emitBlock(bailBlock);
+          // Nothing to do here.
+          IGF.Builder.CreateRetVoid();
+          // Start emitting the "normal" block.
+          IGF.Builder.emitBlock(normalBlock);
+        }
+        auto shouldDeallocateImmediatelyFlag = CoroAllocatorFlags(0);
+        shouldDeallocateImmediatelyFlag.setShouldDeallocateImmediately(true);
+        auto *flagsPtr = IGF.Builder.CreateInBoundsGEP(
+            IGF.IGM.CoroAllocatorTy, allocator,
+            {llvm::ConstantInt::get(IGF.IGM.Int32Ty, 0),
+             llvm::ConstantInt::get(IGF.IGM.Int32Ty, 0)});
+        auto *flags = IGF.Builder.CreateLoad(
+            Address(flagsPtr, IGF.IGM.Int32Ty, Alignment(4)), "");
+        auto *deallocDeferringAllocator = IGF.Builder.CreateAnd(
+            flags,
+            llvm::APInt(IGF.IGM.Int32Ty->getBitWidth(),
+                        shouldDeallocateImmediatelyFlag.getOpaqueValue()));
+        auto *isDeallocDeferringAllocator = IGF.Builder.CreateICmpNE(
+            deallocDeferringAllocator,
+            llvm::ConstantInt::get(IGF.IGM.Int32Ty, 0));
+        auto *deferringAllocatorBlock =
+            IGF.createBasicBlock("deferring_allocator");
+        auto *normalBlock = IGF.createBasicBlock("normal");
+        IGF.Builder.CreateCondBr(isDeallocDeferringAllocator,
+                                 deferringAllocatorBlock, normalBlock);
+        IGF.Builder.emitBlock(deferringAllocatorBlock);
+        // Nothing to do here.
         IGF.Builder.CreateRetVoid();
-        IGF.Builder.emitBlock(forwardBlock);
+        // Start emitting the "normal" block.
-        IGF.Builder.CreateCall(
+        IGF.Builder.emitBlock(normalBlock);
-            IGF.IGM.getCoroDeallocFunctionPointer(), {allocator, ptr});
+        auto *calleePtr = IGF.Builder.CreateInBoundsGEP(
+            IGF.IGM.CoroAllocatorTy, allocator,
+            {llvm::ConstantInt::get(IGF.IGM.Int32Ty, 0),
+             llvm::ConstantInt::get(IGF.IGM.Int32Ty, 2)});
+        auto *callee = IGF.Builder.CreateLoad(
+            Address(calleePtr, IGF.IGM.CoroDeallocateFnTy->getPointerTo(),
+                    IGF.IGM.getPointerAlignment()),
+            "deallocate_fn");
+        auto fnPtr = FunctionPointer::createUnsigned(
+            FunctionPointer::Kind::Function, callee,
+            Signature(cast<llvm::FunctionType>(IGF.IGM.CoroDeallocateFnTy), {},
+                      IGF.IGM.SwiftCC));
+        auto *call = IGF.Builder.CreateCall(fnPtr, {ptr});
+        call->setDoesNotThrow();
+        call->setCallingConv(IGF.IGM.SwiftCC);
         IGF.Builder.CreateRetVoid();
       },
-      /*setIsNoInline=*/false,
+      /*setIsNoInline=*/true,
       /*forPrologue=*/false,
       /*isPerformanceConstraint=*/false,
-      /*optionalLinkageOverride=*/nullptr, llvm::CallingConv::SwiftCoro);
+      /*optionalLinkageOverride=*/nullptr, IGM.SwiftCoroCC);
 }
 
 void irgen::emitYieldOnce2CoroutineEntry(IRGenFunction &IGF,
@@ -5208,12 +5253,8 @@ void irgen::emitYieldOnce2CoroutineEntry(IRGenFunction &IGF,
                                          llvm::Value *allocator,
                                          llvm::GlobalVariable *cfp) {
   IGF.setCoroutineAllocator(allocator);
-  auto isSwiftCoroCCAvailable =
-      IGF.IGM.SwiftCoroCC == llvm::CallingConv::SwiftCoro;
   auto allocFn = IGF.IGM.getOpaquePtr(getCoroAllocFn(IGF.IGM));
-  auto deallocFn = IGF.IGM.getOpaquePtr(isSwiftCoroCCAvailable
+  auto deallocFn = IGF.IGM.getOpaquePtr(getCoroDeallocFn(IGF.IGM));
-                                            ? getCoroDeallocWrapperFn(IGF.IGM)
-                                            : IGF.IGM.getCoroDeallocFn());
   emitRetconCoroutineEntry(
       IGF, fnType, buffer, llvm::Intrinsic::coro_id_retcon_once_dynamic,
       Size(-1) /*dynamic-to-IRGen size*/, IGF.IGM.getCoroStaticFrameAlignment(),

stdlib/public/runtime/Coro.cpp

@@ -15,14 +15,3 @@
 #include "swift/Basic/FlagSet.h"
 
 using namespace swift;
-
-void swift::swift_coro_dealloc(CoroAllocator *allocator, void *ptr) {
-  assert(allocator);
-  // Calls to swift_coro_dealloc are emitted in resume funclets for every
-  // live-across dynamic allocation.  Whether such calls immediately deallocate
-  // memory depends on the allocator.
-  if (!allocator->shouldDeallocateImmediately()) {
-    return;
-  }
-  allocator->deallocate(ptr);
-}

test/IRGen/coroutine_accessors.swift

@@ -39,6 +39,34 @@
 // CHECK:         ret ptr [[ALLOCATION]]
 // CHECK:       }
 
+// CHECK-LABEL: @_swift_coro_dealloc(
+// CHECK-SAME:      ptr [[ALLOCATOR:%[^,]+]]
+// CHECK-SAME:      ptr [[ADDRESS:%[^)]+]]
+// CHECK-SAME:  )
+// CHECK-SAME:  {
+// CHECK:       entry:
+// CHECK:         [[FLAGS_ADDR:%[^,]+]] = getelementptr inbounds %swift.coro_allocator
+// CHECK-SAME:        ptr [[ALLOCATOR]]
+// CHECK-SAME:        i32 0
+// CHECK-SAME:        i32 0
+// CHECK:         [[FLAGS:%[^,]+]] = load i32, ptr [[FLAGS_ADDR]], align 4
+// CHECK:         [[DEALLOC_DEFERRING_ALLOCATOR:%[^,]+]] = and i32 [[FLAGS]], 256
+// CHECK:         [[IS_DEALLOC_DEFERRING_ALLOCATOR:%[^,]+]] = icmp ne i32 [[DEALLOC_DEFERRING_ALLOCATOR]], 0
+// CHECK:         br i1 [[IS_DEALLOC_DEFERRING_ALLOCATOR]]
+// CHECK-SAME:        label %deferring_allocator
+// CHECK-SAME:        label %normal
+// CHECK:         deferring_allocator:
+// CHECK:           ret void
+// CHECK:         normal:
+// CHECK:           [[DEALLOCATE_FN_PTR:%[^,]+]] = getelementptr inbounds %swift.coro_allocator
+// CHECK-SAME:          ptr [[ALLOCATOR]]
+// CHECK-SAME:          i32 0
+// CHECK-SAME:          i32 2
+// CHECK:           [[DEALLOCATE_FN:%[^,]+]] = load ptr, ptr [[DEALLOCATE_FN_PTR]]
+// CHECK:           call swiftcc void [[DEALLOCATE_FN]](ptr [[ADDRESS]])
+// CHECK:         ret void
+// CHECK:       }
+
 @frozen
 public struct S {
 public var o: any AnyObject

test/IRGen/coroutine_accessors_popless.swift

@@ -52,7 +52,7 @@
 // CHECK:         ret ptr [[ALLOCATION]]
 // CHECK:       }
 
-// CHECK-LABEL: @__swift_coro_dealloc_(
+// CHECK-LABEL: @_swift_coro_dealloc(
 // CHECK-SAME:      ptr [[ALLOCATOR:%[^,]+]]
 // CHECK-SAME:      ptr [[ADDRESS:%[^)]+]]
 // CHECK-SAME:  )
@@ -60,15 +60,30 @@
 // CHECK:       entry:
 // CHECK:         [[BAIL:%[^,]+]] = icmp eq ptr [[ALLOCATOR]], null
 // CHECK:         br i1 [[USE_POPLESS]],
-// CHECK-SAME:        label %bail
+// CHECK-SAME:        label %null_allocator
-// CHECK-SAME:        label %forward
+// CHECK-SAME:        label %nonnull_allocator
-// CHECK:       bail:
+// CHECK:       null_allocator:
 // CHECK:         ret void
-// CHECK:       forward:
+// CHECK:       nonnull_allocator:
-// CHECK:         call swiftcc void @swift_coro_dealloc(
+// CHECK:         [[FLAGS_ADDR:%[^,]+]] = getelementptr inbounds %swift.coro_allocator
 // CHECK-SAME:        ptr [[ALLOCATOR]]
-// CHECK-SAME:        ptr [[ADDRESS]]
+// CHECK-SAME:        i32 0
-// CHECK-SAME:    )
+// CHECK-SAME:        i32 0
+// CHECK:         [[FLAGS:%[^,]+]] = load i32, ptr [[FLAGS_ADDR]], align 4
+// CHECK:         [[DEALLOC_DEFERRING_ALLOCATOR:%[^,]+]] = and i32 [[FLAGS]], 256
+// CHECK:         [[IS_DEALLOC_DEFERRING_ALLOCATOR:%[^,]+]] = icmp ne i32 [[DEALLOC_DEFERRING_ALLOCATOR]], 0
+// CHECK:         br i1 [[IS_DEALLOC_DEFERRING_ALLOCATOR]]
+// CHECK-SAME:        label %deferring_allocator
+// CHECK-SAME:        label %normal
+// CHECK:         deferring_allocator:
+// CHECK:           ret void
+// CHECK:         normal:
+// CHECK:           [[DEALLOCATE_FN_PTR:%[^,]+]] = getelementptr inbounds %swift.coro_allocator
+// CHECK-SAME:          ptr [[ALLOCATOR]]
+// CHECK-SAME:          i32 0
+// CHECK-SAME:          i32 2
+// CHECK:           [[DEALLOCATE_FN:%[^,]+]] = load ptr, ptr [[DEALLOCATE_FN_PTR]]
+// CHECK:           call swiftcc void [[DEALLOCATE_FN]](ptr [[ADDRESS]])
 // CHECK:         ret void
 // CHECK:       }
 

test/abi/macOS/arm64/stdlib.swift

@@ -942,7 +942,3 @@ Added: _$ss18EnumeratedSequenceVsSlRzrlE7isEmptySbvpMV
 Added: _$ss18EnumeratedSequenceVsSlRzrlE8endIndexABsSlRzrlE0D0Vyx_GvpMV
 Added: _$ss18EnumeratedSequenceVsSlRzrlEySi6offset_7ElementQz7elementtABsSlRzrlE5IndexVyx_GcipMV
 Added: _$ss18EnumeratedSequenceVyxGSKsSkRzrlMc
-
-// CoroutineAccessors
-Added: _swift_coro_alloc
-Added: _swift_coro_dealloc

test/abi/macOS/x86_64/stdlib.swift

@@ -942,7 +942,3 @@ Added: _$ss18EnumeratedSequenceVsSlRzrlE7isEmptySbvpMV
 Added: _$ss18EnumeratedSequenceVsSlRzrlE8endIndexABsSlRzrlE0D0Vyx_GvpMV
 Added: _$ss18EnumeratedSequenceVsSlRzrlEySi6offset_7ElementQz7elementtABsSlRzrlE5IndexVyx_GcipMV
 Added: _$ss18EnumeratedSequenceVyxGSKsSkRzrlMc
-
-// CoroutineAccessors
-Added: _swift_coro_alloc
-Added: _swift_coro_dealloc