Merge pull request #34595 from DougGregor/get-current-async-task-builtin

[Concurrency] Add a builtin to get the current task in an async function

Author: DougGregor

include/swift/AST/Builtins.def

@@ -717,6 +717,12 @@ BUILTIN_MISC_OPERATION(IntInstrprofIncrement, "int_instrprof_increment", "", Spe
   BUILTIN_MISC_OPERATION(Id, Name, Attrs, Overload)
 #endif
 
+// getCurrentAsyncTask: () -> Builtin.NativeObject
+//
+// Retrieve the pointer to the task in which the current asynchronous
+// function is executing.
+BUILTIN_MISC_OPERATION_WITH_SILGEN(GetCurrentAsyncTask, "getCurrentAsyncTask", "n", 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

lib/AST/Builtins.cpp

@@ -1341,6 +1341,10 @@ static ValueDecl *getConvertUnownedUnsafeToGuaranteed(ASTContext &ctx,
   return builder.build(id);
 }
 
+static ValueDecl *getGetCurrentAsyncTask(ASTContext &ctx, Identifier id) {
+  return getBuiltinFunction(id, { }, ctx.TheNativeObjectType);
+}
+
 static ValueDecl *getPoundAssert(ASTContext &Context, Identifier Id) {
   auto int1Type = BuiltinIntegerType::get(1, Context);
   auto optionalRawPointerType = BoundGenericEnumType::get(
@@ -2467,6 +2471,9 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
   case BuiltinValueKind::ConvertUnownedUnsafeToGuaranteed:
     return getConvertUnownedUnsafeToGuaranteed(Context, Id);
 
+  case BuiltinValueKind::GetCurrentAsyncTask:
+    return getGetCurrentAsyncTask(Context, Id);
+
   case BuiltinValueKind::PoundAssert:
     return getPoundAssert(Context, Id);
 

lib/IRGen/GenBuiltin.cpp

@@ -211,6 +211,13 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
     return;
   }
 
+  // getCurrentAsyncTask has no arguments.
+  if (Builtin.ID == BuiltinValueKind::GetCurrentAsyncTask) {
+    auto task = IGF.getAsyncTask();
+    out.add(IGF.Builder.CreateBitCast(task, IGF.IGM.RefCountedPtrTy));
+    return;
+  }
+
   // Everything else cares about the (rvalue) argument.
 
   // If this is an LLVM IR intrinsic, lower it to an intrinsic call.

lib/SIL/IR/OperandOwnership.cpp

@@ -1035,6 +1035,14 @@ CONSTANT_OWNERSHIP_BUILTIN(Owned, MustBeInvalidated, COWBufferForReading)
 CONSTANT_OWNERSHIP_BUILTIN(Owned, MustBeInvalidated, UnsafeGuaranteed)
 #undef CONSTANT_OWNERSHIP_BUILTIN
 
+#define SHOULD_NEVER_VISIT_BUILTIN(ID)                              \
+  OperandOwnershipKindMap OperandOwnershipKindBuiltinClassifier::visit##ID(    \
+      BuiltinInst *, StringRef) {                                              \
+    llvm_unreachable("Builtin should never be visited! E.x.: It may not have arguments"); \
+  }
+SHOULD_NEVER_VISIT_BUILTIN(GetCurrentAsyncTask)
+#undef SHOULD_NEVER_VISIT_BUILTIN
+
 // Builtins that should be lowered to SIL instructions so we should never see
 // them.
 #define BUILTIN_SIL_OPERATION(ID, NAME, CATEGORY)                              \

lib/SIL/IR/ValueOwnership.cpp

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

lib/SIL/Verifier/SILVerifier.cpp

@@ -1736,6 +1736,14 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
       verifyLLVMIntrinsic(BI, BI->getIntrinsicInfo().ID);
       return;
     }
+
+    // Check that 'getCurrentAsyncTask' only occurs within an async function.
+    if (BI->getBuiltinKind() &&
+        *BI->getBuiltinKind() == BuiltinValueKind::GetCurrentAsyncTask) {
+      require(F.isAsync(),
+          "getCurrentAsyncTask builtin can only be used in an async function");
+      return;
+    }
   }
 
   void checkFunctionRefBaseInst(FunctionRefBaseInst *FRI) {

lib/SILGen/SILGenBuiltin.cpp

@@ -1393,6 +1393,18 @@ static ManagedValue emitBuiltinConvertUnownedUnsafeToGuaranteed(
   return SGF.B.createMarkDependence(loc, guaranteedNonTrivialRefMV, baseMV);
 }
 
+// Emit SIL for the named builtin: getCurrentAsyncTask.
+static ManagedValue emitBuiltinGetCurrentAsyncTask(
+    SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
+    PreparedArguments &&preparedArgs, SGFContext C) {
+  ASTContext &ctx = SGF.getASTContext();
+  auto apply = SGF.B.createBuiltin(
+      loc,
+      ctx.getIdentifier(getBuiltinName(BuiltinValueKind::GetCurrentAsyncTask)),
+      SGF.getLoweredType(ctx.TheNativeObjectType), SubstitutionMap(), { });
+  return SGF.emitManagedRValueWithEndLifetimeCleanup(apply);
+}
+
 Optional<SpecializedEmitter>
 SpecializedEmitter::forDecl(SILGenModule &SGM, SILDeclRef function) {
   // Only consider standalone declarations in the Builtin module.

lib/SILGen/SILGenDecl.cpp

@@ -1326,6 +1326,33 @@ CleanupHandle SILGenFunction::enterDestroyCleanup(SILValue valueOrAddr) {
   return Cleanups.getTopCleanup();
 }
 
+namespace {
+class EndLifetimeCleanup : public Cleanup {
+  SILValue v;
+public:
+  EndLifetimeCleanup(SILValue v) : v(v) {}
+
+  void emit(SILGenFunction &SGF, CleanupLocation l,
+            ForUnwind_t forUnwind) override {
+    SGF.B.createEndLifetime(l, v);
+  }
+
+  void dump(SILGenFunction &) const override {
+#ifndef NDEBUG
+    llvm::errs() << "EndLifetimeCleanup\n"
+                 << "State:" << getState() << "\n"
+                 << "Value:" << v << "\n";
+#endif
+  }
+};
+} // end anonymous namespace
+
+ManagedValue SILGenFunction::emitManagedRValueWithEndLifetimeCleanup(
+    SILValue value) {
+  Cleanups.pushCleanup<EndLifetimeCleanup>(value);
+  return ManagedValue::forUnmanaged(value);
+}
+
 namespace {
   /// A cleanup that deinitializes an opaque existential container
   /// before a value has been stored into it, or after its value was taken.

lib/SILGen/SILGenFunction.h

@@ -2000,7 +2000,21 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
 
   /// Enter a cleanup to emit a ReleaseValue/DestroyAddr of the specified value.
   CleanupHandle enterDestroyCleanup(SILValue valueOrAddr);
-  
+
+  /// Return an owned managed value for \p value that is cleaned up using an end_lifetime instruction.
+  ///
+  /// The end_lifetime cleanup is not placed into the ManagedValue itself and
+  /// thus can not be forwarded. This means that the ManagedValue is treated
+  /// as a +0 value. This means that the owned value will be copied by SILGen
+  /// if it is ever needed as a +1 value (meaning any time that the value
+  /// escapes).
+  ///
+  /// DISCUSSION: end_lifetime ends the lifetime of an owned value in OSSA
+  /// without resulting in a destroy being emitted. This cleanup should only
+  /// be used for owned values that do not need to be destroyed if they do not
+  /// escape the current call frame but need to be copied if they escape.
+  ManagedValue emitManagedRValueWithEndLifetimeCleanup(SILValue value);
+
   /// Enter a cleanup to emit a DeinitExistentialAddr or DeinitExistentialBox
   /// of the specified value.
   CleanupHandle enterDeinitExistentialCleanup(CleanupState state,

lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp

@@ -140,6 +140,7 @@ static bool isBarrier(SILInstruction *inst) {
     case BuiltinValueKind::GlobalStringTablePointer:
     case BuiltinValueKind::COWBufferForReading:
     case BuiltinValueKind::IntInstrprofIncrement:
+    case BuiltinValueKind::GetCurrentAsyncTask:
       return false;
 
     // Handle some rare builtins that may be sensitive to object lifetime

test/IRGen/async/builtins.sil

@@ -0,0 +1,19 @@
+// RUN: %target-swift-frontend -enable-experimental-concurrency -enable-objc-interop  -primary-file %s -emit-ir | %FileCheck %s -DINT=i%target-ptrsize --check-prefixes=CHECK,CHECK-objc
+// RUN: %target-swift-frontend -enable-experimental-concurrency -disable-objc-interop -primary-file %s -emit-ir | %FileCheck %s -DINT=i%target-ptrsize --check-prefixes=CHECK,CHECK-native
+
+// REQUIRES: concurrency
+
+sil_stage canonical
+
+import Builtin
+
+// CHECK-LABEL: define hidden swiftcc void @get_task(%swift.task* %0, %swift.executor* %1, %swift.context* %2)
+sil hidden [ossa] @get_task : $@async @convention(thin) () -> @owned Builtin.NativeObject {
+bb0:
+  // CHECK: [[TASK:%.*]] = bitcast %swift.task* %0 to %swift.refcounted*
+  %0 = builtin "getCurrentAsyncTask"() : $Builtin.NativeObject
+  // CHECK-NEXT: [[TASK_COPY:%.*]] = call %swift.refcounted* @swift_retain(%swift.refcounted* returned [[TASK]])
+  %1 = copy_value %0 : $Builtin.NativeObject
+  end_lifetime %0 : $Builtin.NativeObject
+  return %1 : $Builtin.NativeObject
+}

test/SILGen/async_builtins.swift

@@ -0,0 +1,10 @@
+// RUN: %target-swift-frontend -emit-silgen %s -module-name test -swift-version 5 -enable-experimental-concurrency -parse-stdlib | %FileCheck %s
+// REQUIRES: concurrency
+
+struct X {
+  // CHECK-LABEL: sil hidden [ossa] @$s4test1XV14getCurrentTaskBoyYF
+  func getCurrentTask() async -> Builtin.NativeObject {
+    // CHECK: builtin "getCurrentAsyncTask"() : $Builtin.NativeObject
+    return Builtin.getCurrentAsyncTask()
+  }
+}