[move-only] Fix two small errors around handling capturing of noncopyable self by local functions

lib/SIL/IR/TypeLowering.cpp

@@ -126,7 +126,8 @@ CaptureKind TypeConverter::getDeclCaptureKind(CapturedValue capture,
   if (!var->supportsMutation() && lowering.getLoweredType().isPureMoveOnly() &&
       !capture.isNoEscape()) {
       auto *param = dyn_cast<ParamDecl>(var);
-      if (!param || param->getValueOwnership() != ValueOwnership::Shared) {
+      if (!param || (param->getValueOwnership() != ValueOwnership::Shared &&
+                     !param->isSelfParameter())) {
         return CaptureKind::ImmutableBox;
       }
   }

lib/SILGen/SILGenProlog.cpp

@@ -1141,11 +1141,14 @@ static void emitCaptureArguments(SILGenFunction &SGF,
     fArg->setClosureCapture(true);
     arg = SILValue(fArg);
 
-    // If our capture is no escape and we have a noncopyable value, insert a
-    // consumable and assignable. If we have an escaping closure, we are going
-    // to emit an error later in SIL since it is illegal to capture an inout
-    // value in an escaping closure.
-    if (isInOut && ty.isPureMoveOnly() && capture.isNoEscape()) {
+    // If we have an inout noncopyable paramter, insert a consumable and
+    // assignable.
+    //
+    // NOTE: If we have an escaping closure, we are going to emit an error later
+    // in SIL since it is illegal to capture an inout value in an escaping
+    // closure. The later code knows how to handle that we have the
+    // mark_must_check here.
+    if (isInOut && ty.isPureMoveOnly()) {
       arg = SGF.B.createMarkMustCheckInst(
           Loc, arg, MarkMustCheckInst::CheckKind::ConsumableAndAssignable);
     }

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -2314,6 +2314,18 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
       }
     }
 
+    // If our partial apply takes this parameter as an inout parameter and it
+    // has the no move only diagnostics marker on it, do not emit an error
+    // either.
+    if (auto *f = pas->getCalleeFunction()) {
+      if (f->hasSemanticsAttr(semantics::NO_MOVEONLY_DIAGNOSTICS)) {
+        if (ApplySite(pas).getArgumentOperandConvention(*op).isInoutConvention()) {
+          diagnosticEmitter.emitEarlierPassEmittedDiagnostic(markedValue);
+          return false;
+        }
+      }
+    }
+
     if (pas->isOnStack() ||
         ApplySite(pas).getArgumentConvention(*op).isInoutConvention()) {
       LLVM_DEBUG(llvm::dbgs() << "Found on stack partial apply or inout usage!\n");

test/SILGen/moveonly.swift

@@ -2403,7 +2403,7 @@ public func testSubscriptReadModify_BaseLoadable_ResultAddressOnly_InOut(m: inou
 }
 
 // CHECK-LABEL: sil [ossa] @$s8moveonly61testSubscriptReadModify_BaseLoadable_ResultAddressOnly_GlobalyyF : $@convention(thin) () -> () {
-// CHECK: [[GLOBAL_ADDR:%.*]] = global_addr @$s8moveonly39globalLoadableSubscriptReadModifyTesterAA0cdefG0Vvp : 
+// CHECK: [[GLOBAL_ADDR:%.*]] = global_addr @$s8moveonly39globalLoadableSubscriptReadModifyTesterAA0cdefG0Vvp :
 //
 // The get call
 // CHECK: [[ACCESS:%.*]] = begin_access [read] [dynamic] [[GLOBAL_ADDR]]
@@ -3117,3 +3117,96 @@ public func testSubscriptGetModifyThroughParentClass_BaseLoadable_ResultAddressO
     m.computedTester2[0].mutatingFunc()
 }
 
+//////////////////////////////
+// MARK: Capture Self Tests //
+//////////////////////////////
+
+func testSelfCaptureHandledCorrectly() {
+    struct E {
+        var a: [Int] = []
+    }
+
+    struct Test : ~Copyable {
+        var e: E
+
+        // Test that we capture inits by address.
+        //
+        // CHECK-LABEL: sil private [ossa] @$s8moveonly31testSelfCaptureHandledCorrectlyyyF4TestL_VADycfC : $@convention(method) (@thin Test.Type) -> @owned Test {
+        // CHECK: [[BOX:%.*]] = alloc_box ${ var Test }
+        // CHECK: [[MARK_UNINIT:%.*]] = mark_uninitialized [rootself] [[BOX]]
+        // CHECK: [[BORROW:%.*]] = begin_borrow [lexical] [[MARK_UNINIT]]
+        // CHECK: [[PROJECT:%.*]] = project_box [[BORROW]]
+        // CHECK: apply {{%.*}}([[PROJECT]])
+        // CHECK: } // end sil function '$s8moveonly31testSelfCaptureHandledCorrectlyyyF4TestL_VADycfC'
+        init() {
+            e = E()
+            func capture() {
+                let e = self.e
+            }
+            capture()
+        }
+
+        // CHECK-LABEL: sil private [ossa] @$s8moveonly31testSelfCaptureHandledCorrectlyyyF4TestL_V22captureByLocalFunctionyyF : $@convention(method) (@guaranteed Test) -> () {
+        // CHECK: bb0([[ARG:%.*]] : @guaranteed $Test):
+        // CHECK:   [[COPY:%.*]] = copy_value [[ARG]]
+        // CHECK:   [[MARK:%.*]] = mark_must_check [no_consume_or_assign] [[COPY]]
+        // CHECK:   [[BORROW:%.*]] = begin_borrow [[MARK]]
+        // CHECK:   apply {{%.*}}([[BORROW]])
+        // CHECK:   end_borrow [[BORROW]]
+        // CHECK:   destroy_value [[MARK]]
+        // CHECK: } // end sil function '$s8moveonly31testSelfCaptureHandledCorrectlyyyF4TestL_V22captureByLocalFunctionyyF'
+        func captureByLocalFunction() {
+            func capture() {
+                let e = self.e
+            }
+            capture()
+        }
+
+        // CHECK-LABEL: sil private [ossa] @$s8moveonly31testSelfCaptureHandledCorrectlyyyF4TestL_V17captureByLocalLetyyF : $@convention(method) (@guaranteed Test) -> () {
+        // CHECK: bb0([[ARG:%.*]] : @guaranteed $Test):
+        // CHECK:   [[COPY:%.*]] = copy_value [[ARG]]
+        // CHECK:   [[MARK:%.*]] = mark_must_check [no_consume_or_assign] [[COPY]]
+        // CHECK:   [[COPY2:%.*]] = copy_value [[MARK]]
+        // CHECK:   [[PAI:%.*]] = partial_apply [callee_guaranteed] {{%.*}}([[COPY2]])
+        // CHECK:   destroy_value [[MARK]]
+        // CHECK: } // end sil function '$s8moveonly31testSelfCaptureHandledCorrectlyyyF4TestL_V17captureByLocalLetyyF'
+        func captureByLocalLet() {
+            let f = {
+                let e = self.e
+            }
+
+            f()
+        }
+
+        // CHECK-LABEL: sil private [ossa] @$s8moveonly31testSelfCaptureHandledCorrectlyyyF4TestL_V17captureByLocalVaryyF : $@convention(method) (@guaranteed Test) -> () {
+        // CHECK: bb0([[ARG:%.*]] : @guaranteed $Test):
+        // CHECK:   [[COPY:%.*]] = copy_value [[ARG]]
+        // CHECK:   [[MARK:%.*]] = mark_must_check [no_consume_or_assign] [[COPY]]
+        // CHECK:   [[COPY2:%.*]] = copy_value [[MARK]]
+        // CHECK:   [[PAI:%.*]] = partial_apply [callee_guaranteed] {{%.*}}([[COPY2]])
+        // CHECK:   destroy_value [[MARK]]
+        // CHECK: } // end sil function '$s8moveonly31testSelfCaptureHandledCorrectlyyyF4TestL_V17captureByLocalVaryyF'
+        func captureByLocalVar() {
+            var f = {}
+            f = {
+                let e = self.e
+            }
+            f()
+        }
+
+        // CHECK-LABEL: sil private [ossa] @$s8moveonly31testSelfCaptureHandledCorrectlyyyF4TestL_V27captureByNonEscapingClosureyyF : $@convention(method) (@guaranteed Test) -> () {
+        // CHECK: bb0([[ARG:%.*]] : @guaranteed $Test):
+        // CHECK:   [[COPY:%.*]] = copy_value [[ARG]]
+        // CHECK:   [[MARK:%.*]] = mark_must_check [no_consume_or_assign] [[COPY]]
+        // CHECK:   [[COPY2:%.*]] = copy_value [[MARK]]
+        // CHECK:   [[PAI:%.*]] = partial_apply [callee_guaranteed] {{%.*}}([[COPY2]])
+        // CHECK:   destroy_value [[MARK]]
+        // CHECK: } // end sil function '$s8moveonly31testSelfCaptureHandledCorrectlyyyF4TestL_V27captureByNonEscapingClosureyyF'
+        func captureByNonEscapingClosure() {
+            func useClosure(_ f: () -> ()) {}
+            useClosure {
+                let e = self.e
+            }
+        }
+    }
+}

test/SILGen/moveonly_escaping_closure.swift

@@ -614,7 +614,8 @@ func testConsumingEscapeClosureCaptureLet(_ f: consuming @escaping () -> ()) {
 // CHECK: } // end sil function '$s16moveonly_closure29testGlobalClosureCaptureInOutyyAA9SingleEltVzF'
 
 // CHECK-LABEL: sil private [ossa] @$s16moveonly_closure29testGlobalClosureCaptureInOutyyAA9SingleEltVzFyycfU_ : $@convention(thin) (@inout_aliasable SingleElt) -> () {
-// CHECK: bb0([[PROJECT:%.*]] : @closureCapture
+// CHECK: bb0([[ARG:%.*]] : @closureCapture
+// CHECK:   [[PROJECT:%.*]] = mark_must_check [consumable_and_assignable] [[ARG]]
 //
 // CHECK:   [[ACCESS:%.*]] = begin_access [read] [unknown] [[PROJECT]]
 // CHECK:   [[LOADED:%.*]] = load [copy] [[ACCESS]]
@@ -657,7 +658,8 @@ func testGlobalClosureCaptureInOut(_ x: inout SingleElt) {
 // CHECK: } // end sil function '$s16moveonly_closure31testLocalLetClosureCaptureInOutyyAA9SingleEltVzF'
 //
 // CHECK-LABEL: sil private [ossa] @$s16moveonly_closure31testLocalLetClosureCaptureInOutyyAA9SingleEltVzFyycfU_ : $@convention(thin) (@inout_aliasable SingleElt) -> () {
-// CHECK: bb0([[PROJECT:%.*]] : @closureCapture
+// CHECK: bb0([[ARG:%.*]] : @closureCapture
+// CHECK:   [[PROJECT:%.*]] = mark_must_check [consumable_and_assignable] [[ARG]]
 //
 // CHECK:   [[ACCESS:%.*]] = begin_access [read] [unknown] [[PROJECT]]
 // CHECK:   [[LOADED:%.*]] = load [copy] [[ACCESS]]
@@ -704,7 +706,8 @@ func testLocalLetClosureCaptureInOut(_ x: inout SingleElt) {
 // CHECK: } // end sil function '$s16moveonly_closure31testLocalVarClosureCaptureInOutyyAA9SingleEltVzF'
 //
 // CHECK-LABEL: sil private [ossa] @$s16moveonly_closure31testLocalVarClosureCaptureInOutyyAA9SingleEltVzFyycfU_ : $@convention(thin) (@inout_aliasable SingleElt) -> () {
-// CHECK: bb0([[PROJECT:%.*]] : @closureCapture
+// CHECK: bb0([[ARG:%.*]] : @closureCapture
+// CHECK:   [[PROJECT:%.*]] = mark_must_check [consumable_and_assignable] [[ARG]]
 //
 // CHECK:   [[ACCESS:%.*]] = begin_access [read] [unknown] [[PROJECT]]
 // CHECK:   [[LOADED:%.*]] = load [copy] [[ACCESS]]
@@ -752,7 +755,8 @@ func testLocalVarClosureCaptureInOut(_ x: inout SingleElt) {
 // CHECK: } // end sil function '$s16moveonly_closure026testInOutVarClosureCapturedE0yyyycz_AA9SingleEltVztF'
 //
 // CHECK-LABEL: sil private [ossa] @$s16moveonly_closure026testInOutVarClosureCapturedE0yyyycz_AA9SingleEltVztFyycfU_ : $@convention(thin) (@inout_aliasable SingleElt) -> () {
-// CHECK: bb0([[PROJECT:%.*]] : @closureCapture
+// CHECK: bb0([[ARG:%.*]] : @closureCapture
+// CHECK:   [[PROJECT:%.*]] = mark_must_check [consumable_and_assignable] [[ARG]]
 //
 // CHECK:   [[ACCESS:%.*]] = begin_access [read] [unknown] [[PROJECT]]
 // CHECK:   [[LOADED:%.*]] = load [copy] [[ACCESS]]
@@ -807,7 +811,8 @@ func testInOutVarClosureCaptureInOut(_ f: inout () -> (), _ x: inout SingleElt)
 // CHECK: } // end sil function '$s16moveonly_closure38testConsumingEscapeClosureCaptureInOutyyyycn_AA9SingleEltVztF'
 //
 // CHECK-LABEL: sil private [ossa] @$s16moveonly_closure38testConsumingEscapeClosureCaptureInOutyyyycn_AA9SingleEltVztFyycfU_ : $@convention(thin) (@inout_aliasable SingleElt) -> () {
-// CHECK: bb0([[PROJECT:%.*]] : @closureCapture
+// CHECK: bb0([[ARG:%.*]] : @closureCapture
+// CHECK:   [[PROJECT:%.*]] = mark_must_check [consumable_and_assignable] [[ARG]]
 //
 // CHECK:   [[ACCESS:%.*]] = begin_access [read] [unknown] [[PROJECT]]
 // CHECK:   [[LOADED:%.*]] = load [copy] [[ACCESS]]

test/SILOptimizer/moveonly_self_captures.swift

@@ -0,0 +1,107 @@
+// RUN: %target-swift-emit-sil -sil-verify-all -verify %s
+
+class Klass {}
+
+struct E {
+    var k = Klass()
+}
+
+struct E2 : ~Copyable {
+    var k = Klass()
+}
+
+var g: () -> () = {}
+struct Test : ~Copyable {
+    var e: E
+    var e2: E2
+
+    // Test that we capture inits by address.
+    init() {
+        e = E()
+        e2 = E2()
+        func capture() {
+            let _ = self.e
+        }
+        capture()
+    }
+
+    init(x: ()) {
+        e = E()
+        e2 = E2()
+        func capture() {
+            let _ = self
+            let _ = self.e2 // expected-error {{cannot partially consume 'self'}}
+        }
+        capture()
+    }
+
+    init(y: ()) { // expected-error {{missing reinitialization of closure capture 'self' after consume}}
+        e = E()
+        e2 = E2()
+        func capture() {
+            let _ = self // expected-note {{consumed here}}
+        }
+        capture()
+    }
+
+    init(z: ()) {
+        e = E()
+        e2 = E2()
+        func capture() {
+            let _ = self // expected-note {{captured here}}
+        }
+        capture()
+        g = capture // expected-error {{escaping local function captures mutating 'self' parameter}}
+    }
+
+    func captureByLocalFunction() {
+        func capture() {
+            let _ = self.e
+        }
+        capture()
+    }
+
+    func captureByLocalFunction2() { // expected-error {{noncopyable 'self' cannot be consumed when captured by an escaping closure}}
+        func capture() {
+            let _ = self.e2 // expected-note {{consumed here}}
+        }
+        capture()
+    }
+
+    func captureByLocalFunction3() { // expected-error {{noncopyable 'self' cannot be consumed when captured by an escaping closure}}
+        func capture() {
+            let _ = self // expected-note {{consumed here}}
+        }
+        capture()
+    }
+
+    func captureByLocalLet() { // expected-error {{'self' cannot be captured by an escaping closure since it is a borrowed parameter}}
+        let f = { // expected-note {{capturing 'self' here}}
+            let _ = self.e
+        }
+
+        f()
+    }
+
+    func captureByLocalVar() { // expected-error {{'self' cannot be captured by an escaping closure since it is a borrowed parameter}}
+        var f = {}
+        f = { // expected-note {{closure capturing 'self' here}}
+            let _ = self.e
+        }
+        f()
+    }
+
+    func captureByNonEscapingClosure() {
+        func useClosure(_ f: () -> ()) {}
+        useClosure {
+            let _ = self.e
+        }
+    }
+
+    func captureByNonEscapingClosure2() { // expected-error {{'self' cannot be consumed when captured by an escaping closure}}
+        func useClosure(_ f: () -> ()) {}
+        useClosure {
+            let _ = self // expected-note {{consumed here}}
+        }
+    }
+}