diff --git a/include/swift/ABI/MetadataValues.h b/include/swift/ABI/MetadataValues.h index 4e5e9b79d31b8..54059a6dd7d8e 100644 --- a/include/swift/ABI/MetadataValues.h +++ b/include/swift/ABI/MetadataValues.h @@ -1793,6 +1793,10 @@ namespace SpecialPointerAuthDiscriminators { /// IsCurrentGlobalActor function used between the Swift runtime and /// concurrency runtime. const uint16_t IsCurrentGlobalActorFunction = 0xd1b8; // = 53688 + + /// Function pointers stored in the coro allocator struct. + const uint16_t CoroAllocationFunction = 0x5f95; // = 24469 + const uint16_t CoroDeallocationFunction = 0x9faf; // = 40879 } /// The number of arguments that will be passed directly to a generic diff --git a/include/swift/AST/IRGenOptions.h b/include/swift/AST/IRGenOptions.h index baa4709dc05e1..50d305389c6a1 100644 --- a/include/swift/AST/IRGenOptions.h +++ b/include/swift/AST/IRGenOptions.h @@ -251,6 +251,12 @@ struct PointerAuthOptions : clang::PointerAuthOptions { /// Like PartialApplyCapture but for use with CoroFunctionPointer values. PointerAuthSchema CoroPartialApplyCapture; + + /// Stored in a coro allocator struct, the function used to allocate memory. + PointerAuthSchema CoroAllocationFunction; + + /// Stored in a coro allocator struct, the function used to deallocate memory. + PointerAuthSchema CoroDeallocationFunction; }; enum class JITDebugArtifact : unsigned { diff --git a/lib/IRGen/GenCall.cpp b/lib/IRGen/GenCall.cpp index 6e622da2773d7..0acbb19124a68 100644 --- a/lib/IRGen/GenCall.cpp +++ b/lib/IRGen/GenCall.cpp @@ -48,6 +48,7 @@ #include "EntryPointArgumentEmission.h" #include "Explosion.h" #include "GenCall.h" +#include "GenCoro.h" #include "GenFunc.h" #include "GenHeap.h" #include "GenKeyPath.h" @@ -5131,141 +5132,6 @@ void irgen::emitYieldManyCoroutineEntry( allocFn, deallocFn, {}); } -static llvm::Constant *getCoroAllocFn(IRGenModule &IGM) { - auto isSwiftCoroCCAvailable = IGM.SwiftCoroCC == llvm::CallingConv::SwiftCoro; - return IGM.getOrCreateHelperFunction( - "_swift_coro_alloc", IGM.Int8PtrTy, {IGM.CoroAllocatorPtrTy, IGM.SizeTy}, - [isSwiftCoroCCAvailable](IRGenFunction &IGF) { - auto parameters = IGF.collectParameters(); - auto *allocator = parameters.claimNext(); - auto *size = parameters.claimNext(); - if (isSwiftCoroCCAvailable) { - // swiftcorocc is available, so if there's no allocator pointer, - // allocate storage on the stack and return a pointer to it without - // popping the stack. - auto *nullAllocator = IGF.Builder.CreateCmp( - llvm::CmpInst::Predicate::ICMP_EQ, allocator, - llvm::ConstantPointerNull::get( - cast(allocator->getType()))); - auto *poplessReturn = IGF.createBasicBlock("popless"); - auto *normalReturn = IGF.createBasicBlock("normal"); - IGF.Builder.CreateCondBr(nullAllocator, poplessReturn, normalReturn); - IGF.Builder.emitBlock(poplessReturn); - // Emit the dynamic alloca. - auto *alloca = - IGF.Builder.IRBuilderBase::CreateAlloca(IGF.IGM.Int8Ty, size); - alloca->setAlignment(llvm::Align(MaximumAlignment)); - auto *retPopless = IGF.Builder.CreateIntrinsic( - IGF.IGM.VoidTy, llvm::Intrinsic::ret_popless, {}); - retPopless->setTailCallKind( - llvm::CallInst::TailCallKind::TCK_MustTail); - IGF.Builder.CreateRet(alloca); - // Start emitting the "normal" block. - IGF.Builder.emitBlock(normalReturn); - } - auto *calleePtr = IGF.Builder.CreateInBoundsGEP( - IGF.IGM.CoroAllocatorTy, allocator, - {llvm::ConstantInt::get(IGF.IGM.Int32Ty, 0), - llvm::ConstantInt::get(IGF.IGM.Int32Ty, 1)}); - auto *callee = IGF.Builder.CreateLoad( - Address(calleePtr, IGF.IGM.PtrTy, IGF.IGM.getPointerAlignment()), - "allocate_fn"); - auto fnPtr = FunctionPointer::createUnsigned( - FunctionPointer::Kind::Function, callee, - Signature(cast(IGF.IGM.CoroAllocateFnTy), {}, - IGF.IGM.SwiftCC)); - auto *call = IGF.Builder.CreateCall(fnPtr, {size}); - call->setDoesNotThrow(); - call->setCallingConv(IGF.IGM.SwiftCC); - IGF.Builder.CreateRet(call); - }, - /*setIsNoInline=*/true, - /*forPrologue=*/false, - /*isPerformanceConstraint=*/false, - /*optionalLinkageOverride=*/nullptr, IGM.SwiftCoroCC, - /*transformAttributes=*/ - [&IGM](llvm::AttributeList &attrs) { - IGM.addSwiftCoroAttributes(attrs, 0); - }); -} - -static llvm::Constant *getCoroDeallocFn(IRGenModule &IGM) { - auto isSwiftCoroCCAvailable = IGM.SwiftCoroCC == llvm::CallingConv::SwiftCoro; - return IGM.getOrCreateHelperFunction( - "_swift_coro_dealloc", IGM.VoidTy, - {IGM.CoroAllocatorPtrTy, IGM.Int8PtrTy}, - [isSwiftCoroCCAvailable](IRGenFunction &IGF) { - auto parameters = IGF.collectParameters(); - auto *allocator = parameters.claimNext(); - auto *ptr = parameters.claimNext(); - if (isSwiftCoroCCAvailable) { - // swiftcorocc is available, so if there's no allocator pointer, - // storage was allocated on the stack which will be naturally cleaned - // up when the coroutine's frame is "freed". - auto *nullAllocator = IGF.Builder.CreateCmp( - llvm::CmpInst::Predicate::ICMP_EQ, allocator, - llvm::ConstantPointerNull::get( - cast(allocator->getType()))); - 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(); - // Start emitting the "normal" block. - IGF.Builder.emitBlock(normalBlock); - 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.PtrTy, IGF.IGM.getPointerAlignment()), - "deallocate_fn"); - auto fnPtr = FunctionPointer::createUnsigned( - FunctionPointer::Kind::Function, callee, - Signature(cast(IGF.IGM.CoroDeallocateFnTy), {}, - IGF.IGM.SwiftCC)); - auto *call = IGF.Builder.CreateCall(fnPtr, {ptr}); - call->setDoesNotThrow(); - call->setCallingConv(IGF.IGM.SwiftCC); - IGF.Builder.CreateRetVoid(); - }, - /*setIsNoInline=*/true, - /*forPrologue=*/false, - /*isPerformanceConstraint=*/false, - /*optionalLinkageOverride=*/nullptr, IGM.SwiftCoroCC, - /*transformAttributes=*/ - [&IGM](llvm::AttributeList &attrs) { - IGM.addSwiftCoroAttributes(attrs, 0); - }); -} - void irgen::emitYieldOnce2CoroutineEntry(IRGenFunction &IGF, CanSILFunctionType fnType, llvm::Value *buffer, @@ -5308,85 +5174,6 @@ Address irgen::emitAllocYieldManyCoroutineBuffer(IRGenFunction &IGF) { getYieldManyCoroutineBufferAlignment(IGF.IGM)); } -static llvm::Constant *getAddrOfSwiftCCMalloc(IRGenModule &IGM) { - auto mallocFnPtr = IGM.getMallocFunctionPointer(); - auto sig = mallocFnPtr.getSignature(); - if (sig.getCallingConv() == IGM.SwiftCC) { - return IGM.getMallocFn(); - } - return IGM.getOrCreateHelperFunction( - "_swift_malloc", sig.getType()->getReturnType(), sig.getType()->params(), - [](IRGenFunction &IGF) { - auto parameters = IGF.collectParameters(); - auto *size = parameters.claimNext(); - auto malloc = IGF.IGM.getMallocFunctionPointer(); - auto *call = IGF.Builder.CreateCall(malloc, {size}); - IGF.Builder.CreateRet(call); - }); -} - -static llvm::Constant *getAddrOfSwiftCCFree(IRGenModule &IGM) { - auto freeFnPtr = IGM.getFreeFunctionPointer(); - auto sig = freeFnPtr.getSignature(); - if (sig.getCallingConv() == IGM.SwiftCC) { - return IGM.getFreeFn(); - } - return IGM.getOrCreateHelperFunction( - "_swift_free", sig.getType()->getReturnType(), sig.getType()->params(), - [](IRGenFunction &IGF) { - auto parameters = IGF.collectParameters(); - auto *ptr = parameters.claimNext(); - auto free = IGF.IGM.getFreeFunctionPointer(); - IGF.Builder.CreateCall(free, {ptr}); - IGF.Builder.CreateRetVoid(); - }); -} - -static llvm::Constant *getAddrOfGlobalCoroAllocator( - IRGenModule &IGM, CoroAllocatorKind kind, bool shouldDeallocateImmediately, - llvm::Constant *allocFn, llvm::Constant *deallocFn) { - auto entity = LinkEntity::forCoroAllocator(kind); - auto taskAllocator = IGM.getOrCreateLazyGlobalVariable( - entity, - [&](ConstantInitBuilder &builder) -> ConstantInitFuture { - auto allocator = builder.beginStruct(IGM.CoroAllocatorTy); - auto flags = CoroAllocatorFlags(kind); - flags.setShouldDeallocateImmediately(shouldDeallocateImmediately); - allocator.addInt32(flags.getOpaqueValue()); - allocator.add(allocFn); - allocator.add(deallocFn); - return allocator.finishAndCreateFuture(); - }, - [&](llvm::GlobalVariable *var) { var->setConstant(true); }); - return taskAllocator; -} -llvm::Constant *IRGenModule::getAddrOfGlobalCoroMallocAllocator() { - return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Malloc, - /*shouldDeallocateImmediately=*/true, - getAddrOfSwiftCCMalloc(*this), - getAddrOfSwiftCCFree(*this)); -} -llvm::Constant *IRGenModule::getAddrOfGlobalCoroAsyncTaskAllocator() { - return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Async, - /*shouldDeallocateImmediately=*/false, - getTaskAllocFn(), getTaskDeallocFn()); -} -llvm::Value * -irgen::emitYieldOnce2CoroutineAllocator(IRGenFunction &IGF, - std::optional kind) { - if (!kind) { - return IGF.getCoroutineAllocator(); - } - switch (*kind) { - case CoroAllocatorKind::Stack: - return llvm::ConstantPointerNull::get(IGF.IGM.CoroAllocatorPtrTy); - case CoroAllocatorKind::Async: - return IGF.IGM.getAddrOfGlobalCoroAsyncTaskAllocator(); - case CoroAllocatorKind::Malloc: - return IGF.IGM.getAddrOfGlobalCoroMallocAllocator(); - } - llvm_unreachable("unhandled case"); -} StackAddress irgen::emitAllocYieldOnce2CoroutineFrame(IRGenFunction &IGF, llvm::Value *size) { return emitAllocCoroStaticFrame(IGF, size); diff --git a/lib/IRGen/GenCall.h b/lib/IRGen/GenCall.h index 66194a70bd070..cd6707621a7e5 100644 --- a/lib/IRGen/GenCall.h +++ b/lib/IRGen/GenCall.h @@ -220,9 +220,6 @@ namespace irgen { CanSILFunctionType coroutineType, NativeCCEntryPointArgumentEmission &emission); - llvm::Value * - emitYieldOnce2CoroutineAllocator(IRGenFunction &IGF, - std::optional kind); StackAddress emitAllocYieldOnce2CoroutineFrame(IRGenFunction &IGF, llvm::Value *size); void emitDeallocYieldOnce2CoroutineFrame(IRGenFunction &IGF, diff --git a/lib/IRGen/GenCoro.cpp b/lib/IRGen/GenCoro.cpp index f9ab20348e5b4..faea8e37547e5 100644 --- a/lib/IRGen/GenCoro.cpp +++ b/lib/IRGen/GenCoro.cpp @@ -10,6 +10,7 @@ // //===----------------------------------------------------------------------===// +#include "swift/ABI/Coro.h" #include "swift/ABI/MetadataValues.h" #include "swift/Basic/Assertions.h" #include "swift/IRGen/Linking.h" @@ -17,7 +18,10 @@ #include "llvm/IR/Instructions.h" #include "llvm/IR/Module.h" +#include "ConstantBuilder.h" #include "Explosion.h" +#include "GenCoro.h" +#include "GenPointerAuth.h" #include "IRGenFunction.h" #include "IRGenModule.h" @@ -612,3 +616,304 @@ void IRGenFunction::emitTaskDeallocThrough(Address address) { call->setDoesNotThrow(); call->setCallingConv(IGM.SwiftCC); } + +namespace { +struct Allocator { + struct Field { + enum Kind : uint8_t { + Flags = 0, + Allocate = 1, + Deallocate = 2, + }; + Kind kind; + Field(Kind kind) : kind(kind) {} + operator Kind() { return kind; } + + llvm::Type *getType(IRGenModule &IGM) { + switch (kind) { + case Flags: + return IGM.Int32Ty; + case Field::Allocate: + case Field::Deallocate: + return IGM.PtrTy; + } + } + + Alignment getAlignment(IRGenModule &IGM) { + switch (kind) { + case Flags: + return Alignment(4); + case Field::Allocate: + case Field::Deallocate: + return IGM.getPointerAlignment(); + } + } + + StringRef getName() { + switch (kind) { + case Flags: + return "flags"; + case Field::Allocate: + return "allocate_fn"; + case Field::Deallocate: + return "deallocate_fn"; + } + } + + llvm::FunctionType *getFunctionType(IRGenModule &IGM) { + switch (kind) { + case Flags: + llvm_unreachable("not a function"); + case Field::Allocate: + return IGM.CoroAllocateFnTy; + case Field::Deallocate: + return IGM.CoroDeallocateFnTy; + } + } + + const PointerAuthSchema &getSchema(IRGenModule &IGM) { + switch (kind) { + case Flags: + llvm_unreachable("no schema"); + case Field::Allocate: + return IGM.getOptions().PointerAuth.CoroAllocationFunction; + case Field::Deallocate: + return IGM.getOptions().PointerAuth.CoroDeallocationFunction; + } + } + }; + + llvm::Value *address; + IRGenFunction &IGF; + + Allocator(llvm::Value *address, IRGenFunction &IGF) + : address(address), IGF(IGF) {} + + llvm::Value *getField(Field field) { + auto *fieldAddress = IGF.Builder.CreateInBoundsGEP( + IGF.IGM.CoroAllocatorTy, address, + {llvm::ConstantInt::get(IGF.IGM.Int32Ty, 0), + llvm::ConstantInt::get(IGF.IGM.Int32Ty, field.kind)}); + return IGF.Builder.CreateLoad(Address(fieldAddress, field.getType(IGF.IGM), + field.getAlignment(IGF.IGM)), + field.getName()); + } + + llvm::Value *getFlags() { return getField(Field::Flags); } + + FunctionPointer getAllocate() { return getFunctionPointer(Field::Allocate); } + + FunctionPointer getDeallocate() { + return getFunctionPointer(Field::Deallocate); + } + + void emitOnNullAllocator(StringRef nullBlockName, StringRef nonnullBlockName, + llvm::function_ref body) { + auto *nullAllocator = + IGF.Builder.CreateCmp(llvm::CmpInst::Predicate::ICMP_EQ, address, + llvm::ConstantPointerNull::get( + cast(address->getType()))); + auto *poplessReturn = IGF.createBasicBlock(nullBlockName); + auto *normalReturn = IGF.createBasicBlock(nonnullBlockName); + IGF.Builder.CreateCondBr(nullAllocator, poplessReturn, normalReturn); + IGF.Builder.emitBlock(poplessReturn); + body(IGF); + // Start emitting the "normal" block. + IGF.Builder.emitBlock(normalReturn); + } + +private: + FunctionPointer getFunctionPointer(Field field) { + llvm::Value *callee = getField(field); + if (auto &schema = field.getSchema(IGF.IGM)) { + auto info = + PointerAuthInfo::emit(IGF, schema, nullptr, PointerAuthEntity()); + callee = emitPointerAuthAuth(IGF, callee, info); + } + return FunctionPointer::createUnsigned( + FunctionPointer::Kind::Function, callee, + Signature(field.getFunctionType(IGF.IGM), {}, IGF.IGM.SwiftCC)); + } +}; +} // end anonymous namespace + +llvm::Constant *swift::irgen::getCoroAllocFn(IRGenModule &IGM) { + auto isSwiftCoroCCAvailable = IGM.SwiftCoroCC == llvm::CallingConv::SwiftCoro; + return IGM.getOrCreateHelperFunction( + "_swift_coro_alloc", IGM.Int8PtrTy, {IGM.CoroAllocatorPtrTy, IGM.SizeTy}, + [isSwiftCoroCCAvailable](IRGenFunction &IGF) { + auto parameters = IGF.collectParameters(); + auto allocator = Allocator(parameters.claimNext(), IGF); + auto *size = parameters.claimNext(); + if (isSwiftCoroCCAvailable) { + // swiftcorocc is available, so if there's no allocator pointer, + // allocate storage on the stack and return a pointer to it without + // popping the stack. + allocator.emitOnNullAllocator("popless", "normal", [size](auto &IGF) { + // Emit the dynamic alloca. + auto *alloca = + IGF.Builder.IRBuilderBase::CreateAlloca(IGF.IGM.Int8Ty, size); + alloca->setAlignment(llvm::Align(MaximumAlignment)); + auto *retPopless = IGF.Builder.CreateIntrinsic( + IGF.IGM.VoidTy, llvm::Intrinsic::ret_popless, {}); + retPopless->setTailCallKind( + llvm::CallInst::TailCallKind::TCK_MustTail); + IGF.Builder.CreateRet(alloca); + }); + } + auto fnPtr = allocator.getAllocate(); + auto *call = IGF.Builder.CreateCall(fnPtr, {size}); + call->setDoesNotThrow(); + call->setCallingConv(IGF.IGM.SwiftCC); + IGF.Builder.CreateRet(call); + }, + /*setIsNoInline=*/true, + /*forPrologue=*/false, + /*isPerformanceConstraint=*/false, + /*optionalLinkageOverride=*/nullptr, IGM.SwiftCoroCC, + /*transformAttributes=*/ + [&IGM](llvm::AttributeList &attrs) { + IGM.addSwiftCoroAttributes(attrs, 0); + }); +} + +llvm::Constant *swift::irgen::getCoroDeallocFn(IRGenModule &IGM) { + auto isSwiftCoroCCAvailable = IGM.SwiftCoroCC == llvm::CallingConv::SwiftCoro; + return IGM.getOrCreateHelperFunction( + "_swift_coro_dealloc", IGM.VoidTy, + {IGM.CoroAllocatorPtrTy, IGM.Int8PtrTy}, + [isSwiftCoroCCAvailable](IRGenFunction &IGF) { + auto parameters = IGF.collectParameters(); + auto allocator = Allocator(parameters.claimNext(), IGF); + auto *ptr = parameters.claimNext(); + if (isSwiftCoroCCAvailable) { + // swiftcorocc is available, so if there's no allocator pointer, + // storage was allocated on the stack which will be naturally cleaned + // up when the coroutine's frame is "freed". + allocator.emitOnNullAllocator("null_allocator", "nonnull_allocator", + [&](auto &IGF) { + // Nothing to do here. + IGF.Builder.CreateRetVoid(); + }); + } + auto shouldDeallocateImmediatelyFlag = CoroAllocatorFlags(0); + shouldDeallocateImmediatelyFlag.setShouldDeallocateImmediately(true); + auto *flags = allocator.getFlags(); + auto *isDeallocDeferringAllocator = IGF.Builder.CreateICmpNE( + IGF.Builder.CreateAnd( + flags, + llvm::APInt(IGF.IGM.Int32Ty->getBitWidth(), + shouldDeallocateImmediatelyFlag.getOpaqueValue())), + 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(); + // Start emitting the "normal" block. + IGF.Builder.emitBlock(normalBlock); + auto fnPtr = allocator.getDeallocate(); + auto *call = IGF.Builder.CreateCall(fnPtr, {ptr}); + call->setDoesNotThrow(); + call->setCallingConv(IGF.IGM.SwiftCC); + IGF.Builder.CreateRetVoid(); + }, + /*setIsNoInline=*/true, + /*forPrologue=*/false, + /*isPerformanceConstraint=*/false, + /*optionalLinkageOverride=*/nullptr, IGM.SwiftCoroCC, + /*transformAttributes=*/ + [&IGM](llvm::AttributeList &attrs) { + IGM.addSwiftCoroAttributes(attrs, 0); + }); +} + +static llvm::Constant *getAddrOfGlobalCoroAllocator( + IRGenModule &IGM, CoroAllocatorKind kind, bool shouldDeallocateImmediately, + llvm::Constant *allocFn, llvm::Constant *deallocFn) { + auto entity = LinkEntity::forCoroAllocator(kind); + auto taskAllocator = IGM.getOrCreateLazyGlobalVariable( + entity, + [&](ConstantInitBuilder &builder) -> ConstantInitFuture { + auto allocator = builder.beginStruct(IGM.CoroAllocatorTy); + auto flags = CoroAllocatorFlags(kind); + flags.setShouldDeallocateImmediately(shouldDeallocateImmediately); + allocator.addInt32(flags.getOpaqueValue()); + allocator.addSignedPointer( + allocFn, IGM.getOptions().PointerAuth.CoroAllocationFunction, + PointerAuthEntity::Special::CoroAllocationFunction); + allocator.addSignedPointer( + deallocFn, IGM.getOptions().PointerAuth.CoroDeallocationFunction, + PointerAuthEntity::Special::CoroDeallocationFunction); + return allocator.finishAndCreateFuture(); + }, + [&](llvm::GlobalVariable *var) { var->setConstant(true); }); + return taskAllocator; +} + +static llvm::Constant *getAddrOfSwiftCCMalloc(IRGenModule &IGM) { + auto mallocFnPtr = IGM.getMallocFunctionPointer(); + auto sig = mallocFnPtr.getSignature(); + if (sig.getCallingConv() == IGM.SwiftCC) { + return IGM.getMallocFn(); + } + return IGM.getOrCreateHelperFunction( + "_swift_malloc", sig.getType()->getReturnType(), sig.getType()->params(), + [](IRGenFunction &IGF) { + auto parameters = IGF.collectParameters(); + auto *size = parameters.claimNext(); + auto malloc = IGF.IGM.getMallocFunctionPointer(); + auto *call = IGF.Builder.CreateCall(malloc, {size}); + IGF.Builder.CreateRet(call); + }); +} + +static llvm::Constant *getAddrOfSwiftCCFree(IRGenModule &IGM) { + auto freeFnPtr = IGM.getFreeFunctionPointer(); + auto sig = freeFnPtr.getSignature(); + if (sig.getCallingConv() == IGM.SwiftCC) { + return IGM.getFreeFn(); + } + return IGM.getOrCreateHelperFunction( + "_swift_free", sig.getType()->getReturnType(), sig.getType()->params(), + [](IRGenFunction &IGF) { + auto parameters = IGF.collectParameters(); + auto *ptr = parameters.claimNext(); + auto free = IGF.IGM.getFreeFunctionPointer(); + IGF.Builder.CreateCall(free, {ptr}); + IGF.Builder.CreateRetVoid(); + }); +} + +llvm::Constant *IRGenModule::getAddrOfGlobalCoroMallocAllocator() { + return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Malloc, + /*shouldDeallocateImmediately=*/true, + getAddrOfSwiftCCMalloc(*this), + getAddrOfSwiftCCFree(*this)); +} + +llvm::Constant *IRGenModule::getAddrOfGlobalCoroAsyncTaskAllocator() { + return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Async, + /*shouldDeallocateImmediately=*/false, + getTaskAllocFn(), getTaskDeallocFn()); +} + +llvm::Value * +irgen::emitYieldOnce2CoroutineAllocator(IRGenFunction &IGF, + std::optional kind) { + if (!kind) { + return IGF.getCoroutineAllocator(); + } + switch (*kind) { + case CoroAllocatorKind::Stack: + return llvm::ConstantPointerNull::get(IGF.IGM.CoroAllocatorPtrTy); + case CoroAllocatorKind::Async: + return IGF.IGM.getAddrOfGlobalCoroAsyncTaskAllocator(); + case CoroAllocatorKind::Malloc: + return IGF.IGM.getAddrOfGlobalCoroMallocAllocator(); + } + llvm_unreachable("unhandled case"); +} diff --git a/lib/IRGen/GenCoro.h b/lib/IRGen/GenCoro.h new file mode 100644 index 0000000000000..6dc200c17e7fd --- /dev/null +++ b/lib/IRGen/GenCoro.h @@ -0,0 +1,26 @@ +//===----- GenCoro.h - Code generation related to coroutines --------------===// +// +// This source file is part of the Swift.org open source project +// +// Copyright (c) 2014 - 2025 Apple Inc. and the Swift project authors +// Licensed under Apache License v2.0 with Runtime Library Exception +// +// See https://swift.org/LICENSE.txt for license information +// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors +// +//===----------------------------------------------------------------------===// + +namespace swift { +namespace irgen { + +class IRGenModule; + +llvm::Constant *getCoroAllocFn(IRGenModule &IGM); +llvm::Constant *getCoroDeallocFn(IRGenModule &IGM); + +llvm::Value * +emitYieldOnce2CoroutineAllocator(IRGenFunction &IGF, + std::optional kind); + +} // end namespace irgen +} // end namespace swift diff --git a/lib/IRGen/GenPointerAuth.cpp b/lib/IRGen/GenPointerAuth.cpp index 9d094032747be..abc5af0afd30d 100644 --- a/lib/IRGen/GenPointerAuth.cpp +++ b/lib/IRGen/GenPointerAuth.cpp @@ -400,6 +400,10 @@ PointerAuthEntity::getDeclDiscriminator(IRGenModule &IGM) const { case Special::BlockCopyHelper: case Special::BlockDisposeHelper: llvm_unreachable("no known discriminator for these foreign entities"); + case Special::CoroAllocationFunction: + return SpecialPointerAuthDiscriminators::CoroAllocationFunction; + case Special::CoroDeallocationFunction: + return SpecialPointerAuthDiscriminators::CoroDeallocationFunction; } llvm_unreachable("bad kind"); }; diff --git a/lib/IRGen/GenPointerAuth.h b/lib/IRGen/GenPointerAuth.h index 2473f57365aca..989dee14954ec 100644 --- a/lib/IRGen/GenPointerAuth.h +++ b/lib/IRGen/GenPointerAuth.h @@ -68,6 +68,8 @@ class PointerAuthEntity { OpaqueTypeDescriptorAsArgument, ContextDescriptorAsArgument, TypeLayoutString, + CoroAllocationFunction, + CoroDeallocationFunction, }; private: diff --git a/lib/IRGen/IRGen.cpp b/lib/IRGen/IRGen.cpp index d70df477ef61c..7bf06627f3d9d 100644 --- a/lib/IRGen/IRGen.cpp +++ b/lib/IRGen/IRGen.cpp @@ -990,6 +990,14 @@ static void setPointerAuthOptions(PointerAuthOptions &opts, opts.CoroPartialApplyCapture = PointerAuthSchema(nonABIDataKey, /*address*/ true, Discrimination::Decl); + + opts.CoroAllocationFunction = PointerAuthSchema( + codeKey, /*address*/ false, Discrimination::Constant, + SpecialPointerAuthDiscriminators::CoroAllocationFunction); + + opts.CoroDeallocationFunction = PointerAuthSchema( + codeKey, /*address*/ false, Discrimination::Constant, + SpecialPointerAuthDiscriminators::CoroDeallocationFunction); } std::unique_ptr diff --git a/test/IRGen/coroutine_accessors.swift b/test/IRGen/coroutine_accessors.swift index 914a102153873..f9bd254295723 100644 --- a/test/IRGen/coroutine_accessors.swift +++ b/test/IRGen/coroutine_accessors.swift @@ -23,14 +23,42 @@ // CHECK-SAME: i32 0 // CHECK-SAME: }> +// CHECK-arm64e-LABEL: _swift_malloc.ptrauth = private constant { +// CHECK-arm64e-SAME: ptr @_swift_malloc, +// CHECK-arm64e-SAME: i32 0, +// CHECK-arm64e-SAME: i64 0, +// CHECK-arm64e-SAME: i64 24469 } +// CHECK-arm64e-SAME: section "llvm.ptrauth" +// CHECK-arm64e-SAME: align 8 +// CHECK-arm64e-LABEL: _swift_free.ptrauth = private constant { +// CHECK-arm64e-SAME: ptr @_swift_free, +// CHECK-arm64e-SAME: i32 0, +// CHECK-arm64e-SAME: i64 0, +// CHECK-arm64e-SAME: i64 40879 }, +// CHECK-arm64e-SAME: section "llvm.ptrauth", +// CHECK-arm64e-SAME: align 8 // CHECK-LABEL: _swift_coro_malloc_allocator = linkonce_odr hidden constant %swift.coro_allocator { // CHECK-SAME: i32 258, -// CHECK-SAME: malloc, +// CHECK-SAME: malloc // CHECK-SAME: free // CHECK-SAME: } +// CHECK-arm64e-LABEL: swift_task_alloc.ptrauth = private constant { +// CHECK-arm64e-SAME: ptr @swift_task_alloc, +// CHECK-arm64e-SAME: i32 0, +// CHECK-arm64e-SAME: i64 0, +// CHECK-arm64e-SAME: i64 24469 } +// CHECK-arm64e-SAME: section "llvm.ptrauth" +// CHECK-arm64e-SAME: align 8 +// CHECK-arm64e-LABEL: @swift_task_dealloc.ptrauth = private constant { +// CHECK-arm64e-SAME: ptr @swift_task_dealloc, +// CHECK-arm64e-SAME: i32 0, +// CHECK-arm64e-SAME: i64 0, +// CHECK-arm64e-SAME: i64 40879 }, +// CHECK-arm64e-SAME: section "llvm.ptrauth", +// CHECK-arm64e-SAME: align 8 // CHECK-LABEL: _swift_coro_async_allocator = linkonce_odr hidden constant %swift.coro_allocator { // CHECK-SAME: i32 1, -// CHECK-SAME: swift_task_alloc, +// CHECK-SAME: swift_task_alloc // CHECK-SAME: swift_task_dealloc // CHECK-SAME: } @@ -45,6 +73,9 @@ // CHECK-SAME: i32 0 // CHECK-SAME: i32 1 // CHECK: [[ALLOCATE_FN:%[^,]+]] = load ptr, ptr [[ALLOCATE_FN_PTR]] +// CHECK-arm64e: [[ALLOCATE_FN_BITS:%[^,]+]] = ptrtoint ptr [[ALLOCATE_FN]] to i64 +// CHECK-arm64e: [[ALLOCATE_FN_BITS_AUTHED:%[^,]+]] = call i64 @llvm.ptrauth.auth(i64 [[ALLOCATE_FN_BITS]], i32 0, i64 24469) +// CHECK-arm64e: [[ALLOCATE_FN:%[^,]+]] = inttoptr i64 [[ALLOCATE_FN_BITS_AUTHED]] // CHECK: [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]]([[INT]] [[SIZE]]) // CHECK: ret ptr [[ALLOCATION]] // CHECK: } @@ -73,6 +104,9 @@ // CHECK-SAME: i32 0 // CHECK-SAME: i32 2 // CHECK: [[DEALLOCATE_FN:%[^,]+]] = load ptr, ptr [[DEALLOCATE_FN_PTR]] +// CHECK-arm64e: [[DEALLOCATE_FN_BITS:%[^,]+]] = ptrtoint ptr [[DEALLOCATE_FN]] to i64 +// CHECK-arm64e: [[DEALLOCATE_FN_BITS_AUTHED:%[^,]+]] = call i64 @llvm.ptrauth.auth(i64 [[DEALLOCATE_FN_BITS]], i32 0, i64 40879) +// CHECK-arm64e: [[DEALLOCATE_FN:%[^,]+]] = inttoptr i64 [[DEALLOCATE_FN_BITS_AUTHED]] // CHECK: call swiftcc void [[DEALLOCATE_FN]](ptr [[ADDRESS]]) // CHECK: ret void // CHECK: } diff --git a/test/IRGen/coroutine_accessors_popless.swift b/test/IRGen/coroutine_accessors_popless.swift index 38db9436ea82e..f7672fab7288e 100644 --- a/test/IRGen/coroutine_accessors_popless.swift +++ b/test/IRGen/coroutine_accessors_popless.swift @@ -25,14 +25,42 @@ // CHECK-SAME: i32 0 // CHECK-SAME: }> +// CHECK-arm64e-LABEL: swift_task_alloc.ptrauth = private constant { +// CHECK-arm64e-SAME: ptr @swift_task_alloc, +// CHECK-arm64e-SAME: i32 0, +// CHECK-arm64e-SAME: i64 0, +// CHECK-arm64e-SAME: i64 24469 } +// CHECK-arm64e-SAME: section "llvm.ptrauth" +// CHECK-arm64e-SAME: align 8 +// CHECK-arm64e-LABEL: @swift_task_dealloc.ptrauth = private constant { +// CHECK-arm64e-SAME: ptr @swift_task_dealloc, +// CHECK-arm64e-SAME: i32 0, +// CHECK-arm64e-SAME: i64 0, +// CHECK-arm64e-SAME: i64 40879 }, +// CHECK-arm64e-SAME: section "llvm.ptrauth", +// CHECK-arm64e-SAME: align 8 // CHECK-LABEL: _swift_coro_async_allocator = linkonce_odr hidden constant %swift.coro_allocator { // CHECK-SAME: i32 1, -// CHECK-SAME: swift_task_alloc, +// CHECK-SAME: swift_task_alloc // CHECK-SAME: swift_task_dealloc // CHECK-SAME: } +// CHECK-arm64e-LABEL: _swift_malloc.ptrauth = private constant { +// CHECK-arm64e-SAME: ptr @_swift_malloc, +// CHECK-arm64e-SAME: i32 0, +// CHECK-arm64e-SAME: i64 0, +// CHECK-arm64e-SAME: i64 24469 } +// CHECK-arm64e-SAME: section "llvm.ptrauth" +// CHECK-arm64e-SAME: align 8 +// CHECK-arm64e-LABEL: _swift_free.ptrauth = private constant { +// CHECK-arm64e-SAME: ptr @_swift_free, +// CHECK-arm64e-SAME: i32 0, +// CHECK-arm64e-SAME: i64 0, +// CHECK-arm64e-SAME: i64 40879 }, +// CHECK-arm64e-SAME: section "llvm.ptrauth", +// CHECK-arm64e-SAME: align 8 // CHECK-LABEL: _swift_coro_malloc_allocator = linkonce_odr hidden constant %swift.coro_allocator { // CHECK-SAME: i32 258, -// CHECK-SAME: malloc, +// CHECK-SAME: malloc // CHECK-SAME: free // CHECK-SAME: } @@ -56,7 +84,10 @@ // CHECK-SAME: i32 0 // CHECK-SAME: i32 1 // CHECK: [[ALLOCATE_FN:%[^,]+]] = load ptr, ptr [[ALLOCATE_FN_PTR]] -// CHECK: [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]](i64 [[SIZE]]) +// CHECK-arm64e: [[ALLOCATE_FN_BITS:%[^,]+]] = ptrtoint ptr [[ALLOCATE_FN]] to i64 +// CHECK-arm64e: [[ALLOCATE_FN_BITS_AUTHED:%[^,]+]] = call i64 @llvm.ptrauth.auth(i64 [[ALLOCATE_FN_BITS]], i32 0, i64 24469) +// CHECK-arm64e: [[ALLOCATE_FN:%[^,]+]] = inttoptr i64 [[ALLOCATE_FN_BITS_AUTHED]] +// CHECK: [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]]([[INT]] [[SIZE]]) // CHECK: ret ptr [[ALLOCATION]] // CHECK: } @@ -91,6 +122,9 @@ // CHECK-SAME: i32 0 // CHECK-SAME: i32 2 // CHECK: [[DEALLOCATE_FN:%[^,]+]] = load ptr, ptr [[DEALLOCATE_FN_PTR]] +// CHECK-arm64e: [[DEALLOCATE_FN_BITS:%[^,]+]] = ptrtoint ptr [[DEALLOCATE_FN]] to i64 +// CHECK-arm64e: [[DEALLOCATE_FN_BITS_AUTHED:%[^,]+]] = call i64 @llvm.ptrauth.auth(i64 [[DEALLOCATE_FN_BITS]], i32 0, i64 40879) +// CHECK-arm64e: [[DEALLOCATE_FN:%[^,]+]] = inttoptr i64 [[DEALLOCATE_FN_BITS_AUTHED]] // CHECK: call swiftcc void [[DEALLOCATE_FN]](ptr [[ADDRESS]]) // CHECK: ret void // CHECK: }