Merge pull request #34882 from DougGregor/async-let-silgen

[Concurrency] Implement SIL generation for "async let".

Author: DougGregor

lib/SILGen/SILGen.cpp

@@ -327,10 +327,9 @@ SILGenModule::getConformanceToBridgedStoredNSError(SILLocation loc, Type type) {
   return SwiftModule->lookupConformance(type, proto);
 }
 
-static FuncDecl *
-lookUpResumeContinuationIntrinsic(ASTContext &C,
-                                  Optional<FuncDecl*> &cache,
-                                  StringRef name) {
+static FuncDecl *lookupConcurrencyIntrinsic(ASTContext &C,
+                                            Optional<FuncDecl*> &cache,
+                                            StringRef name) {
   if (cache)
     return *cache;
 
@@ -355,22 +354,43 @@ lookUpResumeContinuationIntrinsic(ASTContext &C,
   return func;
 }
 
+FuncDecl *
+SILGenModule::getRunChildTask() {
+  return lookupConcurrencyIntrinsic(getASTContext(),
+                                    RunChildTask,
+                                    "_runChildTask");
+}
+
+FuncDecl *
+SILGenModule::getTaskFutureGet() {
+  return lookupConcurrencyIntrinsic(getASTContext(),
+                                    TaskFutureGet,
+                                    "_taskFutureGet");
+}
+
+FuncDecl *
+SILGenModule::getTaskFutureGetThrowing() {
+  return lookupConcurrencyIntrinsic(getASTContext(),
+                                    TaskFutureGetThrowing,
+                                    "_taskFutureGetThrowing");
+}
+
 FuncDecl *
 SILGenModule::getResumeUnsafeContinuation() {
-  return lookUpResumeContinuationIntrinsic(getASTContext(),
-                                           ResumeUnsafeContinuation,
-                                           "_resumeUnsafeContinuation");
+  return lookupConcurrencyIntrinsic(getASTContext(),
+                                    ResumeUnsafeContinuation,
+                                    "_resumeUnsafeContinuation");
 }
 FuncDecl *
 SILGenModule::getResumeUnsafeThrowingContinuation() {
-  return lookUpResumeContinuationIntrinsic(getASTContext(),
-                                           ResumeUnsafeThrowingContinuation,
-                                           "_resumeUnsafeThrowingContinuation");
+  return lookupConcurrencyIntrinsic(getASTContext(),
+                                    ResumeUnsafeThrowingContinuation,
+                                    "_resumeUnsafeThrowingContinuation");
 }
 FuncDecl *
 SILGenModule::getResumeUnsafeThrowingContinuationWithError() {
-  return lookUpResumeContinuationIntrinsic(getASTContext(),
-                                 ResumeUnsafeThrowingContinuationWithError,
+  return lookupConcurrencyIntrinsic(getASTContext(),
+                                    ResumeUnsafeThrowingContinuationWithError,
                                  "_resumeUnsafeThrowingContinuationWithError");
 }
 

lib/SILGen/SILGen.h

@@ -119,6 +119,10 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
 
   Optional<ProtocolConformance *> NSErrorConformanceToError;
 
+  Optional<FuncDecl*> RunChildTask;
+  Optional<FuncDecl*> TaskFutureGet;
+  Optional<FuncDecl*> TaskFutureGetThrowing;
+
   Optional<FuncDecl*> ResumeUnsafeContinuation;
   Optional<FuncDecl*> ResumeUnsafeThrowingContinuation;
   Optional<FuncDecl*> ResumeUnsafeThrowingContinuationWithError;
@@ -473,6 +477,15 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// Retrieve the conformance of NSError to the Error protocol.
   ProtocolConformance *getNSErrorConformanceToError();
 
+  /// Retrieve the _Concurrency._runChildTask intrinsic.
+  FuncDecl *getRunChildTask();
+
+  /// Retrieve the _Concurrency._taskFutureGet intrinsic.
+  FuncDecl *getTaskFutureGet();
+
+  /// Retrieve the _Concurrency._taskFutureGetThrowing intrinsic.
+  FuncDecl *getTaskFutureGetThrowing();
+
   /// Retrieve the _Concurrency._resumeUnsafeContinuation intrinsic.
   FuncDecl *getResumeUnsafeContinuation();
   /// Retrieve the _Concurrency._resumeUnsafeThrowingContinuation intrinsic.

lib/SILGen/SILGenApply.cpp

@@ -5671,6 +5671,72 @@ SILGenFunction::emitCoroutineAccessor(SILLocation loc, SILDeclRef accessor,
   return endApplyHandle;
 }
 
+ManagedValue SILGenFunction::emitRunChildTask(
+    SILLocation loc, Type functionType, ManagedValue taskFunction) {
+  auto runChildTaskFn = SGM.getRunChildTask();
+
+  Type resultType = functionType->castTo<FunctionType>()->getResult();
+  Type replacementTypes[] = {resultType};
+  auto subs = SubstitutionMap::get(runChildTaskFn->getGenericSignature(),
+                                   replacementTypes,
+                                   ArrayRef<ProtocolConformanceRef>{});
+
+  CanType origParamType = runChildTaskFn->getParameters()->get(0)
+      ->getInterfaceType()->getCanonicalType();
+  CanType substParamType = origParamType.subst(subs)->getCanonicalType();
+
+  // Ensure that the closure has the appropriate type.
+  AbstractionPattern origParam(
+      runChildTaskFn->getGenericSignature().getCanonicalSignature(),
+      origParamType);
+  taskFunction = emitSubstToOrigValue(
+      loc, taskFunction, origParam, substParamType);
+
+  return emitApplyOfLibraryIntrinsic(
+      loc, runChildTaskFn, subs, {taskFunction}, SGFContext()
+    ).getScalarValue();
+}
+
+ManagedValue SILGenFunction::emitCancelAsyncTask(
+    SILLocation loc, SILValue task) {
+  ASTContext &ctx = getASTContext();
+  auto apply = B.createBuiltin(
+      loc,
+      ctx.getIdentifier(getBuiltinName(BuiltinValueKind::CancelAsyncTask)),
+      getLoweredType(ctx.TheEmptyTupleType), SubstitutionMap(),
+      { task });
+  return ManagedValue::forUnmanaged(apply);
+}
+
+void SILGenFunction::completeAsyncLetChildTask(
+    PatternBindingDecl *patternBinding, unsigned index) {
+  SILValue childTask;
+  bool isThrowing;
+  std::tie(childTask, isThrowing)= AsyncLetChildTasks[{patternBinding, index}];
+
+  Type childResultType = patternBinding->getPattern(index)->getType();
+
+  auto taskFutureGetFn = isThrowing
+      ? SGM.getTaskFutureGetThrowing()
+      : SGM.getTaskFutureGet();
+
+  // Get the result from the future.
+  Type replacementTypes[] = {childResultType};
+  auto subs = SubstitutionMap::get(taskFutureGetFn->getGenericSignature(),
+                                   replacementTypes,
+                                   ArrayRef<ProtocolConformanceRef>{});
+  RValue childResult = emitApplyOfLibraryIntrinsic(
+      SILLocation(patternBinding), taskFutureGetFn, subs,
+      { ManagedValue::forBorrowedObjectRValue(childTask) },
+      SGFContext());
+
+  // Write the child result into the pattern variables.
+  emitAssignToPatternVars(
+      SILLocation(patternBinding), patternBinding->getPattern(index),
+      std::move(childResult));
+}
+
+
 // Create a partial application of a dynamic method, applying bridging thunks
 // if necessary.
 static ManagedValue emitDynamicPartialApply(SILGenFunction &SGF,

lib/SILGen/SILGenBuiltin.cpp

@@ -1409,14 +1409,7 @@ static ManagedValue emitBuiltinGetCurrentAsyncTask(
 static ManagedValue emitBuiltinCancelAsyncTask(
     SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
     ArrayRef<ManagedValue> args, SGFContext C) {
-  ASTContext &ctx = SGF.getASTContext();
-  auto argument = args[0].borrow(SGF, loc).forward(SGF);
-  auto apply = SGF.B.createBuiltin(
-      loc,
-      ctx.getIdentifier(getBuiltinName(BuiltinValueKind::CancelAsyncTask)),
-      SGF.getLoweredType(ctx.TheEmptyTupleType), SubstitutionMap(),
-      { argument });
-  return ManagedValue::forUnmanaged(apply);
+  return SGF.emitCancelAsyncTask(loc, args[0].borrow(SGF, loc).forward(SGF));
 }
 
 // Emit SIL for the named builtin: createAsyncTask.

lib/SILGen/SILGenDecl.cpp

@@ -428,8 +428,8 @@ class LetValueInitialization : public Initialization {
     auto &lowering = SGF.getTypeLowering(vd->getType());
 
     // Decide whether we need a temporary stack buffer to evaluate this 'let'.
-    // There are three cases we need to handle here: parameters, initialized (or
-    // bound) decls, and uninitialized ones.
+    // There are four cases we need to handle here: parameters, initialized (or
+    // bound) decls, uninitialized ones, and async let declarations.
     bool needsTemporaryBuffer;
     bool isUninitialized = false;
 
@@ -440,6 +440,13 @@ class LetValueInitialization : public Initialization {
       // buffer.  DI will make sure it is only assigned to once.
       needsTemporaryBuffer = true;
       isUninitialized = true;
+    } else if (vd->isAsyncLet()) {
+      // If this is an async let, treat it like a let-value without an
+      // initializer. The initializer runs concurrently in a child task,
+      // and value will be initialized at the point the variable in the
+      // async let is used.
+      needsTemporaryBuffer = true;
+      isUninitialized = true;
     } else {
       // If this is a let with an initializer or bound value, we only need a
       // buffer if the type is address only.
@@ -1130,12 +1137,53 @@ SILGenFunction::emitInitializationForVarDecl(VarDecl *vd, bool forceImmutable) {
 
 void SILGenFunction::emitPatternBinding(PatternBindingDecl *PBD,
                                         unsigned idx) {
+
+
   auto initialization = emitPatternBindingInitialization(PBD->getPattern(idx),
                                                          JumpDest::invalid());
 
-  // If an initial value expression was specified by the decl, emit it into
-  // the initialization. Otherwise, mark it uninitialized for DI to resolve.
-  if (auto *Init = PBD->getExecutableInit(idx)) {
+  // If this is an async let, create a child task to compute the initializer
+  // value.
+  if (PBD->isAsyncLet()) {
+    // Look through the implicit await (if present), try (if present), and
+    // call to reach the autoclosure that computes the value.
+    auto *init = PBD->getExecutableInit(idx);
+    if (auto awaitExpr = dyn_cast<AwaitExpr>(init))
+      init = awaitExpr->getSubExpr();
+    if (auto tryExpr = dyn_cast<TryExpr>(init))
+      init = tryExpr->getSubExpr();
+    init = cast<CallExpr>(init)->getFn();
+    assert(isa<AutoClosureExpr>(init) &&
+           "Could not find async let autoclosure");
+    bool isThrowing = init->getType()->castTo<AnyFunctionType>()->isThrowing();
+
+    // Emit the closure for the child task.
+    SILValue childTask;
+    {
+      FullExpr Scope(Cleanups, CleanupLocation(init));
+      SILLocation loc(PBD);
+      childTask = emitRunChildTask(
+          loc,
+          init->getType(),
+          emitRValue(init).getScalarValue()
+        ).forward(*this);
+    }
+
+    // Destroy the task at the end of the scope.
+    enterDestroyCleanup(childTask);
+
+    // Push a cleanup that will cancel the child task at the end of the scope.
+    enterCancelAsyncTaskCleanup(childTask);
+
+    // Save the child task so we can await it as needed.
+    AsyncLetChildTasks[{PBD, idx}] = { childTask, isThrowing };
+
+    // Mark as uninitialized; actual initialization will occur when the
+    // variables are referenced.
+    initialization->finishUninitialized(*this);
+  } else if (auto *Init = PBD->getExecutableInit(idx)) {
+    // If an initial value expression was specified by the decl, emit it into
+    // the initialization.
     FullExpr Scope(Cleanups, CleanupLocation(Init));
 
     auto *var = PBD->getSingleVar();
@@ -1155,6 +1203,7 @@ void SILGenFunction::emitPatternBinding(PatternBindingDecl *PBD,
 
     emitExprInto(Init, initialization.get(), SILLocation(PBD));
   } else {
+    // Otherwise, mark it uninitialized for DI to resolve.
     initialization->finishUninitialized(*this);
   }
 }
@@ -1413,6 +1462,35 @@ CleanupHandle SILGenFunction::enterDeinitExistentialCleanup(
   return Cleanups.getTopCleanup();
 }
 
+namespace {
+  /// A cleanup that cancels an asynchronous task.
+  class CancelAsyncTaskCleanup: public Cleanup {
+    SILValue task;
+  public:
+    CancelAsyncTaskCleanup(SILValue task) : task(task) { }
+
+    void emit(SILGenFunction &SGF, CleanupLocation l,
+              ForUnwind_t forUnwind) override {
+      SILValue borrowedTask = SGF.B.createBeginBorrow(l, task);
+      SGF.emitCancelAsyncTask(l, borrowedTask);
+      SGF.B.createEndBorrow(l, borrowedTask);
+    }
+
+    void dump(SILGenFunction &) const override {
+#ifndef NDEBUG
+      llvm::errs() << "CancelAsyncTaskCleanup\n"
+                   << "Task:" << task << "\n";
+#endif
+    }
+  };
+} // end anonymous namespace
+
+CleanupHandle SILGenFunction::enterCancelAsyncTaskCleanup(SILValue task) {
+  Cleanups.pushCleanupInState<CancelAsyncTaskCleanup>(
+      CleanupState::Active, task);
+  return Cleanups.getTopCleanup();
+}
+
 /// Create a LocalVariableInitialization for the uninitialized var.
 InitializationPtr SILGenFunction::emitLocalVariableWithCleanup(
     VarDecl *vd, Optional<MarkUninitializedInst::Kind> kind, unsigned ArgNo) {