SILGen: Emit async let entry points at correct abstraction level.

Avoid a reabstraction thunk every time an async let entry point is emitted,
by setting the context abstraction level while we emit the implicit closure.
Adjust some surrounding logic that breaks when we do this:

- When lowering a non-throwing function type against a throwing abstraction
  pattern, include the error type in the lowered substituted type. Async
  throwing and nonthrowing functions would require another thunk to convert
  away the throwingness of the async context, which would defeat the purpose.
- Adjust the code in IRGen that pads the initial context size for `async let`
  entry points so that it works when the entry point has not yet emitted, by
  marking the async function pointer to be padded later if it isn't defined
  yet.

Author: jckarter

lib/IRGen/GenCall.cpp

@@ -3897,23 +3897,58 @@ emitRetconCoroutineEntry(IRGenFunction &IGF, CanSILFunctionType fnType,
   IGF.setEarliestInsertionPoint(pt);
 }
 
+void IRGenModule::addAsyncCoroIDMapping(llvm::GlobalVariable *asyncFunctionPointer,
+                                        llvm::CallInst *coro_id_builtin) {
+  AsyncCoroIDsForPadding[asyncFunctionPointer] = coro_id_builtin;
+}
+
+llvm::CallInst *
+IRGenModule::getAsyncCoroIDMapping(llvm::GlobalVariable *asyncFunctionPointer) {
+  auto found = AsyncCoroIDsForPadding.find(asyncFunctionPointer);
+  if (found == AsyncCoroIDsForPadding.end())
+    return nullptr;
+  return found->second;
+}
+
+void IRGenModule::markAsyncFunctionPointerForPadding(
+                                   llvm::GlobalVariable *asyncFunctionPointer) {
+  AsyncCoroIDsForPadding[asyncFunctionPointer] = nullptr;
+}
+
+bool IRGenModule::isAsyncFunctionPointerMarkedForPadding(
+                                   llvm::GlobalVariable *asyncFunctionPointer) {
+  auto found = AsyncCoroIDsForPadding.find(asyncFunctionPointer);
+  if (found == AsyncCoroIDsForPadding.end())
+    return false;
+  return found->second == nullptr;
+}
+
 void irgen::emitAsyncFunctionEntry(IRGenFunction &IGF,
                                    const AsyncContextLayout &layout,
                                    LinkEntity asyncFunction,
                                    unsigned asyncContextIndex) {
   auto &IGM = IGF.IGM;
   auto size = layout.getSize();
   auto asyncFuncPointerVar = cast<llvm::GlobalVariable>(IGM.getAddrOfAsyncFunctionPointer(asyncFunction));
+  bool isPadded = IGM
+    .isAsyncFunctionPointerMarkedForPadding(asyncFuncPointerVar);
   auto asyncFuncPointer = IGF.Builder.CreateBitOrPointerCast(
                                            asyncFuncPointerVar, IGM.Int8PtrTy);
+  
+  if (isPadded) {
+    size = std::max(layout.getSize(),
+                    NumWords_AsyncLet * IGM.getPointerSize());
+  }
+  
   auto *id = IGF.Builder.CreateIntrinsicCall(
       llvm::Intrinsic::coro_id_async,
       {llvm::ConstantInt::get(IGM.Int32Ty, size.getValue()),
        llvm::ConstantInt::get(IGM.Int32Ty, 16),
        llvm::ConstantInt::get(IGM.Int32Ty, asyncContextIndex),
        asyncFuncPointer});
-  auto inserted = IGM.AsyncCoroIDs.insert({asyncFuncPointerVar, id});
-  assert(inserted.second);
+  
+  IGM.addAsyncCoroIDMapping(asyncFuncPointerVar, id);
+
   // Call 'llvm.coro.begin', just for consistency with the normal pattern.
   // This serves as a handle that we can pass around to other intrinsics.
   auto hdl = IGF.Builder.CreateIntrinsicCall(

lib/IRGen/GenConcurrency.cpp

@@ -181,6 +181,9 @@ llvm::Value *irgen::emitBuiltinStartAsyncLet(IRGenFunction &IGF,
                                              llvm::Value *localContextInfo,
                                              llvm::Value *localResultBuffer,
                                              SubstitutionMap subs) {
+  localContextInfo = IGF.Builder.CreateBitCast(localContextInfo,
+                                               IGF.IGM.OpaquePtrTy);
+  
   // stack allocate AsyncLet, and begin lifetime for it (until EndAsyncLet)
   auto ty = llvm::ArrayType::get(IGF.IGM.Int8PtrTy, NumWords_AsyncLet);
   auto address = IGF.createAlloca(ty, Alignment(Alignment_AsyncLet));
@@ -204,31 +207,35 @@ llvm::Value *irgen::emitBuiltinStartAsyncLet(IRGenFunction &IGF,
   auto deploymentAvailability
     = AvailabilityContext::forDeploymentTarget(IGF.IGM.Context);
   if (!deploymentAvailability.isContainedIn(
-                                   IGF.IGM.Context.getSwift57Availability())) {
+                                   IGF.IGM.Context.getSwift57Availability()))
+  {
     auto taskAsyncFunctionPointer
                 = cast<llvm::GlobalVariable>(taskFunction->stripPointerCasts());
 
-    auto taskAsyncIDIter = IGF.IGM.AsyncCoroIDs.find(taskAsyncFunctionPointer);
-    assert(taskAsyncIDIter != IGF.IGM.AsyncCoroIDs.end()
-           && "async let entry point not emitted locally");
-    auto taskAsyncID = taskAsyncIDIter->second;
-    
-    // Pad out the initial context size in the async function pointer record
-    // and ID intrinsic so that it will never fit in the preallocated space.
-    uint64_t origSize = cast<llvm::ConstantInt>(taskAsyncID->getArgOperand(0))
-      ->getValue().getLimitedValue();
-    
-    uint64_t paddedSize = std::max(origSize,
+    if (auto taskAsyncID
+          = IGF.IGM.getAsyncCoroIDMapping(taskAsyncFunctionPointer)) {
+      // If the entry point function has already been emitted, retroactively
+      // pad out the initial context size in the async function pointer record
+      // and ID intrinsic so that it will never fit in the preallocated space.
+      uint64_t origSize = cast<llvm::ConstantInt>(taskAsyncID->getArgOperand(0))
+        ->getValue().getLimitedValue();
+      
+      uint64_t paddedSize = std::max(origSize,
                      (NumWords_AsyncLet * IGF.IGM.getPointerSize()).getValue());
-    auto paddedSizeVal = llvm::ConstantInt::get(IGF.IGM.Int32Ty, paddedSize);
-    taskAsyncID->setArgOperand(0, paddedSizeVal);
-    
-    auto origInit = taskAsyncFunctionPointer->getInitializer();
-    auto newInit = llvm::ConstantStruct::get(
+      auto paddedSizeVal = llvm::ConstantInt::get(IGF.IGM.Int32Ty, paddedSize);
+      taskAsyncID->setArgOperand(0, paddedSizeVal);
+      
+      auto origInit = taskAsyncFunctionPointer->getInitializer();
+      auto newInit = llvm::ConstantStruct::get(
                                    cast<llvm::StructType>(origInit->getType()),
                                    origInit->getAggregateElement(0u),
                                    paddedSizeVal);
-    taskAsyncFunctionPointer->setInitializer(newInit);
+      taskAsyncFunctionPointer->setInitializer(newInit);
+    } else {
+      // If it hasn't been emitted yet, mark it to get the padding when it does
+      // get emitted.
+      IGF.IGM.markAsyncFunctionPointerForPadding(taskAsyncFunctionPointer);
+    }
   }
 
   llvm::CallInst *call;

lib/IRGen/GenMeta.cpp

@@ -5611,6 +5611,12 @@ llvm::GlobalValue *irgen::emitAsyncFunctionPointer(IRGenModule &IGM,
                                                    llvm::Function *function,
                                                    LinkEntity entity,
                                                    Size size) {
+  auto afp = cast<llvm::GlobalVariable>(IGM.getAddrOfAsyncFunctionPointer(entity));
+  if (IGM.isAsyncFunctionPointerMarkedForPadding(afp)) {
+    size = std::max(size,
+                    NumWords_AsyncLet * IGM.getPointerSize());
+  }
+  
   ConstantInitBuilder initBuilder(IGM);
   ConstantStructBuilder builder(
       initBuilder.beginStruct(IGM.AsyncFunctionPointerTy));

lib/IRGen/IRGenModule.h

@@ -893,10 +893,18 @@ class IRGenModule {
   std::string GetObjCSectionName(StringRef Section, StringRef MachOAttributes);
   void SetCStringLiteralSection(llvm::GlobalVariable *GV, ObjCLabelType Type);
 
-  // Mapping of AsyncFunctionPointer records to their corresponding
-  // `@llvm.coro.id.async` intrinsic tag in the function implementation.
-  llvm::DenseMap<llvm::GlobalVariable*, llvm::CallInst*> AsyncCoroIDs;
-
+  void addAsyncCoroIDMapping(llvm::GlobalVariable *asyncFunctionPointer,
+                             llvm::CallInst *coro_id_builtin);
+  
+  llvm::CallInst *getAsyncCoroIDMapping(
+                                    llvm::GlobalVariable *asyncFunctionPointer);
+  
+  void markAsyncFunctionPointerForPadding(
+                                    llvm::GlobalVariable *asyncFunctionPointer);
+  
+  bool isAsyncFunctionPointerMarkedForPadding(
+                                    llvm::GlobalVariable *asyncFunctionPointer);
+  
 private:
   Size PtrSize;
   Size AtomicBoolSize;
@@ -911,7 +919,18 @@ class IRGenModule {
   llvm::PointerType *EnumValueWitnessTablePtrTy = nullptr;
 
   llvm::DenseMap<llvm::Type *, SpareBitVector> SpareBitsForTypes;
-  
+
+  // Mapping of AsyncFunctionPointer records to their corresponding
+  // `@llvm.coro.id.async` intrinsic tag in the function implementation.
+  // This is used for a runtime bug workaround where we need to pad the initial
+  // context size for tasks used as `async let` entry points.
+  //
+  // An entry in the map may have a null value, to indicate that a not-yet-
+  // emitted async function pointer should get the padding applied when it is
+  // emitted.
+  llvm::DenseMap<llvm::GlobalVariable*, llvm::CallInst*> AsyncCoroIDsForPadding;
+
+
 //--- Types -----------------------------------------------------------------
 public:
   const ProtocolInfo &getProtocolInfo(ProtocolDecl *D, ProtocolInfoKind kind);

lib/SIL/IR/SILFunctionType.cpp

@@ -1933,7 +1933,16 @@ static CanSILFunctionType getSILFunctionType(
   }
 
   // Map 'throws' to the appropriate error convention.
-  if (substFnInterfaceType->getExtInfo().isThrowing() && !foreignInfo.error &&
+  // Give the type an error argument whether the substituted type semantically
+  // `throws` or if the abstraction pattern specifies a Swift function type
+  // that also throws. This prevents the need for a second possibly-thunking
+  // conversion when using a non-throwing function in more abstract throwing
+  // context.
+  bool isThrowing = substFnInterfaceType->getExtInfo().isThrowing();
+  if (auto origFnType = origType.getAs<AnyFunctionType>()) {
+    isThrowing |= origFnType->getExtInfo().isThrowing();
+  }
+  if (isThrowing && !foreignInfo.error &&
       !foreignInfo.async) {
     assert(!origType.isForeign()
            && "using native Swift error convention for foreign type!");

lib/SILGen/SILGenApply.cpp

@@ -5948,10 +5948,11 @@ SILGenFunction::emitCoroutineAccessor(SILLocation loc, SILDeclRef accessor,
 ManagedValue SILGenFunction::emitAsyncLetStart(
     SILLocation loc,
     SILValue taskOptions,
-    Type functionType, ManagedValue taskFunction,
+    AbstractClosureExpr *asyncLetEntryPoint,
     SILValue resultBuf) {
   ASTContext &ctx = getASTContext();
-  Type resultType = functionType->castTo<FunctionType>()->getResult();
+  Type resultType = asyncLetEntryPoint->getType()
+    ->castTo<FunctionType>()->getResult();
   Type replacementTypes[] = {resultType};
   auto startBuiltin = cast<FuncDecl>(
       getBuiltinValueDecl(ctx, ctx.getIdentifier("startAsyncLet")));
@@ -5966,8 +5967,15 @@ ManagedValue SILGenFunction::emitAsyncLetStart(
   AbstractionPattern origParam(
       startBuiltin->getGenericSignature().getCanonicalSignature(),
       origParamType);
-  taskFunction = emitSubstToOrigValue(
-      loc, taskFunction, origParam, substParamType);
+  
+  auto conversion = Conversion::getSubstToOrig(origParam, substParamType,
+                                     getLoweredType(origParam, substParamType));
+  ConvertingInitialization convertingInit(conversion, SGFContext());
+  auto taskFunction = emitRValue(asyncLetEntryPoint,
+                                 SGFContext(&convertingInit))
+    .getAsSingleValue(*this, loc);
+  taskFunction = emitSubstToOrigValue(loc, taskFunction,
+                                      origParam, substParamType);
 
   auto apply = B.createBuiltin(
       loc,

lib/SILGen/SILGenDecl.cpp

@@ -1212,8 +1212,7 @@ void SILGenFunction::emitPatternBinding(PatternBindingDecl *PBD,
     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");
+    auto initClosure = cast<AutoClosureExpr>(init);
     bool isThrowing = init->getType()->castTo<AnyFunctionType>()->isThrowing();
 
     // Allocate space to receive the child task's result.
@@ -1240,8 +1239,7 @@ void SILGenFunction::emitPatternBinding(PatternBindingDecl *PBD,
       alet = emitAsyncLetStart(
           loc,
           options.forward(*this), // options is B.createManagedOptionalNone
-          init->getType(),
-          emitRValue(init).getScalarValue(),
+          initClosure,
           resultBufPtr
         ).forward(*this);
     }

lib/SILGen/SILGenFunction.h

@@ -1456,7 +1456,7 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
 
   ManagedValue emitAsyncLetStart(SILLocation loc,
                                  SILValue taskOptions,
-                                 Type functionType, ManagedValue taskFunction,
+                                 AbstractClosureExpr *asyncLetEntryPoint,
                                  SILValue resultBuf);
 
   void emitFinishAsyncLet(SILLocation loc, SILValue asyncLet, SILValue resultBuf);

test/IRGen/async_let_back_deploy_workaround.swift

@@ -1,5 +1,5 @@
-// RUN: %target-swift-frontend -emit-ir -target x86_64-apple-macos99.99 %s | %FileCheck --check-prefix=CHECK --check-prefix=CHECK-sans-workaround %s
-// RUN: %target-swift-frontend -emit-ir -target x86_64-apple-macos12.3 %s | %FileCheck --check-prefix=CHECK --check-prefix=CHECK-with-workaround %s
+// RUN: %target-swift-frontend -emit-ir -target %target-cpu-apple-macos99.99 %s | %FileCheck --check-prefix=CHECK --check-prefix=CHECK-sans-workaround %s
+// RUN: %target-swift-frontend -emit-ir -target %target-cpu-apple-macos12.3 %s | %FileCheck --check-prefix=CHECK --check-prefix=CHECK-with-workaround %s
 
 // REQUIRES: OS=macosx
 
@@ -11,12 +11,26 @@
 // initial context sizes of any async functions used as `async let` entry points
 // to ensure that the preallocated space is never used for the initial context.
 
-// CHECK: [[ASYNC_LET_ENTRY:@"\$sSiIeghHd_Sis5Error_pIegHrzo_TRTATu"]]
+// CHECK: [[ASYNC_LET_ENTRY_FOO:@"\$s32async_let_back_deploy_workaround3foo1x1yS2i_SitYaFSiyYaYbcfu_TATu"]]
 // CHECK-with-workaround-SAME: = internal {{.*}} %swift.async_func_pointer <{ {{.*}}, i32 6{{[0-9][0-9]}} }>
 // CHECK-sans-workaround-SAME: = internal {{.*}} %swift.async_func_pointer <{ {{.*}}, i32 {{[0-9][0-9]}} }>
 
-// CHECK: swift_asyncLet_begin{{.*}}[[ASYNC_LET_ENTRY]]
+// CHECK: [[ASYNC_LET_ENTRY_BAR:@"\$s32async_let_back_deploy_workaround3bar1x1yS2i_SitYaFSiyYaYbcfu_Tu"]]
+// CHECK-with-workaround-SAME: = internal {{.*}} %swift.async_func_pointer <{ {{.*}}, i32 6{{[0-9][0-9]}} }>
+// CHECK-sans-workaround-SAME: = internal {{.*}} %swift.async_func_pointer <{ {{.*}}, i32 {{[0-9][0-9]}} }>
+
+// CHECK-LABEL: define {{.*}}3foo
+// CHECK: swift_asyncLet_begin{{.*}}[[ASYNC_LET_ENTRY_FOO]]
 public func foo(x: Int, y: Int) async -> Int {
     async let z = x + y
     return await z
 }
+
+@_silgen_name("bar_work") func bar_work() -> Int
+
+// CHECK-LABEL: define {{.*}}3bar
+// CHECK: swift_asyncLet_begin{{.*}}[[ASYNC_LET_ENTRY_BAR]]
+public func bar(x: Int, y: Int) async -> Int {
+    async let z = bar_work()
+    return await z
+}