[IRGen] Fix ABI for thick async functions.

Previously, thick async functions were represented sometimes as a pair
of (AsyncFunctionPointer, nullptr)--when the thick function was produced
via a thin_to_thick_function, e.g.--and sometimes as a pair of
(FunctionPointer, ThickContext)--when the thick function was produced by
a partial_apply--with the size stored in the slot of the ThickContext.

That optimized for the wrong case: partial applies of dynamic async
functions; in that case, there is no appropriate AsyncFunctionPointer to
form when lowering the partial_apply instruction.  The far more common
case is to know exactly which function is being partially applied.  In
that case, we can form the appropriate AsyncFunctionPointer.

Furthermore, the previous representation made calling a thick function
more complex: it was always necessary to check whether the context was
in fact null and then proceed along two different paths depending.

Here, that behavior is corrected by creating a thunk in a mandatory
IRGen SIL pass in the case that the function that is being partially
applied is dynamic.  That new thunk is then partially applied in place
of the original partial_apply of the dynamic function.

Author: nate-chandler

docs/ABI/Mangling.rst

@@ -210,6 +210,7 @@ types where the metadata itself has unknown layout.)
   global ::= global 'Td'                 // direct method reference thunk
   global ::= global 'TI'                 // implementation of a dynamic_replaceable function
   global ::= global 'Tu'                 // async function pointer of a function
+  global ::= global 'Tw'                 // async partial apply thunk for a non-constant function
   global ::= global 'TX'                 // function pointer of a dynamic_replaceable function
   global ::= entity entity 'TV'          // vtable override thunk, derived followed by base
   global ::= type label-list? 'D'        // type mangling for the debugger with label list for function types.

include/swift/IRGen/IRGenSILPasses.h

@@ -19,6 +19,7 @@ namespace irgen {
 /// Create a pass to hoist alloc_stack instructions with non-fixed size.
 SILTransform *createAllocStackHoisting();
 SILTransform *createLoadableByAddress();
+SILTransform *createPartialApplyLowering();
 
 } // end namespace irgen
 } // end namespace swift

include/swift/SILOptimizer/PassManager/Passes.def

@@ -296,6 +296,8 @@ IRGEN_PASS(LoadableByAddress, "loadable-address",
      "SIL Large Loadable type by-address lowering.")
 PASS(MandatorySILLinker, "mandatory-linker",
      "Deserialize all referenced SIL functions that are shared or transparent")
+IRGEN_PASS(PartialApplyLowering, "partial-apply-lowering",
+           "Partial Apply Lowering")
 PASS(PerformanceSILLinker, "performance-linker",
      "Deserialize all referenced SIL functions")
 PASS(RawSILInstLowering, "raw-sil-inst-lowering",

lib/IRGen/CMakeLists.txt

@@ -52,6 +52,7 @@ add_swift_host_library(swiftIRGen STATIC
   MetadataLayout.cpp
   MetadataRequest.cpp
   Outlining.cpp
+  PartialApplyLowering.cpp
   StructLayout.cpp
   SwiftTargetInfo.cpp
   TypeLayout.cpp

lib/IRGen/GenCall.cpp

@@ -1926,208 +1926,38 @@ std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
   assert(functionPointer.getKind() != FunctionPointer::Kind::Function);
   bool emitFunction = values.first;
   bool emitSize = values.second;
-  // TODO: This calculation should be extracted out into standalone functions
-  //       emitted on-demand per-module to improve codesize.
-  switch (representation) {
-  case SILFunctionTypeRepresentation::Thick: {
-    assert(!functionPointer.useStaticContextSize());
-    // If the called function is thick, the size of the called function's
-    // async context is not statically knowable.
-    //
-    // Specifically, if the thick function was produced by a partial_apply,
-    // the function which was originally partially applied determines the
-    // size of the needed async context.  That original function isn't known
-    // statically.  The dynamic size is available within the context as an
-    // i32 at the first index: <{ %swift.refcounted*, /*size*/ i32, ... }>.
-    // In this case, the function pointer is actually a pointer to an llvm
-    // function.
-    //
-    // On the other hand, if the thick function was produced by a
-    // thin_to_thick_function, then the context will be nullptr.  In that
-    // case, the dynamic size of the needed async context is available within
-    // the struct, an AsyncFunctionPointer pointed to by the "function" pointer
-    // as an i32 at the second index: <{ /*fn rel addr*/ i32, /*size*/ i32 }>.
-    //
-    // We are currently emitting into some basic block.  To handle these two
-    // cases, we need to branch based on whether the context is nullptr; each
-    // branch must then determine the size and function pointer in the manner
-    // appropriate to it.  Finally, both blocks must join back together to make
-    // the call:
-    //
-    //                      +-------------------------+
-    //                      |%cond = %ctx == nullptr  |
-    //     +----------------|br %cond, thin, thick    |----------------------+
-    //     |                +-------------------------+                      |
-    //     |                                                                 |
-    //     V                                                                 |
-    // +-thin-------------------------------------------+                    |
-    // |%afp = bitcast %fp to %swift.async_func_pointer*|                    |
-    // |%size_ptr = getelementptr %afp, i32 0, i32 1    |                    |
-    // |%size = load %size_ptr                          |                    |
-    // |%offset_ptr = getelementptr %afp, i32 0, i32 1  |                    |
-    // |%offset = load i32 %offset_ptr                  |                    |
-    // |%offset64 = sext %offset to i64                 |                    |
-    // |%raw_fp = add %offset64, %offset_ptr            |                    |
-    // |br join(%raw_fp, %size)                         |                    |
-    // +------------------------------------------------+                    |
-    //     |                                                                 |
-    //     |                                                                 V
-    //     |                +-thick--------------------------------------------+
-    //     |                |%layout = bitcast %ctx to <{%swift.context*, i32}>|
-    //     |                |%size_addr = getelementptr %layout, i32 0, i32 1  |
-    //     |                |%size = load %size_addr                           |
-    //     |                |br join(%fp, %size)                               |
-    //     |                +---/----------------------------------------------+
-    //     |                   /
-    //     |                  /
-    //     V                 V
-    // +-join(%fn, %size)------------------------------------------------------+
-    // |%dataAddr = swift_taskAlloc(%task, %size)                              |
-    // |%async_context = bitcast %dataAddr to ASYNC_CONTEXT(static_callee_type)|
-    // |... // populate the fields %ctx with arguments                         |
-    // |call %fn(%async_context, %ctx)                                         |
-    // +-----------------------------------------------------------------------+
-    auto *thinBlock = llvm::BasicBlock::Create(IGF.IGM.getLLVMContext());
-    auto *thickBlock = llvm::BasicBlock::Create(IGF.IGM.getLLVMContext());
-    auto *joinBlock = llvm::BasicBlock::Create(IGF.IGM.getLLVMContext());
-
-    auto hasThickContext =
-        IGF.Builder.CreateICmpNE(thickContext, IGF.IGM.RefCountedNull);
-    IGF.Builder.CreateCondBr(hasThickContext, thickBlock, thinBlock);
-
-    SmallVector<std::pair<llvm::BasicBlock *, llvm::Value *>, 2> fnPhiValues;
-    SmallVector<std::pair<llvm::BasicBlock *, llvm::Value *>, 2> sizePhiValues;
-    { // thin
-      IGF.Builder.emitBlock(thinBlock);
-      auto *ptr = functionPointer.getRawPointer();
-      if (auto authInfo = functionPointer.getAuthInfo()) {
-        ptr = emitPointerAuthAuth(IGF, ptr, authInfo);
-      }
-      auto *afpPtr =
-          IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
-      if (emitFunction) {
-        llvm::Value *addrPtr = IGF.Builder.CreateStructGEP(afpPtr, 0);
-        auto *uncastFnPtr = IGF.emitLoadOfRelativePointer(
-            Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
-            /*expectedType*/ functionPointer.getFunctionType()->getPointerTo());
-        auto *fnPtr = IGF.Builder.CreateBitCast(uncastFnPtr, IGF.IGM.Int8PtrTy);
-        if (auto authInfo = functionPointer.getAuthInfo()) {
-          fnPtr = emitPointerAuthSign(IGF, fnPtr, authInfo);
-        }
-        fnPhiValues.push_back({thinBlock, fnPtr});
-      }
-      if (emitSize) {
-        auto *sizePtr = IGF.Builder.CreateStructGEP(afpPtr, 1);
-        auto *size =
-            IGF.Builder.CreateLoad(sizePtr, IGF.IGM.getPointerAlignment());
-        sizePhiValues.push_back({thinBlock, size});
-      }
-      IGF.Builder.CreateBr(joinBlock);
-    }
-
-    { // thick
-      IGF.Builder.emitBlock(thickBlock);
-      if (emitFunction) {
-        auto *uncastFnPtr = functionPointer.getRawPointer();
-        auto *fnPtr = IGF.Builder.CreateBitCast(uncastFnPtr, IGF.IGM.Int8PtrTy);
-        fnPhiValues.push_back({thickBlock, fnPtr});
-      }
-      if (emitSize) {
-        SmallVector<const TypeInfo *, 4> argTypeInfos;
-        SmallVector<SILType, 4> argValTypes;
-        auto int32ASTType =
-            BuiltinIntegerType::get(32, IGF.IGM.IRGen.SIL.getASTContext())
-                ->getCanonicalType();
-        auto int32SILType = SILType::getPrimitiveObjectType(int32ASTType);
-        const TypeInfo &int32TI = IGF.IGM.getTypeInfo(int32SILType);
-        argValTypes.push_back(int32SILType);
-        argTypeInfos.push_back(&int32TI);
-        HeapLayout layout(IGF.IGM, LayoutStrategy::Optimal, argValTypes,
-                          argTypeInfos,
-                          /*typeToFill*/ nullptr, NecessaryBindings());
-        auto castThickContext =
-            layout.emitCastTo(IGF, thickContext, "context.prefix");
-        auto sizeLayout = layout.getElement(0);
-        auto sizeAddr = sizeLayout.project(IGF, castThickContext,
-                                           /*NonFixedOffsets*/ llvm::None);
-        auto *sizeValue = IGF.Builder.CreateLoad(sizeAddr);
-        sizePhiValues.push_back({thickBlock, sizeValue});
-      }
-      IGF.Builder.CreateBr(joinBlock);
-    }
-
-    { // join
-      IGF.Builder.emitBlock(joinBlock);
-      llvm::Value *fn = nullptr;
-      llvm::PHINode *fnPhi = nullptr;
-      llvm::PHINode *sizePhi = nullptr;
-      if (emitFunction) {
-        fnPhi = IGF.Builder.CreatePHI(IGF.IGM.Int8PtrTy, fnPhiValues.size());
-      }
-      if (emitSize) {
-        sizePhi = IGF.Builder.CreatePHI(IGF.IGM.Int32Ty, sizePhiValues.size());
-      }
-      if (emitFunction) {
-        assert(fnPhi);
-        for (auto &entry : fnPhiValues) {
-          fnPhi->addIncoming(entry.second, entry.first);
-        }
-        fn = IGF.Builder.CreateBitCast(
-            fnPhi, functionPointer.getFunctionType()->getPointerTo());
-      }
-      llvm::Value *size = nullptr;
-      if (emitSize) {
-        assert(sizePhi);
-        for (auto &entry : sizePhiValues) {
-          sizePhi->addIncoming(entry.second, entry.first);
-        }
-        size = sizePhi;
-      }
-      return {fn, size};
+  auto *ptr = functionPointer.getRawPointer();
+  if (auto authInfo = functionPointer.getAuthInfo()) {
+    ptr = emitPointerAuthAuth(IGF, ptr, authInfo);
+  }
+  auto *afpPtr =
+      IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
+  llvm::Value *fn = nullptr;
+  if (emitFunction) {
+    if (functionPointer.useStaticContextSize()) {
+      fn = functionPointer.getRawPointer();
+    } else {
+      llvm::Value *addrPtr = IGF.Builder.CreateStructGEP(afpPtr, 0);
+      fn = IGF.emitLoadOfRelativePointer(
+          Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
+          /*expectedType*/ functionPointer.getFunctionType()->getPointerTo());
     }
-  }
-  case SILFunctionTypeRepresentation::Thin:
-  case SILFunctionTypeRepresentation::CFunctionPointer:
-  case SILFunctionTypeRepresentation::Method:
-  case SILFunctionTypeRepresentation::ObjCMethod:
-  case SILFunctionTypeRepresentation::WitnessMethod:
-  case SILFunctionTypeRepresentation::Closure:
-  case SILFunctionTypeRepresentation::Block: {
-    auto *ptr = functionPointer.getRawPointer();
     if (auto authInfo = functionPointer.getAuthInfo()) {
-      ptr = emitPointerAuthAuth(IGF, ptr, authInfo);
+      fn = emitPointerAuthSign(IGF, fn, authInfo);
     }
-    auto *afpPtr =
-        IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
-    llvm::Value *fn = nullptr;
-    if (emitFunction) {
-      if (functionPointer.useStaticContextSize()) {
-        fn = functionPointer.getRawPointer();
-      } else {
-        llvm::Value *addrPtr = IGF.Builder.CreateStructGEP(afpPtr, 0);
-        fn = IGF.emitLoadOfRelativePointer(
-            Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
-            /*expectedType*/ functionPointer.getFunctionType()->getPointerTo());
-      }
-      if (auto authInfo = functionPointer.getAuthInfo()) {
-        fn = emitPointerAuthSign(IGF, fn, authInfo);
-      }
-    }
-    llvm::Value *size = nullptr;
-    if (emitSize) {
-      if (functionPointer.useStaticContextSize()) {
-        size = llvm::ConstantInt::get(IGF.IGM.Int32Ty,
-                                      initialContextSize.getValue());
-      } else {
-        assert(!functionPointer.useStaticContextSize());
-        auto *sizePtr = IGF.Builder.CreateStructGEP(afpPtr, 1);
-        size = IGF.Builder.CreateLoad(sizePtr, IGF.IGM.getPointerAlignment());
-      }
-    }
-    return {fn, size};
   }
+  llvm::Value *size = nullptr;
+  if (emitSize) {
+    if (functionPointer.useStaticContextSize()) {
+      size = llvm::ConstantInt::get(IGF.IGM.Int32Ty,
+                                    initialContextSize.getValue());
+    } else {
+      assert(!functionPointer.useStaticContextSize());
+      auto *sizePtr = IGF.Builder.CreateStructGEP(afpPtr, 1);
+      size = IGF.Builder.CreateLoad(sizePtr, IGF.IGM.getPointerAlignment());
+    }
   }
-  llvm_unreachable("unhandled case");
+  return {fn, size};
 }
 
 static void externalizeArguments(IRGenFunction &IGF, const Callee &callee,