IRGen: Properly adjust the closure type of a partial_apply of an objc_method

It needs to match with the (large loadable) lowered closure type in the rest of
the program: Large types in the signature need to be passed indirectly.

rdar://127367321

lib/IRGen/IRGenSIL.cpp

@@ -6674,6 +6674,20 @@ void IRGenSILFunction::visitEndUnpairedAccessInst(EndUnpairedAccessInst *i) {
 }
 
 void IRGenSILFunction::visitConvertFunctionInst(swift::ConvertFunctionInst *i) {
+  auto &lv = getLoweredValue(i->getOperand());
+  if (lv.kind == LoweredValue::Kind::ObjCMethod) {
+    // LoadableByAddress lowering will insert convert_function instructions to
+    // change the type of a partial_apply instruction involving a objc_method
+    // convention, to change the partial_apply's SIL type (rewriting large types
+    // to @in_guaranteed/@out). This is important for pointer authentication.
+
+    // The convert_function instruction will carry the desired SIL type.
+    // Here we just forward the objective-c method.
+    auto &objcMethod = lv.getObjCMethod();
+    setLoweredObjCMethod(i, objcMethod.getMethod());
+    return;
+  }
+
   // This instruction is specified to be a no-op.
   Explosion temp = getLoweredExplosion(i->getOperand());
 

lib/IRGen/LoadableByAddress.cpp

@@ -81,10 +81,12 @@ public:
                                             irgen::IRGenModule &IGM);
   SmallVector<SILResultInfo, 2> getNewResults(GenericEnvironment *GenericEnv,
                                               CanSILFunctionType fnType,
-                                              irgen::IRGenModule &Mod);
+                                              irgen::IRGenModule &Mod,
+                                              bool mustTransform = false);
   CanSILFunctionType getNewSILFunctionType(GenericEnvironment *env,
                                            CanSILFunctionType fnType,
-                                           irgen::IRGenModule &IGM);
+                                           irgen::IRGenModule &IGM,
+                                           bool mustTransform = false);
   SILType getNewOptionalFunctionType(GenericEnvironment *GenericEnv,
                                      SILType storageType,
                                      irgen::IRGenModule &Mod);
@@ -241,8 +243,9 @@ bool LargeSILTypeMapper::newResultsDiffer(GenericEnvironment *GenericEnv,
 
 static bool modNonFuncTypeResultType(GenericEnvironment *genEnv,
                                      CanSILFunctionType loweredTy,
-                                     irgen::IRGenModule &Mod) {
+                                     irgen::IRGenModule &Mod,
-  if (!modifiableFunction(loweredTy)) {
+                                     bool mustTransform = false) {
+  if (!modifiableFunction(loweredTy) && !mustTransform) {
     return false;
   }
   if (loweredTy->getNumResults() != 1) {
@@ -259,7 +262,8 @@ static bool modNonFuncTypeResultType(GenericEnvironment *genEnv,
 SmallVector<SILResultInfo, 2>
 LargeSILTypeMapper::getNewResults(GenericEnvironment *GenericEnv,
                                   CanSILFunctionType fnType,
-                                  irgen::IRGenModule &Mod) {
+                                  irgen::IRGenModule &Mod,
+                                  bool mustTransform) {
   // Get new SIL Function results - same as old results UNLESS:
   // 1) Function type results might have a different signature
   // 2) Large loadables are replaced by @out version
@@ -268,7 +272,7 @@ LargeSILTypeMapper::getNewResults(GenericEnvironment *GenericEnv,
   for (auto result : origResults) {
     SILType currResultTy = result.getSILStorageInterfaceType();
     SILType newSILType = getNewSILType(GenericEnv, currResultTy, Mod);
-    if (modNonFuncTypeResultType(GenericEnv, fnType, Mod)) {
+    if (modNonFuncTypeResultType(GenericEnv, fnType, Mod, mustTransform)) {
       // Case (2) Above
       SILResultInfo newSILResultInfo(newSILType.getASTType(),
                                      ResultConvention::Indirect);
@@ -288,8 +292,9 @@ LargeSILTypeMapper::getNewResults(GenericEnvironment *GenericEnv,
 CanSILFunctionType
 LargeSILTypeMapper::getNewSILFunctionType(GenericEnvironment *env,
                                           CanSILFunctionType fnType,
-                                          irgen::IRGenModule &IGM) {
+                                          irgen::IRGenModule &IGM,
-  if (!modifiableFunction(fnType)) {
+                                          bool mustTransform) {
+  if (!modifiableFunction(fnType) && !mustTransform) {
     return fnType;
   }
   
@@ -301,7 +306,7 @@ LargeSILTypeMapper::getNewSILFunctionType(GenericEnvironment *env,
   
   auto newParams = getNewParameters(env, fnType, IGM);
   auto newYields = getNewYields(env, fnType, IGM);
-  auto newResults = getNewResults(env, fnType, IGM);
+  auto newResults = getNewResults(env, fnType, IGM, mustTransform);
   auto newFnType = SILFunctionType::get(
       fnType->getInvocationGenericSignature(),
       fnType->getExtInfo(),
@@ -2623,7 +2628,20 @@ void LoadableByAddress::recreateSingleApply(
     // Change the type of the Closure
     auto partialApplyConvention = castedApply->getCalleeConvention();
     auto resultIsolation = castedApply->getResultIsolation();
-
+    // We do need to update the closure's funtion type to match with the other
+    // uses inside of the binary. Pointer auth cares about the SIL function
+    // type.
+    if (callee->getType().castTo<SILFunctionType>()->getExtInfo().getRepresentation() ==
+        SILFunctionTypeRepresentation::ObjCMethod) {
+      CanSILFunctionType newFnType =
+        MapperCache.getNewSILFunctionType(
+          genEnv,
+          callee->getType().castTo<SILFunctionType>(), *currIRMod,
+          /*mustTransform*/ true);
+      SILType newType = SILType::getPrimitiveObjectType(newFnType);
+      callee = applyBuilder.createConvertFunction(castedApply->getLoc(),
+                                                  callee, newType, false);
+    }
     auto newApply = applyBuilder.createPartialApply(
         castedApply->getLoc(), callee, applySite.getSubstitutionMap(), callArgs,
         partialApplyConvention, resultIsolation, castedApply->isOnStack());

test/IRGen/Inputs/large_c.h

@@ -30,6 +30,7 @@ struct SamplesType {
     void* AA;
 };
 
+struct SamplesType samples();
 
 typedef struct _ContainedType {
   unsigned int f1;

test/IRGen/big_types_corner_cases.sil

@@ -1,9 +1,10 @@
-// RUN: %target-swift-frontend -Xllvm -sil-disable-pass=simplification -I %S/Inputs/abi %s -emit-ir | %FileCheck %s
+// RUN: %target-swift-frontend -sil-verify-none -Xllvm -sil-disable-pass=simplification -I %S/Inputs/abi %s -emit-ir | %FileCheck %s
 
 // REQUIRES: CPU=x86_64
 // REQUIRES: OS=macosx
 
-sil_stage canonical
+sil_stage lowered
+
 import c_layout
 import Builtin
 import Swift
@@ -307,13 +308,17 @@ bb0(%0 : $AnyObject):
   dynamic_method_br %2 : $@opened("E5D03528-36AD-11E8-A0AB-D0817AD47398", AnyObject) Self, #X.foo!foreign, bb1, bb2
 
 bb1(%6 : $@convention(objc_method) (@opened("E5D03528-36AD-11E8-A0AB-D0817AD47398", AnyObject) Self) -> BitfieldOne): // Preds: bb0
+  %addr = alloc_stack $BitfieldOne
   strong_retain %2 : $@opened("E5D03528-36AD-11E8-A0AB-D0817AD47398", AnyObject) Self
-  %8 = partial_apply [callee_guaranteed] %6(%2) : $@convention(objc_method) (@opened("E5D03528-36AD-11E8-A0AB-D0817AD47398", AnyObject) Self) -> BitfieldOne
+  %cvt = convert_function %6 : $@convention(objc_method) (@opened("E5D03528-36AD-11E8-A0AB-D0817AD47398", AnyObject) Self) -> BitfieldOne to $@convention(objc_method) (@opened("E5D03528-36AD-11E8-A0AB-D0817AD47398", AnyObject) Self) -> @out BitfieldOne
-  %9 = apply %8() : $@callee_guaranteed () -> BitfieldOne
+  %8 = partial_apply [callee_guaranteed] %cvt(%2) : $@convention(objc_method) (@opened("E5D03528-36AD-11E8-A0AB-D0817AD47398", AnyObject) Self) -> @out BitfieldOne
+  %9 = apply %8(%addr) : $@callee_guaranteed () -> @out BitfieldOne
   %10 = init_enum_data_addr %4 : $*Optional<BitfieldOne>, #Optional.some!enumelt
-  store %9 to %10 : $*BitfieldOne
+  %val = load %addr : $*BitfieldOne
+  store %val  to %10 : $*BitfieldOne
   inject_enum_addr %4 : $*Optional<BitfieldOne>, #Optional.some!enumelt
-  strong_release %8 : $@callee_guaranteed () -> BitfieldOne
+  strong_release %8 : $@callee_guaranteed () -> @out BitfieldOne
+  dealloc_stack %addr : $*BitfieldOne
   br bb3
 
 bb2:                                              // Preds: bb0

test/IRGen/loadable_by_address_objc_method.swift

@@ -0,0 +1,29 @@
+// RUN: %target-swift-frontend -I %t -emit-ir %s -import-objc-header %S/Inputs/large_c.h | %FileCheck %s
+// RUN: %target-swift-frontend -I %t -emit-ir %s -import-objc-header %S/Inputs/large_c.h -Xllvm -sil-print-after=loadable-address 2>&1 | %FileCheck %s --check-prefix=SIL
+
+// REQUIRES: OS=ios && CPU=arm64e
+
+import Foundation
+
+@objc protocol P { @objc optional func testFunction(_ i: SamplesType) -> SamplesType }
+
+class C: P { func testFunction(_ i: SamplesType) -> SamplesType { samples() } }
+
+func test() {
+_ = (C() as P).testFunction?(samples())
+}
+
+// Make sure the ptrauth discriminator at closure build and invocation time match.
+
+// CHECK: @"$sTa.ptrauth" = private constant {{.*}} ptr @"$sTa"{{.*}} i64 55683 }, section "llvm.ptrauth"
+// CHECK: define hidden swiftcc void @"$s31loadable_by_address_objc_method4testyyF"()
+// CHECK:   store {{.*}} @"$sTa.ptrauth"
+// CHECK:   call swiftcc void {{.*}}(ptr {{.*}}sret(%TSo11SamplesTypeV) {{.*}} [ "ptrauth"(i32 0, i64 55683) ]
+// CHECK: }
+
+test()
+
+
+// SIL: sil hidden @$s31loadable_by_address_objc_method4testyyF : $@convention(thin) () -> () {
+// SIL: [[C:%.*]] = convert_function {{.*}} : $@convention(objc_method) (SamplesType, @opened({{.*}}, any P) Self) -> SamplesType to $@convention(objc_method) (@in_guaranteed SamplesType, @opened({{.*}}, any P) Self) -> @out SamplesType
+// SIL:             partial_apply [callee_guaranteed] [[C]]({{.*}}) : $@convention(objc_method) (@in_guaranteed SamplesType, @opened({{.*}}, any P) Self) -> @out SamplesType