Index: llvm/lib/Transforms/Coroutines/CoroInstr.h =================================================================== --- llvm/lib/Transforms/Coroutines/CoroInstr.h +++ llvm/lib/Transforms/Coroutines/CoroInstr.h @@ -25,6 +25,7 @@ #ifndef LLVM_LIB_TRANSFORMS_COROUTINES_COROINSTR_H #define LLVM_LIB_TRANSFORMS_COROUTINES_COROINSTR_H +#include "llvm/IR/GlobalPtrAuthInfo.h" #include "llvm/IR/GlobalVariable.h" #include "llvm/IR/IntrinsicInst.h" #include "llvm/Support/raw_ostream.h" @@ -223,7 +224,16 @@ /// attributes, and calling convention of the continuation function(s) /// are taken from this declaration. Function *getPrototype() const { - return cast(getArgOperand(PrototypeArg)->stripPointerCasts()); + Value *rawPrototype = getArgOperand(PrototypeArg)->stripPointerCasts(); + if (auto global = GlobalPtrAuthInfo::analyze(rawPrototype)) { + rawPrototype = + const_cast(global->getPointer()->stripPointerCasts()); + } + return cast(rawPrototype); + } + + Optional getPtrAuthInfo() const { + return GlobalPtrAuthInfo::analyze(getArgOperand(PrototypeArg)); } /// Return the function to use for allocating memory. Index: llvm/lib/Transforms/Coroutines/CoroInternal.h =================================================================== --- llvm/lib/Transforms/Coroutines/CoroInternal.h +++ llvm/lib/Transforms/Coroutines/CoroInternal.h @@ -118,6 +118,7 @@ }; struct RetconLoweringStorage { + Optional ResumePtrAuthInfo; Function *ResumePrototype; Function *Alloc; Function *Dealloc; @@ -152,8 +153,18 @@ PointerType *getSwitchResumePointerType() const { assert(ABI == coro::ABI::Switch); - assert(FrameTy && "frame type not assigned"); - return cast(FrameTy->getElementType(SwitchFieldIndex::Resume)); + assert(FrameTy && "frame type not assigned"); + return cast(FrameTy->getElementType(SwitchFieldIndex::Resume)); + } + + Optional getResumePtrAuthInfo() const { + switch (ABI) { + case coro::ABI::Switch: + return None; + case coro::ABI::Retcon: + case coro::ABI::RetconOnce: + return RetconLowering.ResumePtrAuthInfo; + } } FunctionType *getResumeFunctionType() const { @@ -231,7 +242,6 @@ /// \param CG - if non-null, will be updated for the new call void emitDealloc(IRBuilder<> &Builder, Value *Ptr, CallGraph *CG) const; - Shape() = default; explicit Shape(Function &F) { buildFrom(F); } void buildFrom(Function &F); }; Index: llvm/lib/Transforms/Coroutines/CoroSplit.cpp =================================================================== --- llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -1313,9 +1313,16 @@ Builder.CreateRet(RetV); } + // Sign the continuation value if requested. + Constant *ContinuationValue = Continuation; + if (auto PtrAuthInfo = Shape.getResumePtrAuthInfo()) { + ContinuationValue = + PtrAuthInfo->createWithSameSchema(*F.getParent(), ContinuationValue); + } + // Branch to the return block. Branch->setSuccessor(0, ReturnBB); - ReturnPHIs[0]->addIncoming(Continuation, SuspendBB); + ReturnPHIs[0]->addIncoming(ContinuationValue, SuspendBB); size_t NextPHIIndex = 1; for (auto &VUse : Suspend->value_operands()) ReturnPHIs[NextPHIIndex++]->addIncoming(&*VUse, SuspendBB); Index: llvm/lib/Transforms/Coroutines/Coroutines.cpp =================================================================== --- llvm/lib/Transforms/Coroutines/Coroutines.cpp +++ llvm/lib/Transforms/Coroutines/Coroutines.cpp @@ -382,11 +382,21 @@ : coro::ABI::RetconOnce); auto Prototype = ContinuationId->getPrototype(); this->RetconLowering.ResumePrototype = Prototype; + this->RetconLowering.ResumePtrAuthInfo = ContinuationId->getPtrAuthInfo(); this->RetconLowering.Alloc = ContinuationId->getAllocFunction(); this->RetconLowering.Dealloc = ContinuationId->getDeallocFunction(); this->RetconLowering.ReturnBlock = nullptr; this->RetconLowering.IsFrameInlineInStorage = false; + if (RetconLowering.ResumePtrAuthInfo && + RetconLowering.ResumePtrAuthInfo->hasAddressDiversity()) { +#ifndef NDEBUG + RetconLowering.ResumePtrAuthInfo->getGV()->dump(); +#endif + report_fatal_error("ptrauth-signed prototype must not have address " + "diversity"); + } + // Determine the result value types, and make sure they match up with // the values passed to the suspends. auto ResultTys = getRetconResultTypes(); @@ -553,6 +563,9 @@ /// Check that the given value is a well-formed prototype for the /// llvm.coro.id.retcon.* intrinsics. static void checkWFRetconPrototype(const AnyCoroIdRetconInst *I, Value *V) { + if (auto PtrAuth = GlobalPtrAuthInfo::analyze(V)) { + V = const_cast(PtrAuth->getPointer()); + } auto F = dyn_cast(V->stripPointerCasts()); if (!F) fail(I, "llvm.coro.id.retcon.* prototype not a Function", V); Index: llvm/test/Transforms/Coroutines/coro-retcon-ptrauth.ll =================================================================== --- /dev/null +++ llvm/test/Transforms/Coroutines/coro-retcon-ptrauth.ll @@ -0,0 +1,60 @@ +; RUN: opt -coro-early -coro-split -coro-cleanup -S %s | FileCheck %s +target datalayout = "E-p:64:64" + +%swift.type = type { i64 } +%swift.opaque = type opaque +%T4red215EmptyCollectionV = type opaque +%TSi = type <{ i64 }> + +@prototype.signed = private constant { i8*, i32, i64, i64 } { i8* bitcast (i8* (i8*, i1)* @prototype to i8*), i32 2, i64 0, i64 12867 }, section "llvm.ptrauth" + +define i8* @f(i8* %buffer, i32 %n) { +entry: + %id = call token @llvm.coro.id.retcon(i32 8, i32 4, i8* %buffer, i8* bitcast ({ i8*, i32, i64, i64 }* @prototype.signed to i8*), i8* bitcast (i8* (i32)* @allocate to i8*), i8* bitcast (void (i8*)* @deallocate to i8*)) + %hdl = call i8* @llvm.coro.begin(token %id, i8* null) + br label %loop + +loop: + %n.val = phi i32 [ %n, %entry ], [ %inc, %resume ] + call void @print(i32 %n.val) + %unwind0 = call i1 (...) @llvm.coro.suspend.retcon.i1() + br i1 %unwind0, label %cleanup, label %resume + +resume: + %inc = add i32 %n.val, 1 + br label %loop + +cleanup: + call i1 @llvm.coro.end(i8* %hdl, i1 0) + unreachable +} + +; CHECK: @prototype.signed = private constant +; CHECK: [[GLOBAL:@.*]] = private constant { i8*, i32, i64, i64 } { i8* bitcast (i8* (i8*, i1)* [[RESUME:@.*]] to i8*), i32 2, i64 0, i64 12867 }, section "llvm.ptrauth" + +; CHECK-LABEL: define i8* @f(i8* %buffer, i32 %n) +; CHECK-NEXT: entry: +; CHECK-NEXT: [[T0:%.*]] = bitcast i8* %buffer to [[FRAME_T:%.*]]* +; CHECK-NEXT: [[T1:%.*]] = getelementptr inbounds [[FRAME_T]], [[FRAME_T]]* [[T0]], i32 0, i32 0 +; CHECK-NEXT: store i32 %n, i32* [[T1]] +; CHECK-NEXT: call void @print(i32 %n) +; CHECK-NEXT: ret i8* bitcast ({ i8*, i32, i64, i64 }* [[GLOBAL]] to i8*) + +; CHECK: define internal i8* [[RESUME]](i8* noalias nonnull, i1 zeroext) { +; CHECK: bitcast ({ i8*, i32, i64, i64 }* [[GLOBAL]] to + +declare noalias i8* @malloc(i64) #5 +declare void @free(i8* nocapture) #5 + +declare token @llvm.coro.id.retcon(i32, i32, i8*, i8*, i8*, i8*) +declare i8* @llvm.coro.begin(token, i8*) +declare i1 @llvm.coro.suspend.retcon.i1(...) +declare i1 @llvm.coro.end(i8*, i1) +declare i8* @llvm.coro.prepare.retcon(i8*) + +declare i8* @prototype(i8*, i1 zeroext) + +declare noalias i8* @allocate(i32 %size) +declare void @deallocate(i8* %ptr) + +declare void @print(i32)