IRGen: Artificially pad async let entry point contexts.

We fixed a bug in the concurrency runtime (rdar://problem/90357994) that would
lead to memory corruption when a new `async let` child task tried to use the
last 16 bytes of the preallocated slab from its parent to seed its own task
allocator. To work around this bug in older OSes that shipped with the bug,
artificially pad the async coroutine context size for any async functions
used as an `async let` entry point, which will prevent the runtime from
going down the path of trying to use the preallocated storage at all.
rdar://90506708

Author: jckarter

lib/IRGen/GenCall.cpp

@@ -3903,14 +3903,17 @@ void irgen::emitAsyncFunctionEntry(IRGenFunction &IGF,
                                    unsigned asyncContextIndex) {
   auto &IGM = IGF.IGM;
   auto size = layout.getSize();
+  auto asyncFuncPointerVar = cast<llvm::GlobalVariable>(IGM.getAddrOfAsyncFunctionPointer(asyncFunction));
   auto asyncFuncPointer = IGF.Builder.CreateBitOrPointerCast(
-      IGM.getAddrOfAsyncFunctionPointer(asyncFunction), IGM.Int8PtrTy);
+                                           asyncFuncPointerVar, IGM.Int8PtrTy);
   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);
   // 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

@@ -26,6 +26,7 @@
 #include "IRGenModule.h"
 #include "LoadableTypeInfo.h"
 #include "ScalarPairTypeInfo.h"
+#include "swift/AST/ASTContext.h"
 #include "swift/AST/ProtocolConformanceRef.h"
 #include "swift/ABI/MetadataValues.h"
 
@@ -192,6 +193,44 @@ llvm::Value *irgen::emitBuiltinStartAsyncLet(IRGenFunction &IGF,
   auto futureResultType = subs.getReplacementTypes()[0]->getCanonicalType();
   auto futureResultTypeMetadata = IGF.emitAbstractTypeMetadataRef(futureResultType);
 
+  // The concurrency runtime for older Apple OSes has a bug in task formation
+  // for `async let`s that may manifest when trying to use room in the
+  // parent task's preallocated `async let` buffer for the child task's
+  // initial task allocator slab. If targeting those older OSes, pad the
+  // context size for async let entry points to never fit in the preallocated
+  // space, so that we don't run into that bug. We leave a note on the
+  // declaration so that coroutine splitting can pad out the final context
+  // size after splitting.
+  auto deploymentAvailability
+    = AvailabilityContext::forDeploymentTarget(IGF.IGM.Context);
+  if (!deploymentAvailability.isContainedIn(
+                                   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,
+                     (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(
+                                   cast<llvm::StructType>(origInit->getType()),
+                                   origInit->getAggregateElement(0u),
+                                   paddedSizeVal);
+    taskAsyncFunctionPointer->setInitializer(newInit);
+  }
+  
   llvm::CallInst *call;
   if (localResultBuffer) {
     // This is @_silgen_name("swift_asyncLet_begin")

lib/IRGen/IRGenModule.h

@@ -328,10 +328,10 @@ class IRGenerator {
 
   /// The queue of IRGenModules for multi-threaded compilation.
   SmallVector<IRGenModule *, 8> Queue;
-
+  
   std::atomic<int> QueueIndex;
 
-  friend class CurrentIGMPtr;  
+  friend class CurrentIGMPtr;
 public:
   explicit IRGenerator(const IRGenOptions &opts, SILModule &module);
 
@@ -893,6 +893,10 @@ 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;
+
 private:
   Size PtrSize;
   Size AtomicBoolSize;

test/IRGen/async_let_back_deploy_workaround.swift

@@ -0,0 +1,22 @@
+// 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
+
+// REQUIRES: OS=macosx
+ 
+// rdar://90506708: Prior to Swift 5.7, the Swift concurrency runtime had a bug
+// that led to memory corruption in cases when an `async let` child task
+// would try to use the last 16 bytes of the preallocated slab from the parent
+// to seed its own task allocator. When targeting older Apple OSes that shipped
+// with this bug in their runtime, we work around the bug by inflating the
+// 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-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]]
+public func foo(x: Int, y: Int) async -> Int {
+    async let z = x + y
+    return await z
+}