[SILGen] InitAccessors: Make sure that assign_or_init always directly references self

`nonmutating set` gets a copy of "self" in `GetterSetterComponent`
which is expected for partial application of the setter but doesn't
work for "self" reference that `assign_or_init` instruction needs
to emit references to stored properties during lowering. We need to
make sure that "self" is always a reference to rootself of the
constructor before passing it to `assign_or_init`.

Resolves: #67827
Resolves: rdar://114433261
(cherry picked from commit 10947de)

Author: xedin

lib/SILGen/SILGenFunction.cpp

@@ -1859,6 +1859,19 @@ void SILGenFunction::emitAssignOrInit(SILLocation loc, ManagedValue selfValue,
     setterFRef = SILUndef::get(initFRef->getType(), F);
   }
 
-  B.createAssignOrInit(loc, selfValue.getValue(), newValue.forward(*this),
+  auto isValueSelf = !selfValue.getType().getASTType()->mayHaveSuperclass();
+  // If we are emitting `assign_or_init` instruction for a value
+  // type, we need to make sure that "self" is always a l-value
+  // reference to "rootself" because `nonmutating set` loads "self"
+  // and referencing `selfValue` in such case is incorrect because
+  // it's a copy which is going to be skipped by DI.
+  auto selfRef = selfValue;
+  if (isValueSelf && !selfRef.isLValue()) {
+    auto *ctor = cast<ConstructorDecl>(FunctionDC->getAsDecl());
+    selfRef = maybeEmitValueOfLocalVarDecl(ctor->getImplicitSelfDecl(),
+                                           AccessKind::ReadWrite);
+  }
+
+  B.createAssignOrInit(loc, selfRef.getValue(), newValue.forward(*this),
                        initFRef, setterFRef, AssignOrInitInst::Unknown);
 }

test/Interpreter/init_accessors.swift

@@ -846,3 +846,77 @@ do {
 // CHECK-NEXT: TestMixedDefaultInitalizable(a: nil, b: Optional(""))
 // CHECK-NEXT: TestMixedDefaultInitalizable(a: nil, b: Optional("Hello"))
 // CHECK-NEXT: TestMixedDefaultInitalizable(a: nil, b: Optional(""))
+
+do {
+  struct TestNonMutatingSetDefault {
+    var _count: Int
+
+    var count: Int = 42 {
+      @storageRestrictions(initializes: _count)
+      init {
+        _count = newValue
+      }
+
+      get { _count }
+      nonmutating set {}
+    }
+  }
+
+  struct TestNonMutatingSetNoDefault {
+    var _count: Int
+
+    var count: Int {
+      @storageRestrictions(initializes: _count)
+      init {
+        print("init accessor is called: \(newValue)")
+        _count = newValue
+      }
+
+      get { _count }
+
+      nonmutating set {
+        print("nonmutating set called: \(newValue)")
+      }
+    }
+
+    init(value: Int) {
+      self.count = value
+      self.count = value + 1
+    }
+  }
+
+  struct TestNonMutatingSetCustom {
+    var _count: Int
+
+    var count: Int = 42 {
+      @storageRestrictions(initializes: _count)
+      init {
+        print("init accessor is called: \(newValue)")
+        _count = newValue
+      }
+
+      get { _count }
+
+      nonmutating set {
+        print("nonmutating set called: \(newValue)")
+      }
+    }
+
+    init(custom: Int) {
+      count = custom
+    }
+  }
+
+  print("test-nonmutating-set-1: \(TestNonMutatingSetDefault())")
+  print("test-nonmutating-set-2: \(TestNonMutatingSetDefault(count: 0))")
+  print("test-nonmutating-set-3: \(TestNonMutatingSetNoDefault(value: -1))")
+  print("test-nonmutating-set-4: \(TestNonMutatingSetCustom(custom: 0))")
+}
+// CHECK: test-nonmutating-set-1: TestNonMutatingSetDefault(_count: 42)
+// CHECK-NEXT: test-nonmutating-set-2: TestNonMutatingSetDefault(_count: 0)
+// CHECK-NEXT: init accessor is called: -1
+// CHECK-NEXT: nonmutating set called: 0
+// CHECK-NEXT: test-nonmutating-set-3: TestNonMutatingSetNoDefault(_count: -1)
+// CHECK-NEXT: init accessor is called: 42
+// CHECK-NEXT: nonmutating set called: 0
+// CHECK-NEXT: test-nonmutating-set-4: TestNonMutatingSetCustom(_count: 42)

test/SILOptimizer/init_accessor_raw_sil_lowering.swift

@@ -225,7 +225,7 @@ func test_default_inits() {
     // CHECK: [[INIT_ACCESSOR:%.*]] = function_ref @$s23assign_or_init_lowering18test_default_initsyyF5Test1L_V1xSivi : $@convention(thin) (Int) -> @out Int
     // CHECK: [[SETTER_REF:%.*]] = function_ref @$s23assign_or_init_lowering18test_default_initsyyF5Test1L_V1xSivs : $@convention(method) (Int, @inout Test1) -> ()
     // CHECK-NEXT: [[SETTER:%.*]] = partial_apply [callee_guaranteed] [[SETTER_REF]]([[SELF_REF]]) : $@convention(method) (Int, @inout Test1) -> ()
-    // CHECK-NEXT: assign_or_init [init] self [[SELF_REF]] : $*Test1, value [[INIT_VAL]] : $Int, init [[INIT_ACCESSOR]] : $@convention(thin) (Int) -> @out Int, set [[SETTER]] : $@callee_guaranteed (Int) -> ()
+    // CHECK-NEXT: assign_or_init [init] self %1 : $*Test1, value [[INIT_VAL]] : $Int, init [[INIT_ACCESSOR]] : $@convention(thin) (Int) -> @out Int, set [[SETTER]] : $@callee_guaranteed (Int) -> ()
     // CHECK-NEXT: end_access [[SELF_REF]] : $*Test1
     // CHECK-NEXT: destroy_value [[SETTER]] : $@callee_guaranteed (Int) -> ()
     init() {
@@ -238,7 +238,7 @@ func test_default_inits() {
     // CHECK: [[INIT_ACCESSOR:%.*]] = function_ref @$s23assign_or_init_lowering18test_default_initsyyF5Test1L_V1xSivi : $@convention(thin) (Int) -> @out Int
     // CHECK: [[SETTER_REF:%.*]] = function_ref @$s23assign_or_init_lowering18test_default_initsyyF5Test1L_V1xSivs : $@convention(method) (Int, @inout Test1) -> ()
     // CHECK-NEXT: [[SETTER:%.*]] = partial_apply [callee_guaranteed] [[SETTER_REF]]([[SELF_REF]]) : $@convention(method) (Int, @inout Test1) -> ()
-    // CHECK-NEXT: assign_or_init [init] self [[SELF_REF]] : $*Test1, value [[INIT_VAL]] : $Int, init [[INIT_ACCESSOR]] : $@convention(thin) (Int) -> @out Int, set [[SETTER]] : $@callee_guaranteed (Int) -> ()
+    // CHECK-NEXT: assign_or_init [init] self %2 : $*Test1, value [[INIT_VAL]] : $Int, init [[INIT_ACCESSOR]] : $@convention(thin) (Int) -> @out Int, set [[SETTER]] : $@callee_guaranteed (Int) -> ()
     // CHECK-NEXT: end_access [[SELF_REF]] : $*Test1
     // CHECK-NEXT: destroy_value [[SETTER]] : $@callee_guaranteed (Int) -> ()
     //
@@ -396,3 +396,33 @@ func test_handling_of_superclass_properties() {
     }
   }
 }
+
+func test_handling_of_nonmutating_set() {
+  struct Test {
+    private var _count: Int
+
+    var count: Int = 42 {
+      @storageRestrictions(initializes: _count)
+      init {
+        _count = newValue
+      }
+      get {
+        _count
+      }
+      nonmutating set {
+        // Update store
+      }
+    }
+
+    // CHECK-LABEL: sil private [ossa] @$s23assign_or_init_lowering32test_handling_of_nonmutating_setyyF4TestL_V5countADSi_tcfC : $@convention(method) (Int, @thin Test.Type) -> Test
+    // CHECK: [[INIT_VALUE:%.*]] = function_ref @$s23assign_or_init_lowering32test_handling_of_nonmutating_setyyF4TestL_V5countSivpfi : $@convention(thin) () -> Int
+    // CHECK-NEXT: [[VALUE:%.*]] = apply [[INIT_VALUE]]() : $@convention(thin) () -> Int
+    // CHECK: assign_or_init [init] self %2 : $*Test, value [[VALUE]] : $Int, init {{.*}} : $@convention(thin) (Int) -> @out Int, set {{.*}} : $@callee_guaranteed (Int) -> ()
+    // CHECK: assign_or_init [set] self %2 : $*Test, value %0 : $Int, init {{.*}} : $@convention(thin) (Int) -> @out Int, set {{.*}} : $@callee_guaranteed (Int) -> ()
+    // CHECK: assign_or_init [set] self %2 : $*Test, value [[ZERO:%.*]] : $Int, init {{.*}} : $@convention(thin) (Int) -> @out Int, set {{.*}} : $@callee_guaranteed (Int) -> ()
+    init(count: Int) {
+      self.count = count
+      self.count = 0
+    }
+  }
+}