Merge pull request #36011 from aschwaighofer/irgen_fix_partial_apply_self

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

Author: aschwaighofer

lib/IRGen/IRGenSIL.cpp

@@ -2974,7 +2974,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))