[SILGen] InitAccessors: Prevent setter closure from escaping

The closure with applied base is not escaping and gets applied
only once (when self is fully initialized). Let's make sure that
the partial application results in on-stack closure that borrows
"self" instead of copying it.

Author: xedin

lib/SILGen/SILGenFunction.cpp

@@ -1821,8 +1821,8 @@ SILGenFunction::emitApplyOfSetterToBase(SILLocation loc, SILDeclRef setter,
     SmallVector<ManagedValue, 4> captures;
     emitCaptures(loc, setter, CaptureEmission::AssignByWrapper, captures);
 
-    for (auto capture : captures)
-      capturedArgs.push_back(capture.forward(*this));
+    llvm::transform(captures, std::back_inserter(capturedArgs),
+                    [](auto &capture) { return capture.getValue(); });
   } else {
     assert(base);
 
@@ -1840,16 +1840,20 @@ SILGenFunction::emitApplyOfSetterToBase(SILLocation loc, SILDeclRef setter,
       // nonmutating setter, for example.
       capturedBase = B.createTrivialLoadOr(loc, base.getValue(),
                                            LoadOwnershipQualifier::Copy);
+      // On-stack partial apply doesn't take ownership of the base, so
+      // we have to destroy it manually.
+      enterDestroyCleanup(capturedBase);
     } else {
-      capturedBase = base.copy(*this, loc).forward(*this);
+      capturedBase = base.borrow(*this, loc).getValue();
     }
 
     capturedArgs.push_back(capturedBase);
   }
 
   PartialApplyInst *setterPAI =
       B.createPartialApply(loc, setterFRef, substitutions, capturedArgs,
-                           ParameterConvention::Direct_Guaranteed);
+                           ParameterConvention::Direct_Guaranteed,
+                           PartialApplyInst::OnStackKind::OnStack);
   return emitManagedRValueWithCleanup(setterPAI).getValue();
 }
 

test/SILGen/objc_properties.swift

@@ -289,11 +289,11 @@ class SomeWrapperTests {
 
 // CHECK-LABEL: sil hidden [ossa] @$s15objc_properties16SomeWrapperTestsCyACSScfc : $@convention(method) (@owned String, @owned SomeWrapperTests) -> @owned SomeWrapperTests {
 // CHECK:  [[M:%.*]] = function_ref @$s15objc_properties16SomeWrapperTestsC04someD0SivsTD
-// CHECK:  [[C:%.*]] = partial_apply [callee_guaranteed] [[M]]({{.*}})
-// CHECK: assign_by_wrapper {{%.*}}: $Int to {{%.*}} : $*SomeWrapper, init {{.*}} : $@callee_guaranteed (Int) -> SomeWrapper, set [[C]] : $@callee_guaranteed (Int) -> ()
+// CHECK:  [[C:%.*]] = partial_apply [callee_guaranteed] [on_stack] [[M]]({{.*}})
+// CHECK: assign_by_wrapper {{%.*}}: $Int to {{%.*}} : $*SomeWrapper, init {{.*}} : $@callee_guaranteed (Int) -> SomeWrapper, set [[C]] : $@noescape @callee_guaranteed (Int) -> ()
 // CHECK: [[M:%.*]] = function_ref @$s15objc_properties16SomeWrapperTestsC1sSSSgvsTD
-// CHECK: [[C:%.*]] = partial_apply [callee_guaranteed] [[M]](
-// CHECK:  assign_by_wrapper {{.*}} : $Optional<String> to {{.*}} : $*W<Optional<String>>, init {{.*}} : $@callee_guaranteed (@owned Optional<String>) -> @owned W<Optional<String>>, set [[C]] : $@callee_guaranteed (@owned Optional<String>) -> ()
+// CHECK: [[C:%.*]] = partial_apply [callee_guaranteed] [on_stack] [[M]](
+// CHECK:  assign_by_wrapper {{.*}} : $Optional<String> to {{.*}} : $*W<Optional<String>>, init {{.*}} : $@callee_guaranteed (@owned Optional<String>) -> @owned W<Optional<String>>, set [[C]] : $@noescape @callee_guaranteed (@owned Optional<String>) -> ()
   init(_ s: String) {
     someWrapper = 1000
     self.s = s

test/SILGen/property_wrapper_local.swift

@@ -25,8 +25,8 @@ func testLocalWrapper() {
   // CHECK: [[S:%.*]] = function_ref @$s22property_wrapper_local16testLocalWrapperyyF5valueL_Sivs : $@convention(thin) (Int, @guaranteed { var Wrapper<Int> }) -> ()
   // CHECK-NEXT: [[C:%.*]] = copy_value [[W]] : ${ var Wrapper<Int> }
   // CHECK-NOT: mark_function_escape
-  // CHECK-NEXT: [[SPA:%.*]] = partial_apply [callee_guaranteed] [[S]]([[C]]) : $@convention(thin) (Int, @guaranteed { var Wrapper<Int> }) -> ()
-  // CHECK-NEXT: assign_by_wrapper {{%.*}} : $Int to [[P]] : $*Wrapper<Int>, init [[IPA]] : $@callee_guaranteed (Int) -> Wrapper<Int>, set [[SPA]] : $@callee_guaranteed (Int) -> ()
+  // CHECK-NEXT: [[SPA:%.*]] = partial_apply [callee_guaranteed] [on_stack] [[S]]([[C]]) : $@convention(thin) (Int, @guaranteed { var Wrapper<Int> }) -> ()
+  // CHECK-NEXT: assign_by_wrapper {{%.*}} : $Int to [[P]] : $*Wrapper<Int>, init [[IPA]] : $@callee_guaranteed (Int) -> Wrapper<Int>, set [[SPA]] : $@noescape @callee_guaranteed (Int) -> ()
 
   _ = value
   // CHECK: mark_function_escape [[P]] : $*Wrapper<Int>

test/SILGen/property_wrappers.swift

@@ -903,8 +903,8 @@ struct NonMutatingSetterWrapper<Value> {
 struct NonMutatingWrapperTestStruct {
     // CHECK-LABEL: sil hidden [ossa] @$s17property_wrappers28NonMutatingWrapperTestStructV3valACSi_tcfC : $@convention(method) (Int, @thin NonMutatingWrapperTestStruct.Type) -> NonMutatingWrapperTestStruct {
     // CHECK: %[[LOAD:[0-9]+]] = load [trivial] %[[SRC:[0-9]+]] : $*NonMutatingWrapperTestStruct
-    // CHECK-NEXT: %[[SET_PA:[0-9]+]] = partial_apply [callee_guaranteed] %[[PW_SETTER:[0-9]+]](%[[LOAD]]) : $@convention(method) (Int, NonMutatingWrapperTestStruct) -> ()
-    // CHECK-NEXT: assign_by_wrapper %[[SETVAL:[0-9]+]] : $Int to %[[ADDR:[0-9]+]] : $*NonMutatingSetterWrapper<Int>, init %[[INIT_PA:[0-9]+]] : $@callee_guaranteed (Int) -> NonMutatingSetterWrapper<Int>, set %[[SET_PA]] : $@callee_guaranteed (Int) -> ()
+    // CHECK-NEXT: %[[SET_PA:[0-9]+]] = partial_apply [callee_guaranteed] [on_stack] %[[PW_SETTER:[0-9]+]](%[[LOAD]]) : $@convention(method) (Int, NonMutatingWrapperTestStruct) -> ()
+    // CHECK-NEXT: assign_by_wrapper %[[SETVAL:[0-9]+]] : $Int to %[[ADDR:[0-9]+]] : $*NonMutatingSetterWrapper<Int>, init %[[INIT_PA:[0-9]+]] : $@callee_guaranteed (Int) -> NonMutatingSetterWrapper<Int>, set %[[SET_PA]] : $@noescape @callee_guaranteed (Int) -> ()
     @NonMutatingSetterWrapper var SomeProp: Int
     init(val: Int) {
         SomeProp = val

test/SILGen/resilient_assign_by_wrapper.swift

@@ -33,9 +33,9 @@ public class AddressOnlySetter {
 
     // CHECK: [[E2:%.*]] = alloc_stack $AddressOnlyEnum
     // CHECK-NEXT: inject_enum_addr [[E2]] : $*AddressOnlyEnum, #AddressOnlyEnum.some!enumelt
-    // CHECK: [[S:%.*]] = partial_apply [callee_guaranteed] {{%.*}}({{%.*}}) : $@convention(method) (@in AddressOnlyEnum, @guaranteed AddressOnlySetter) -> ()
+    // CHECK: [[S:%.*]] = partial_apply [callee_guaranteed] [on_stack] {{%.*}}({{%.*}}) : $@convention(method) (@in AddressOnlyEnum, @guaranteed AddressOnlySetter) -> ()
     // CHECK: assign_by_wrapper [[E2]] : $*AddressOnlyEnum
-    // CHECK-SAME: set [[S]] : $@callee_guaranteed (@in AddressOnlyEnum) -> ()
+    // CHECK-SAME: set [[S]] : $@noescape @callee_guaranteed (@in AddressOnlyEnum) -> ()
     self.value = .some
   }
 
@@ -60,7 +60,7 @@ extension SubstitutedSetter where T == Bool {
     // CHECK: [[B:%.*]] = alloc_stack $Bool
     // CHECK: assign_by_wrapper [[B]] : $*Bool to [[W]] : $*WrapGod<Bool>
     // CHECK-SAME: init {{%.*}} : $@callee_guaranteed (@in Bool) -> @out WrapGod<Bool>
-    // CHECK-SAME: set {{%.*}} : $@callee_guaranteed (@in Bool) -> ()
+    // CHECK-SAME: set {{%.*}} : $@noescape @callee_guaranteed (@in Bool) -> ()
     self.value = true
   }
 }
@@ -96,7 +96,7 @@ extension ObjectifiedSetter where T == SomeObject {
     // CHECK: [[STORAGE:%.*]] = struct_element_addr {{%.*}} : $*ObjectifiedSetter<SomeObject>, #ObjectifiedSetter._value
     // CHECK: assign_by_wrapper [[OBJ]] : $SomeObject to [[STORAGE]] : $*WrapGod<SomeObject>
     // CHECK-SAME: init {{%.*}} : $@callee_guaranteed (@owned SomeObject) -> @out WrapGod<SomeObject>
-    // CHECK-SAME: set {{%.*}} : $@callee_guaranteed (@owned SomeObject) -> ()
+    // CHECK-SAME: set {{%.*}} : $@noescape @callee_guaranteed (@owned SomeObject) -> ()
     self.value = SomeObject()
   }
 }