[5.4] IRGen: We can only reuse methods for partial apply thunks if their abi uses swiftself for self

Explanation: A mismatch between the ABI of methods and the partial apply
thunk they are used for in an optimization -- the swiftself argument of
the closure context is is not reflected in the method's abi -- causes
crashes. Only certain methods use the swiftself register for the self
parameter. The code determining reuse of the method for the partial
apply thunk did not correctly compute the match.

Origination: This was introduced by an optimization in a commit in May.

Risk: Low. The fix uses the function that method signature lowering use
to determine whether swiftself is used in the method signature.

Testing: A regression test was added.

rdar://73777202

Author: aschwaighofer

lib/IRGen/IRGenSIL.cpp

@@ -2956,7 +2956,12 @@ static bool isSimplePartialApply(IRGenFunction &IGF, PartialApplyInst *i) {
   // handled by a simplification pass in SIL.)
   if (i->getNumArguments() != 1)
     return false;
-  
+  // The closure application is going to expect to pass the context in swiftself
+  // only methods where the call to `hasSelfContextParameter` returns true will
+  // use swiftself for the self parameter.
+  if (!hasSelfContextParameter(calleeTy))
+    return false;
+
   auto appliedParam = calleeTy->getParameters().back();
   if (resultTy->isNoEscape()) {
     // A trivial closure accepts an unowned or guaranteed argument, possibly

test/IRGen/simple_partial_apply.sil

@@ -23,10 +23,11 @@ entry(%body : $@convention(method) (Int, @guaranteed C) -> Int, %context : $C):
   return %closure : $@callee_guaranteed (Int) -> Int
 }
 
+// Can't reuse the method because it does not have swiftself.
+
 // CHECK-LABEL: define {{.*}} @escape_partial_apply_swift_single_refcount_struct
-// CHECK-arm64e:  call i64 @llvm.ptrauth.resign.i64
-// CHECK:      [[FPTR:%.*]] = insertvalue { i8*, %swift.refcounted* } undef, i8* {{.*}}, 0
-// CHECK-NEXT: [[FCTX:%.*]] = insertvalue { i8*, %swift.refcounted* } [[FPTR]], %swift.refcounted* {{.*}}, 1
+// CHECK: [[CTXT:%.*]] = call {{.*}} @swift_allocObject
+// CHECK: [[FCTX:%.*]] = insertvalue { i8*, %swift.refcounted* } { i8* bitcast (i{{(64|32)}} (i{{(64|32)}}, %swift.refcounted*)* @"$sTA" to i8*), %swift.refcounted* undef }, %swift.refcounted* [[CTXT]], 1
 // CHECK-NEXT: ret { i8*, %swift.refcounted* } [[FCTX]]
 sil @escape_partial_apply_swift_single_refcount_struct : $@convention(thin) (@convention(method) (Int, @guaranteed SingleRefcounted) -> Int, @guaranteed SingleRefcounted) -> @callee_guaranteed (Int) -> Int {
 entry(%body : $@convention(method) (Int, @guaranteed SingleRefcounted) -> Int, %context : $SingleRefcounted):
@@ -47,11 +48,9 @@ entry(%body : $@convention(method) (Int, @in_guaranteed C) -> Int, %context : $*
   return undef : $()
 }
 
+// Can't reuse the method because it does not have swiftself.
 // CHECK-LABEL: define {{.*}} @noescape_partial_apply_swift_direct_word
-// CHECK-arm64e:  call i64 @llvm.ptrauth.resign.i64
-// CHECK:         [[CTX:%.*]] = inttoptr i{{.*}} %1 to %swift.opaque*
-// CHECK-NEXT:    [[CONT:%.*]] = bitcast i8* %2
-// CHECK-NEXT:    call {{.*}}void [[CONT]](i8* {{.*}}, %swift.opaque* [[CTX]], %swift.refcounted* {{.*}}%3)
+// CHECK: call swiftcc void %11(i8* bitcast (i{{(64|32)}} (i{{(64|32)}}, %swift.refcounted*)* @"$sTA.1" to i8*), %swift.opaque* {{.*}}, %swift.refcounted* swiftself {{.*}})
 sil @noescape_partial_apply_swift_direct_word : $@convention(thin) (@convention(method) (Int, Int) -> Int, Int, @guaranteed @callee_guaranteed (@noescape @callee_guaranteed (Int) -> Int) -> ()) -> () {
 entry(%body : $@convention(method) (Int, Int) -> Int, %context : $Int, %cont : $@callee_guaranteed (@noescape @callee_guaranteed (Int) -> Int) -> ()):
   %closure = partial_apply [callee_guaranteed] [on_stack] %body(%context) : $@convention(method) (Int, Int) -> Int

test/IRGen/simple_partial_apply_or_not.swift

@@ -0,0 +1,42 @@
+// RUN: %target-swift-emit-ir -module-name test %s | %FileCheck %s
+// RUN: %target-run-simple-swift %s | %FileCheck %s --check-prefix=CHECK-EXEC
+
+@propertyWrapper
+struct State<T> {
+  private class Reference {
+    var value: T
+    init(value: T) { self.value = value }
+  }
+
+  private let ref: Reference
+
+  init(wrappedValue: T) {
+    ref = Reference(value: wrappedValue)
+  }
+
+  var wrappedValue: T {
+    get { ref.value }
+    nonmutating set { ref.value = newValue }
+  }
+}
+
+struct S {
+  @State var value: Int = 1
+
+  init() {
+    value = 10 // CRASH
+  }
+}
+
+print("Hello!")
+let s = S()
+print(s)
+
+// CHECK: define {{.*}}swiftcc %T4test5StateV9Reference33_C903A018FCE7355FD30EF8324850EB90LLCySi_G* @"$s4test1SVACycfC"()
+// CHECK:  call swiftcc void @"$s4test1SV5valueSivsTA"(i{{(32|64)}} 10, %swift.refcounted* swiftself {{.*}})
+// CHECK:  ret %T4test5StateV9Reference33_C903A018FCE7355FD30EF8324850EB90LLCySi_G
+
+// This used to crash.
+
+// CHECK-EXEC: Hello!
+// CHECK-EXEC: S(_value: main.State<Swift.Int>(ref: main.State<Swift.Int>.(unknown context at {{.*}}).Reference))