Merge pull request swiftlang#34600 from DougGregor/async-task-builtins

[Concurrency] Add cancelAsyncTask() builtin.

Author: DougGregor

include/swift/ABI/Task.h

@@ -182,7 +182,7 @@ class AsyncTask : public HeapObject, public Job {
   /// prevent it from being corrupted in flight.
   AsyncContext * __ptrauth_swift_task_resume_context ResumeContext;
 
-  /// The currntly-active information about cancellation.
+  /// The currently-active information about cancellation.
   std::atomic<ActiveTaskStatus> Status;
 
   /// Reserved for the use of the task-local stack allocator.

include/swift/AST/Builtins.def

@@ -723,6 +723,12 @@ BUILTIN_MISC_OPERATION(IntInstrprofIncrement, "int_instrprof_increment", "", Spe
 // function is executing.
 BUILTIN_MISC_OPERATION_WITH_SILGEN(GetCurrentAsyncTask, "getCurrentAsyncTask", "n", Special)
 
+/// cancelAsyncTask(): (Builtin.NativeObject) -> Void
+///
+/// Cancel the given asynchronous task.
+BUILTIN_MISC_OPERATION_WITH_SILGEN(CancelAsyncTask, "cancelAsyncTask", "", Special)
+
+
 /// globalStringTablePointer has type String -> Builtin.RawPointer.
 /// It returns an immortal, global string table pointer for strings constructed
 /// from string literals. We consider it effects as readnone meaning that it

include/swift/Runtime/RuntimeFunctions.def

@@ -1489,6 +1489,14 @@ FUNCTION(TaskDealloc,
          ARGS(SwiftTaskPtrTy, Int8PtrTy),
          ATTRS(NoUnwind, ArgMemOnly))
 
+// void swift_task_dealloc(AsyncTask *task, void *ptr);
+FUNCTION(TaskCancel,
+         swift_task_cancel, SwiftCC,
+         ConcurrencyAvailability,
+         RETURNS(VoidTy),
+         ARGS(SwiftTaskPtrTy),
+         ATTRS(NoUnwind, ArgMemOnly))
+
 #undef RETURNS
 #undef ARGS
 #undef ATTRS

lib/AST/Builtins.cpp

@@ -1345,6 +1345,11 @@ static ValueDecl *getGetCurrentAsyncTask(ASTContext &ctx, Identifier id) {
   return getBuiltinFunction(id, { }, ctx.TheNativeObjectType);
 }
 
+static ValueDecl *getCancelAsyncTask(ASTContext &ctx, Identifier id) {
+  return getBuiltinFunction(
+      id, { ctx.TheNativeObjectType }, ctx.TheEmptyTupleType);
+}
+
 static ValueDecl *getPoundAssert(ASTContext &Context, Identifier Id) {
   auto int1Type = BuiltinIntegerType::get(1, Context);
   auto optionalRawPointerType = BoundGenericEnumType::get(
@@ -2474,6 +2479,9 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
   case BuiltinValueKind::GetCurrentAsyncTask:
     return getGetCurrentAsyncTask(Context, Id);
 
+  case BuiltinValueKind::CancelAsyncTask:
+    return getCancelAsyncTask(Context, Id);
+
   case BuiltinValueKind::PoundAssert:
     return getPoundAssert(Context, Id);
 

lib/IRGen/GenBuiltin.cpp

@@ -220,6 +220,11 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
 
   // Everything else cares about the (rvalue) argument.
 
+  if (Builtin.ID == BuiltinValueKind::CancelAsyncTask) {
+    emitTaskCancel(IGF, args.claimNext());
+    return;
+  }
+
   // If this is an LLVM IR intrinsic, lower it to an intrinsic call.
   const IntrinsicInfo &IInfo = IGF.getSILModule().getIntrinsicInfo(FnId);
   llvm::Intrinsic::ID IID = IInfo.ID;

lib/IRGen/GenCall.cpp

@@ -3458,6 +3458,18 @@ void irgen::emitTaskDealloc(IRGenFunction &IGF, Address address,
                      llvm::Attribute::ReadNone);
 }
 
+void irgen::emitTaskCancel(IRGenFunction &IGF, llvm::Value *task) {
+  if (task->getType() != IGF.IGM.SwiftTaskPtrTy) {
+    task = IGF.Builder.CreateBitCast(task, IGF.IGM.SwiftTaskPtrTy);
+  }
+
+  auto *call = IGF.Builder.CreateCall(IGF.IGM.getTaskCancelFn(), {task});
+  call->setDoesNotThrow();
+  call->setCallingConv(IGF.IGM.SwiftCC);
+  call->addAttribute(llvm::AttributeList::FunctionIndex,
+                     llvm::Attribute::ReadNone);
+}
+
 std::pair<Address, Size> irgen::emitAllocAsyncContext(IRGenFunction &IGF,
                                                       AsyncContextLayout layout,
                                                       llvm::Value *sizeValue,

lib/IRGen/GenCall.h

@@ -388,6 +388,8 @@ namespace irgen {
   Address emitTaskAlloc(IRGenFunction &IGF, llvm::Value *size,
                         Alignment alignment);
   void emitTaskDealloc(IRGenFunction &IGF, Address address, llvm::Value *size);
+  void emitTaskCancel(IRGenFunction &IGF, llvm::Value *task);
+
   /// Allocate task local storage for the specified layout but using the
   /// provided dynamic size.  Allowing the size to be specified dynamically is
   /// necessary for applies of thick functions the sizes of whose async contexts

lib/SIL/IR/OperandOwnership.cpp

@@ -231,6 +231,7 @@ CONSTANT_OWNERSHIP_INST(None, MustBeLive, UncheckedTakeEnumDataAddr)
 CONSTANT_OWNERSHIP_INST(None, MustBeLive, UnconditionalCheckedCastAddr)
 CONSTANT_OWNERSHIP_INST(None, MustBeLive, AllocValueBuffer)
 CONSTANT_OWNERSHIP_INST(None, MustBeLive, DeallocValueBuffer)
+
 #define NEVER_LOADABLE_CHECKED_REF_STORAGE(Name, ...)                          \
   CONSTANT_OWNERSHIP_INST(None, MustBeLive, Load##Name)
 #define ALWAYS_LOADABLE_CHECKED_REF_STORAGE(Name, ...)                         \
@@ -1033,6 +1034,8 @@ ANY_OWNERSHIP_BUILTIN(IntInstrprofIncrement)
   }
 CONSTANT_OWNERSHIP_BUILTIN(Owned, MustBeInvalidated, COWBufferForReading)
 CONSTANT_OWNERSHIP_BUILTIN(Owned, MustBeInvalidated, UnsafeGuaranteed)
+CONSTANT_OWNERSHIP_BUILTIN(Guaranteed, MustBeLive, CancelAsyncTask)
+
 #undef CONSTANT_OWNERSHIP_BUILTIN
 
 #define SHOULD_NEVER_VISIT_BUILTIN(ID)                              \

lib/SIL/IR/ValueOwnership.cpp

@@ -542,6 +542,7 @@ CONSTANT_OWNERSHIP_BUILTIN(None, TypePtrAuthDiscriminator)
 CONSTANT_OWNERSHIP_BUILTIN(None, IntInstrprofIncrement)
 CONSTANT_OWNERSHIP_BUILTIN(None, GlobalStringTablePointer)
 CONSTANT_OWNERSHIP_BUILTIN(Owned, GetCurrentAsyncTask)
+CONSTANT_OWNERSHIP_BUILTIN(None, CancelAsyncTask)
 
 #undef CONSTANT_OWNERSHIP_BUILTIN
 

lib/SIL/Utils/MemAccessUtils.cpp

@@ -1799,6 +1799,7 @@ static void visitBuiltinAddress(BuiltinInst *builtin,
     case BuiltinValueKind::PoundAssert:
     case BuiltinValueKind::IntInstrprofIncrement:
     case BuiltinValueKind::TSanInoutAccess:
+    case BuiltinValueKind::CancelAsyncTask:
       return;
 
     // General memory access to a pointer in first operand position.