[IRGen] Restored non-constant async function calls in PAFs.

Previously, because partial apply forwarders for async functions were
not themselves fully-fledged async functions, they were not able to
handle dynamic functions.  Specifically, the reason was that it was not
possible to produce an async function pointer for the partial apply
forwarder because the size to be used was not knowable.

Thanks to swiftlang#36700, that cause has been
eliminated.  With it, partial apply forwarders are fully-fledged async
functions and in particular have their own async function pointers.
Consequently, it is again possible for these partial apply forwarders to
handle non-constant function pointers.

Here, that behavior is restored, by way of reverting part of
ee63777 while preserving the ABI it
introduced.

rdar://76122027

Author: nate-chandler

include/swift/AST/IRGenOptions.h

@@ -159,6 +159,9 @@ struct PointerAuthOptions : clang::PointerAuthOptions {
   /// values.
   PointerAuthSchema AsyncSwiftDynamicReplacements;
 
+  /// Like PartialApplyCapture but for use with AsyncFunctionPointer values.
+  PointerAuthSchema AsyncPartialApplyCapture;
+
   /// The parent async context stored within a child async context.
   PointerAuthSchema AsyncContextParent;
 

include/swift/IRGen/IRGenSILPasses.h

@@ -19,7 +19,6 @@ 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

@@ -298,8 +298,6 @@ 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,7 +52,6 @@ add_swift_host_library(swiftIRGen STATIC
   MetadataLayout.cpp
   MetadataRequest.cpp
   Outlining.cpp
-  PartialApplyLowering.cpp
   StructLayout.cpp
   SwiftTargetInfo.cpp
   TypeLayout.cpp

lib/IRGen/Explosion.h

@@ -97,6 +97,10 @@ class Explosion {
     Values.append(values.begin(), values.end());
   }
 
+  void insert(unsigned index, llvm::Value *value) {
+    Values.insert(Values.begin() + index, value);
+  }
+
   /// Return an array containing the given range of values.  The values
   /// are not claimed.
   ArrayRef<llvm::Value*> getRange(unsigned from, unsigned to) const {

lib/IRGen/GenFunc.cpp

@@ -768,15 +768,12 @@ class PartialApplicationForwarderEmission {
         origParams(subIGF.collectParameters()) {}
 
 public:
-  enum class DynamicFunctionKind {
-    Witness,
-    PartialApply,
-  };
   virtual void begin(){};
 
   virtual void gatherArgumentsFromApply() = 0;
 
-  virtual void mapAsyncParameters() {}
+  virtual void mapAsyncParameters(FunctionPointer fnPtr) {}
+  virtual void recordAsyncParametersInsertionPoint(){};
 
   void gatherArgumentsFromApply(bool isAsync) {
     // Lower the forwarded arguments in the original function's generic context.
@@ -831,7 +828,7 @@ class PartialApplicationForwarderEmission {
     }
 
     if (isAsync)
-      mapAsyncParameters();
+      recordAsyncParametersInsertionPoint();
 
     // Reemit the parameters as unsubstituted.
     for (unsigned i = 0; i < outType->getParameters().size(); ++i) {
@@ -921,12 +918,10 @@ class PartialApplicationForwarderEmission {
 
   llvm::Value *getContext() { return origParams.claimNext(); }
 
-  virtual llvm::Value *getDynamicFunctionPointer(PointerAuthInfo &authInfo) = 0;
+  virtual llvm::Value *getDynamicFunctionPointer() = 0;
   virtual llvm::Value *getDynamicFunctionContext() = 0;
-  virtual void addDynamicFunctionContext(Explosion &explosion,
-                                         DynamicFunctionKind kind) = 0;
-  virtual void addDynamicFunctionPointer(Explosion &explosion,
-                                         DynamicFunctionKind kind) = 0;
+  virtual void addDynamicFunctionContext(Explosion &explosion) = 0;
+  virtual void addDynamicFunctionPointer(Explosion &explosion) = 0;
 
   void addSelf(Explosion &explosion) { addArgument(explosion); }
   void addWitnessSelfMetadata(llvm::Value *value) {
@@ -965,16 +960,12 @@ class SyncPartialApplicationForwarderEmission
   void gatherArgumentsFromApply() override {
     super::gatherArgumentsFromApply(false);
   }
-  llvm::Value *getDynamicFunctionPointer(PointerAuthInfo &authInfo) override {
-    return args.takeLast();
-  }
+  llvm::Value *getDynamicFunctionPointer() override { return args.takeLast(); }
   llvm::Value *getDynamicFunctionContext() override { return args.takeLast(); }
-  void addDynamicFunctionContext(Explosion &explosion,
-                                 DynamicFunctionKind kind) override {
+  void addDynamicFunctionContext(Explosion &explosion) override {
     addArgument(explosion);
   }
-  void addDynamicFunctionPointer(Explosion &explosion,
-                                 DynamicFunctionKind kind) override {
+  void addDynamicFunctionPointer(Explosion &explosion) override {
     addArgument(explosion);
   }
   void forwardErrorResult() override {
@@ -1045,13 +1036,6 @@ class AsyncPartialApplicationForwarderEmission
   Address context;
   Address calleeContextBuffer;
   unsigned currentArgumentIndex;
-  struct DynamicFunction {
-    using Kind = DynamicFunctionKind;
-    Kind kind;
-    llvm::Value *pointer;
-    llvm::Value *context;
-  };
-  Optional<DynamicFunction> dynamicFunction = llvm::None;
   struct Self {
     enum class Kind {
       Method,
@@ -1061,6 +1045,7 @@ class AsyncPartialApplicationForwarderEmission
     llvm::Value *value;
   };
   Optional<Self> self = llvm::None;
+  unsigned asyncParametersInsertionIndex = 0;
 
   void saveValue(ElementLayout layout, Explosion &explosion) {
     Address addr = layout.project(subIGF, context, /*offsets*/ llvm::None);
@@ -1088,24 +1073,28 @@ class AsyncPartialApplicationForwarderEmission
 
   void begin() override { super::begin(); }
 
-  void mapAsyncParameters() override {
+  void recordAsyncParametersInsertionPoint() override {
     // Ignore the original context.
     (void)origParams.claimNext();
 
+    asyncParametersInsertionIndex = args.size();
+  }
+  void mapAsyncParameters(FunctionPointer fnPtr) override {
     llvm::Value *dynamicContextSize32;
     auto initialContextSize = Size(0);
     std::tie(calleeFunction, dynamicContextSize32) = getAsyncFunctionAndSize(
-        subIGF, origType->getRepresentation(), *staticFnPtr,
-        nullptr, std::make_pair(true, true), initialContextSize);
+        subIGF, origType->getRepresentation(), fnPtr, nullptr,
+        std::make_pair(true, true), initialContextSize);
     auto *dynamicContextSize =
         subIGF.Builder.CreateZExt(dynamicContextSize32, subIGF.IGM.SizeTy);
     calleeContextBuffer =
         emitAllocAsyncContext(subIGF, dynamicContextSize);
     context = layout.emitCastTo(subIGF, calleeContextBuffer.getAddress());
     auto calleeContext =
         layout.emitCastTo(subIGF, calleeContextBuffer.getAddress());
-    args.add(subIGF.Builder.CreateBitOrPointerCast(
-        calleeContextBuffer.getAddress(), IGM.SwiftContextPtrTy));
+    args.insert(asyncParametersInsertionIndex,
+                subIGF.Builder.CreateBitOrPointerCast(
+                    calleeContextBuffer.getAddress(), IGM.SwiftContextPtrTy));
 
     // Set caller info into the context.
     { // caller context
@@ -1146,20 +1135,14 @@ class AsyncPartialApplicationForwarderEmission
   void gatherArgumentsFromApply() override {
     super::gatherArgumentsFromApply(true);
   }
-  llvm::Value *getDynamicFunctionPointer(PointerAuthInfo &authInfo) override {
-    llvm_unreachable(
-        "async partial applies never have dynamic function pointers");
-  }
+  llvm::Value *getDynamicFunctionPointer() override { return args.takeLast(); }
   llvm::Value *getDynamicFunctionContext() override {
     return args.takeLast();
   }
-  void addDynamicFunctionContext(Explosion &explosion,
-                                 DynamicFunction::Kind kind) override {
+  void addDynamicFunctionContext(Explosion &explosion) override {
     addArgument(explosion);
   }
-  void addDynamicFunctionPointer(Explosion &explosion,
-                                 DynamicFunction::Kind kind) override {
-    assert(false);
+  void addDynamicFunctionPointer(Explosion &explosion) override {
     addArgument(explosion);
   }
 
@@ -1669,20 +1652,10 @@ static llvm::Value *emitPartialApplicationForwarder(IRGenModule &IGM,
       } else {
         switch (extraFieldIndex) {
         case 0:
-          emission->addDynamicFunctionContext(
-              param, isWitnessMethodCallee
-                         ? PartialApplicationForwarderEmission::
-                               DynamicFunctionKind::Witness
-                         : PartialApplicationForwarderEmission::
-                               DynamicFunctionKind::PartialApply);
+          emission->addDynamicFunctionContext(param);
           break;
         case 1:
-          emission->addDynamicFunctionPointer(
-              param, isWitnessMethodCallee
-                         ? PartialApplicationForwarderEmission::
-                               DynamicFunctionKind::Witness
-                         : PartialApplicationForwarderEmission::
-                               DynamicFunctionKind::PartialApply);
+          emission->addDynamicFunctionPointer(param);
           break;
         default:
           llvm_unreachable("unexpected extra field in thick context");
@@ -1724,22 +1697,29 @@ static llvm::Value *emitPartialApplicationForwarder(IRGenModule &IGM,
     // Otherwise, it was the last thing we added to the layout.
 
     assert(lastCapturedFieldPtr);
-    auto authInfo = PointerAuthInfo::emit(subIGF,
-                            IGM.getOptions().PointerAuth.PartialApplyCapture,
-                            lastCapturedFieldPtr,
-                            PointerAuthEntity::Special::PartialApplyCapture);
+    auto authInfo = PointerAuthInfo::emit(
+        subIGF,
+        origType->isAsync()
+            ? IGM.getOptions().PointerAuth.AsyncPartialApplyCapture
+            : IGM.getOptions().PointerAuth.PartialApplyCapture,
+        lastCapturedFieldPtr, PointerAuthEntity::Special::PartialApplyCapture);
 
     // The dynamic function pointer is packed "last" into the context,
     // and we pulled it out as an argument.  Just pop it off.
-    auto fnPtr = emission->getDynamicFunctionPointer(authInfo);
+    auto fnPtr = emission->getDynamicFunctionPointer();
 
     // It comes out of the context as an i8*. Cast to the function type.
     fnPtr = subIGF.Builder.CreateBitCast(fnPtr, fnTy);
 
-    return FunctionPointer(FunctionPointer::Kind::Function, fnPtr, authInfo,
-                           origSig);
+    return FunctionPointer(origType->isAsync()
+                               ? FunctionPointer::Kind::AsyncFunctionPointer
+                               : FunctionPointer::Kind::Function,
+                           fnPtr, authInfo, origSig);
   }();
 
+  if (origType->isAsync())
+    emission->mapAsyncParameters(fnPtr);
+
   // Derive the context argument if needed.  This is either:
   //   - the saved context argument, in which case it was the last
   //     thing we added to the layout other than a possible non-static
@@ -1770,12 +1750,7 @@ static llvm::Value *emitPartialApplicationForwarder(IRGenModule &IGM,
     if (isMethodCallee) {
       emission->addSelf(explosion);
     } else {
-      emission->addDynamicFunctionContext(
-          explosion, isWitnessMethodCallee
-                         ? PartialApplicationForwarderEmission::
-                               DynamicFunctionKind::Witness
-                         : PartialApplicationForwarderEmission::
-                               DynamicFunctionKind::PartialApply);
+      emission->addDynamicFunctionContext(explosion);
     }
 
   // Pass a placeholder for thin function calls.
@@ -2120,16 +2095,18 @@ Optional<StackAddress> irgen::emitFunctionPartialApplication(
       // We don't add non-constant function pointers to the explosion above,
       // so we need to handle them specially now.
       if (i == nonStaticFnIndex) {
-        llvm::Value *fnPtr;
-        if (auto &schema = IGF.getOptions().PointerAuth.PartialApplyCapture) {
-          auto schemaAuthInfo =
-            PointerAuthInfo::emit(IGF, schema, fieldAddr.getAddress(),
-                           PointerAuthEntity::Special::PartialApplyCapture);
+        llvm::Value *fnPtr = fn.getRawPointer();
+        if (auto &schema =
+                origType->isAsync()
+                    ? IGF.getOptions().PointerAuth.AsyncPartialApplyCapture
+                    : IGF.getOptions().PointerAuth.PartialApplyCapture) {
+          auto schemaAuthInfo = PointerAuthInfo::emit(
+              IGF, schema, fieldAddr.getAddress(),
+              PointerAuthEntity::Special::PartialApplyCapture);
           fnPtr =
               emitPointerAuthResign(IGF, fn, schemaAuthInfo).getRawPointer();
-        } else {
-          fnPtr = fn.getRawPointer();
         }
+
         fnPtr = IGF.Builder.CreateBitCast(fnPtr, IGF.IGM.Int8PtrTy);
         IGF.Builder.CreateStore(fnPtr, fieldAddr);
         continue;

lib/IRGen/IRGen.cpp

@@ -644,6 +644,11 @@ static void setPointerAuthOptions(PointerAuthOptions &opts,
   // might end up on a global constant initializer.
   auto nonABICodeKey = PointerAuthSchema::ARM8_3Key::ASIB;
 
+  // A key suitable for data pointers that are only used in private
+  // situations.  Do not use this key for any sort of signature that
+  // might end up on a global constant initializer.
+  auto nonABIDataKey = PointerAuthSchema::ARM8_3Key::ASDB;
+
   // If you change anything here, be sure to update <ptrauth.h>.
   opts.SwiftFunctionPointers =
     PointerAuthSchema(codeKey, /*address*/ false, Discrimination::Type);
@@ -712,6 +717,9 @@ static void setPointerAuthOptions(PointerAuthOptions &opts,
   opts.AsyncSwiftDynamicReplacements =
       PointerAuthSchema(dataKey, /*address*/ true, Discrimination::Decl);
 
+  opts.AsyncPartialApplyCapture =
+      PointerAuthSchema(nonABIDataKey, /*address*/ true, Discrimination::Decl);
+
   opts.AsyncContextParent =
       PointerAuthSchema(dataKey, /*address*/ true, Discrimination::Constant,
                         SpecialPointerAuthDiscriminators::AsyncContextParent);