Consistently mark borrowing and consuming parameters for move-only checking when captured.

When a `borrowing` or `consuming` parameter is captured by a closure,
we emit references to the binding within the closure as if it is non-implicitly
copyable, but we didn't mark the bindings inside the closure for move-only
checking to ensure the uses were correct, so improper consumes would go
undiagnosed and lead to assertion failures, compiler crashes, and/or
miscompiles. Fixes rdar://127382105

Author: jckarter

lib/SILGen/SILGenLValue.cpp

@@ -1150,8 +1150,12 @@ namespace {
       // See if we have a noncopyable address from a project_box or global.
       if (Value.getType().isAddress() && Value.getType().isMoveOnly()) {
         SILValue addr = Value.getValue();
-        auto box = dyn_cast<ProjectBoxInst>(addr);
-        if (box || isa<GlobalAddrInst>(addr) || IsLazyInitializedGlobal) {
+        auto boxProject = addr;
+        if (auto m = dyn_cast<CopyableToMoveOnlyWrapperAddrInst>(boxProject)) {
+          boxProject = m->getOperand();
+        }
+        auto box = dyn_cast<ProjectBoxInst>(boxProject);
+        if (box || isa<GlobalAddrInst>(boxProject) || IsLazyInitializedGlobal) {
           if (Enforcement)
             addr = enterAccessScope(SGF, loc, base, addr, getTypeData(),
                                     getAccessKind(), *Enforcement,

lib/SILGen/SILGenProlog.cpp

@@ -1054,6 +1054,7 @@ static void emitCaptureArguments(SILGenFunction &SGF,
                                  uint16_t ArgNo) {
 
   auto *VD = cast<VarDecl>(capture.getDecl());
+  
   SILLocation Loc(VD);
   Loc.markAsPrologue();
 
@@ -1082,6 +1083,30 @@ static void emitCaptureArguments(SILGenFunction &SGF,
   auto &lowering = SGF.getTypeLowering(getVarTypeInCaptureContext());
   SILType ty = lowering.getLoweredType();
 
+  bool isNoImplicitCopy;
+  
+  if (ty.isTrivial(SGF.F) || ty.isMoveOnly()) {
+    isNoImplicitCopy = false;
+  } else if (VD->isNoImplicitCopy()) {
+    isNoImplicitCopy = true;
+  } else if (auto pd = dyn_cast<ParamDecl>(VD)) {
+    switch (pd->getSpecifier()) {
+    case ParamSpecifier::Borrowing:
+    case ParamSpecifier::Consuming:
+      isNoImplicitCopy = true;
+      break;
+    case ParamSpecifier::ImplicitlyCopyableConsuming:
+    case ParamSpecifier::Default:
+    case ParamSpecifier::InOut:
+    case ParamSpecifier::LegacyOwned:
+    case ParamSpecifier::LegacyShared:
+      isNoImplicitCopy = false;
+      break;
+    }
+  } else {
+    isNoImplicitCopy = false;
+  }
+    
   SILValue arg;
   SILFunctionArgument *box = nullptr;
 
@@ -1116,6 +1141,10 @@ static void emitCaptureArguments(SILGenFunction &SGF,
       addr->finishInitialization(SGF);
       val = addr->getManagedAddress();
     }
+    
+    if (isNoImplicitCopy && !val.getType().isMoveOnly()) {
+      val = SGF.B.createGuaranteedCopyableToMoveOnlyWrapperValue(VD, val);
+    }
 
     // If this constant is a move only type, we need to add no_consume_or_assign checking to
     // ensure that we do not consume this captured value in the function. This
@@ -1149,6 +1178,9 @@ static void emitCaptureArguments(SILGenFunction &SGF,
         SILType::getPrimitiveObjectType(boxTy), VD);
     box->setClosureCapture(true);
     arg = SGF.B.createProjectBox(VD, box, 0);
+    if (isNoImplicitCopy && !arg->getType().isMoveOnly()) {
+      arg = SGF.B.createCopyableToMoveOnlyWrapperAddr(VD, arg);
+    }
     break;
   }
   case CaptureKind::StorageAddress:
@@ -1169,6 +1201,33 @@ static void emitCaptureArguments(SILGenFunction &SGF,
     auto *fArg = SGF.F.begin()->createFunctionArgument(ty, VD);
     fArg->setClosureCapture(true);
     arg = SILValue(fArg);
+    
+    if (isNoImplicitCopy && !arg->getType().isMoveOnly()) {
+      switch (argConv) {
+      case SILArgumentConvention::Indirect_Inout:
+      case SILArgumentConvention::Indirect_InoutAliasable:
+      case SILArgumentConvention::Indirect_In:
+      case SILArgumentConvention::Indirect_In_Guaranteed:
+      case SILArgumentConvention::Pack_Inout:
+      case SILArgumentConvention::Pack_Owned:
+      case SILArgumentConvention::Pack_Guaranteed:
+        arg = SGF.B.createCopyableToMoveOnlyWrapperAddr(VD, arg);
+        break;
+        
+      case SILArgumentConvention::Direct_Owned:
+        arg = SGF.B.createOwnedCopyableToMoveOnlyWrapperValue(VD, arg);
+        break;
+      
+      case SILArgumentConvention::Direct_Guaranteed:
+        arg = SGF.B.createGuaranteedCopyableToMoveOnlyWrapperValue(VD, arg);
+        break;
+      
+      case SILArgumentConvention::Direct_Unowned:
+      case SILArgumentConvention::Indirect_Out:
+      case SILArgumentConvention::Pack_Out:
+        llvm_unreachable("should be impossible");
+      }
+    }
 
     // If we have an inout noncopyable parameter, insert a consumable and
     // assignable.
@@ -1177,7 +1236,7 @@ static void emitCaptureArguments(SILGenFunction &SGF,
     // 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_unresolved_non_copyable_value here.
-    if (isInOut && ty.isMoveOnly(/*orWrapped=*/false)) {
+    if (isInOut && arg->getType().isMoveOnly()) {
       arg = SGF.B.createMarkUnresolvedNonCopyableValueInst(
           Loc, arg,
           MarkUnresolvedNonCopyableValueInst::CheckKind::

lib/SILOptimizer/Mandatory/AccessEnforcementSelection.cpp

@@ -707,9 +707,11 @@ SourceAccess AccessEnforcementSelection::getSourceAccess(SILValue address) {
   if (auto *mmci = dyn_cast<MarkUnresolvedNonCopyableValueInst>(address))
     return getSourceAccess(mmci->getOperand());
 
-  // Recurse through moveonlywrapper_to_copyable_addr.
+  // Recur through moveonlywrapper_to_copyable_addr or vice versa.
   if (auto *m = dyn_cast<MoveOnlyWrapperToCopyableAddrInst>(address))
     return getSourceAccess(m->getOperand());
+  if (auto *c = dyn_cast<CopyableToMoveOnlyWrapperAddrInst>(address))
+    return getSourceAccess(c->getOperand());
 
   // Recurse through drop_deinit.
   if (auto *ddi = dyn_cast<DropDeinitInst>(address))

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -417,6 +417,11 @@ static bool visitScopeEndsRequiringInit(
     llvm_unreachable("invalid check!?");
   }
 
+  // Look through wrappers.
+  if (auto m = dyn_cast<CopyableToMoveOnlyWrapperAddrInst>(operand)) {
+    operand = m->getOperand();
+  }
+
   // Check for inout types of arguments that are marked with consumable and
   // assignable.
   if (auto *fArg = dyn_cast<SILFunctionArgument>(operand)) {
@@ -1007,6 +1012,7 @@ void UseState::initializeLiveness(
   // Then check if our markedValue is from an argument that is in,
   // in_guaranteed, inout, or inout_aliasable, consider the marked address to be
   // the initialization point.
+  bool beginsInitialized = false;
   {
     SILValue operand = address->getOperand();
     if (auto *c = dyn_cast<CopyableToMoveOnlyWrapperAddrInst>(operand))
@@ -1025,8 +1031,7 @@ void UseState::initializeLiveness(
                "an init... adding mark_unresolved_non_copyable_value as "
                "init!\n");
         // We cheat here slightly and use our address's operand.
-        recordInitUse(address, address, liveness.getTopLevelSpan());
-        liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
+        beginsInitialized = true;
         break;
       case swift::SILArgumentConvention::Indirect_Out:
         llvm_unreachable("Should never have out addresses here");
@@ -1042,6 +1047,22 @@ void UseState::initializeLiveness(
     }
   }
 
+  // A read or write access always begins on an initialized value.
+  if (auto access = dyn_cast<BeginAccessInst>(address->getOperand())) {
+    switch (access->getAccessKind()) {
+    case SILAccessKind::Deinit:
+    case SILAccessKind::Read:
+    case SILAccessKind::Modify:
+      LLVM_DEBUG(llvm::dbgs()
+                 << "Found move only arg closure box use... "
+                    "adding mark_unresolved_non_copyable_value as init!\n");
+      beginsInitialized = true;
+      break;
+    case SILAccessKind::Init:
+      break;
+    }
+  }
+
   // See if our address is from a closure guaranteed box that we did not promote
   // to an address. In such a case, just treat our
   // mark_unresolved_non_copyable_value as the init of our value.
@@ -1057,16 +1078,14 @@ void UseState::initializeLiveness(
         LLVM_DEBUG(llvm::dbgs()
                    << "Found move only arg closure box use... "
                       "adding mark_unresolved_non_copyable_value as init!\n");
-        recordInitUse(address, address, liveness.getTopLevelSpan());
-        liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
+        beginsInitialized = true;
       }
     } else if (auto *box = dyn_cast<AllocBoxInst>(
                    lookThroughOwnershipInsts(projectBox->getOperand()))) {
       LLVM_DEBUG(llvm::dbgs()
                  << "Found move only var allocbox use... "
                     "adding mark_unresolved_non_copyable_value as init!\n");
-      recordInitUse(address, address, liveness.getTopLevelSpan());
-      liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
+      beginsInitialized = true;
     }
   }
 
@@ -1077,8 +1096,7 @@ void UseState::initializeLiveness(
     LLVM_DEBUG(llvm::dbgs()
                << "Found ref_element_addr use... "
                   "adding mark_unresolved_non_copyable_value as init!\n");
-    recordInitUse(address, address, liveness.getTopLevelSpan());
-    liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
+    beginsInitialized = true;
   }
 
   // Check if our address is from a global_addr. In such a case, we treat the
@@ -1088,8 +1106,7 @@ void UseState::initializeLiveness(
     LLVM_DEBUG(llvm::dbgs()
                << "Found global_addr use... "
                   "adding mark_unresolved_non_copyable_value as init!\n");
-    recordInitUse(address, address, liveness.getTopLevelSpan());
-    liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
+    beginsInitialized = true;
   }
 
   if (auto *ptai = dyn_cast<PointerToAddressInst>(
@@ -1098,24 +1115,21 @@ void UseState::initializeLiveness(
     LLVM_DEBUG(llvm::dbgs()
                << "Found pointer to address use... "
                   "adding mark_unresolved_non_copyable_value as init!\n");
-    recordInitUse(address, address, liveness.getTopLevelSpan());
-    liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
+    beginsInitialized = true;
   }
 
   if (auto *bai = dyn_cast_or_null<BeginApplyInst>(
         stripAccessMarkers(address->getOperand())->getDefiningInstruction())) {
     LLVM_DEBUG(llvm::dbgs()
                << "Adding accessor coroutine begin_apply as init!\n");
-    recordInitUse(address, address, liveness.getTopLevelSpan());
-    liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
+    beginsInitialized = true;
   }
 
   if (auto *eai = dyn_cast<UncheckedTakeEnumDataAddrInst>(
           stripAccessMarkers(address->getOperand()))) {
     LLVM_DEBUG(llvm::dbgs()
                << "Adding enum projection as init!\n");
-    recordInitUse(address, address, liveness.getTopLevelSpan());
-    liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
+    beginsInitialized = true;
   }
 
   // Assume a strict check of a temporary or formal access is initialized
@@ -1125,32 +1139,33 @@ void UseState::initializeLiveness(
       asi && address->isStrict()) {
     LLVM_DEBUG(llvm::dbgs()
                << "Adding strict-marked alloc_stack as init!\n");
-    recordInitUse(address, address, liveness.getTopLevelSpan());
-    liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
+    beginsInitialized = true;
   }
 
   // Assume a strict-checked value initialized before the check.
   if (address->isStrict()) {
     LLVM_DEBUG(llvm::dbgs()
                << "Adding strict marker as init!\n");
-    recordInitUse(address, address, liveness.getTopLevelSpan());
-    liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
+    beginsInitialized = true;
   }
 
   // Assume a value whose deinit has been dropped has been initialized.
   if (auto *ddi = dyn_cast<DropDeinitInst>(address->getOperand())) {
     LLVM_DEBUG(llvm::dbgs()
                << "Adding copyable_to_move_only_wrapper as init!\n");
-    recordInitUse(address, address, liveness.getTopLevelSpan());
-    liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
+    beginsInitialized = true;
   }
 
   // Assume a value wrapped in a MoveOnlyWrapper is initialized.
   if (auto *m2c = dyn_cast<CopyableToMoveOnlyWrapperAddrInst>(address->getOperand())) {
     LLVM_DEBUG(llvm::dbgs()
                << "Adding copyable_to_move_only_wrapper as init!\n");
+    beginsInitialized = true;
+  }
+  
+  if (beginsInitialized) {
     recordInitUse(address, address, liveness.getTopLevelSpan());
-    liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());    
+    liveness.initializeDef(SILValue(address), liveness.getTopLevelSpan());
   }
 
   // Now that we have finished initialization of defs, change our multi-maps

test/SILOptimizer/moveonly_copyable_wrapper_capture.swift

@@ -0,0 +1,40 @@
+// RUN: %target-swift-frontend -emit-sil -verify %s
+
+class Class {}
+//struct Class : ~Copyable {}
+
+func consume(_: consuming Class) {}
+func nonescapingClosure(_ body: () -> ()) {
+    body()
+}
+
+func testNonescapingCaptureConsuming(x: consuming Class) { // expected-error{{}}
+    nonescapingClosure { consume(x) } // expected-note{{consumed here}}
+}
+
+// TODO: `defer` should be allowed to consume local bindings
+func testDeferCaptureConsuming(x: consuming Class) { // expected-error{{}}
+    defer { consume(x) } // expected-note{{consumed here}}
+    do {}
+}
+
+func testLocalFunctionCaptureConsuming(x: consuming Class) {
+    func local() {
+        consume(x) // expected-error{{cannot be consumed when captured by an escaping closure}}
+    }
+}
+
+func testNonescapingCaptureBorrowing(x: borrowing Class) { // expected-error{{}}
+    nonescapingClosure { consume(x) } // expected-note{{consumed here}}
+}
+
+func testDeferCaptureBorrowing(x: borrowing Class) { // expected-error{{}}
+    defer { consume(x) } // expected-note{{consumed here}}
+    do {}
+}
+
+func testLocalFunctionCaptureBorrowing(x: borrowing Class) { // expected-error{{borrowed and cannot be consumed}}
+    func local() {
+        consume(x) // expected-note{{consumed here}}
+    }
+}