[move-only] Teach SILGen that lets captured by a closure must be no copy.

The current diagnostic that we are emitting is not perfect, but at least this
prevents us from saying that an error is not occuring here. I am going to file a
bug to track the QoI work of improving the diagnostic.

rdar://103313305

Author: gottesmm

include/swift/AST/DiagnosticsSIL.def

@@ -731,6 +731,8 @@ ERROR(sil_moveonlychecker_value_used_after_consume, none,
       "'%0' used after consume. Lifetime extension of variable requires a copy", (StringRef))
 ERROR(sil_moveonlychecker_guaranteed_value_consumed, none,
       "'%0' has guaranteed ownership but was consumed", (StringRef))
+ERROR(sil_moveonlychecker_guaranteed_value_captured_by_closure, none,
+      "'%0' has guaranteed ownership but was consumed due to being captured by a closure", (StringRef))
 ERROR(sil_moveonlychecker_inout_not_reinitialized_before_end_of_function, none,
       "'%0' consumed but not reinitialized before end of function", (StringRef))
 ERROR(sil_moveonlychecker_value_consumed_in_a_loop, none,
@@ -740,6 +742,8 @@ ERROR(sil_moveonlychecker_exclusivity_violation, none,
 
 NOTE(sil_moveonlychecker_consuming_use_here, none,
      "consuming use", ())
+NOTE(sil_moveonlychecker_consuming_closure_use_here, none,
+     "closure capture", ())
 NOTE(sil_moveonlychecker_nonconsuming_use_here, none,
      "non-consuming use", ())
 ERROR(sil_moveonlychecker_not_understand_no_implicit_copy, none,

lib/SILGen/SILGenBuilder.cpp

@@ -954,3 +954,12 @@ SILGenBuilder::createMarkMustCheckInst(SILLocation loc, ManagedValue value,
       loc, value.forward(getSILGenFunction()), kind);
   return cloner.clone(mdi);
 }
+
+ManagedValue SILGenBuilder::emitCopyValueOperation(SILLocation loc,
+                                                   ManagedValue value) {
+  auto cvi = SILBuilder::emitCopyValueOperation(loc, value.getValue());
+  // Trivial type.
+  if (cvi == value.getValue())
+    return value;
+  return SGF.emitManagedRValueWithCleanup(cvi);
+}

lib/SILGen/SILGenBuilder.h

@@ -430,6 +430,9 @@ class SILGenBuilder : public SILBuilder {
   using SILBuilder::createMarkMustCheckInst;
   ManagedValue createMarkMustCheckInst(SILLocation loc, ManagedValue value,
                                        MarkMustCheckInst::CheckKind kind);
+
+  using SILBuilder::emitCopyValueOperation;
+  ManagedValue emitCopyValueOperation(SILLocation Loc, ManagedValue v);
 };
 
 } // namespace Lowering

lib/SILGen/SILGenProlog.cpp

@@ -520,34 +520,47 @@ static void emitCaptureArguments(SILGenFunction &SGF,
     auto &lowering = SGF.getTypeLowering(type);
     // Constant decls are captured by value.
     SILType ty = lowering.getLoweredType();
-    SILValue val = SGF.F.begin()->createFunctionArgument(ty, VD);
-
-    bool NeedToDestroyValueAtExit = false;
+    ManagedValue val = ManagedValue::forUnmanaged(
+        SGF.F.begin()->createFunctionArgument(ty, VD));
 
     // If the original variable was settable, then Sema will have treated the
     // VarDecl as an lvalue, even in the closure's use.  As such, we need to
     // allow formation of the address for this captured value.  Create a
     // temporary within the closure to provide this address.
     if (VD->isSettable(VD->getDeclContext())) {
-      auto addr = SGF.emitTemporaryAllocation(VD, ty);
+      auto addr = SGF.emitTemporary(VD, lowering);
       // We have created a copy that needs to be destroyed.
       val = SGF.B.emitCopyValueOperation(Loc, val);
-      NeedToDestroyValueAtExit = true;
-      lowering.emitStore(SGF.B, VD, val, addr, StoreOwnershipQualifier::Init);
-      val = addr;
+      // We use the SILValue version of this because the SILGenBuilder version
+      // will create a cloned cleanup, which we do not want since our temporary
+      // already has a cleanup.
+      //
+      // MG: Is this the right semantics for createStore? Seems like that
+      // should be potentially a different API.
+      SGF.B.emitStoreValueOperation(VD, val.forward(SGF), addr->getAddress(),
+                                    StoreOwnershipQualifier::Init);
+      addr->finishInitialization(SGF);
+      val = addr->getManagedAddress();
+    }
+
+    // If this constant is a move only type, we need to add no_copy checking to
+    // ensure that we do not consume this captured value in the function. This
+    // is because closures can be invoked multiple times which is inconsistent
+    // with consuming the move only type.
+    if (val.getType().isMoveOnly()) {
+      val = val.ensurePlusOne(SGF, Loc);
+      val = SGF.B.createMarkMustCheckInst(Loc, val,
+                                          MarkMustCheckInst::CheckKind::NoCopy);
     }
 
-    SGF.VarLocs[VD] = SILGenFunction::VarLoc::get(val);
-    if (auto *AllocStack = dyn_cast<AllocStackInst>(val))
+    SGF.VarLocs[VD] = SILGenFunction::VarLoc::get(val.getValue());
+    if (auto *AllocStack = dyn_cast<AllocStackInst>(val.getValue())) {
       AllocStack->setArgNo(ArgNo);
-    else {
+    } else {
       SILDebugVariable DbgVar(VD->isLet(), ArgNo);
-      SGF.B.createDebugValue(Loc, val, DbgVar);
+      SGF.B.createDebugValue(Loc, val.getValue(), DbgVar);
     }
 
-    // TODO: Closure contexts should always be guaranteed.
-    if (NeedToDestroyValueAtExit && !lowering.isTrivial())
-      SGF.enterDestroyCleanup(val);
     break;
   }
 

lib/SILOptimizer/Mandatory/MoveOnlyObjectChecker.cpp

@@ -99,6 +99,26 @@ struct OSSACanonicalizer {
   bool foundConsumingUseRequiringCopy() const {
     return consumingUsesNeedingCopy.size();
   }
+
+  bool hasPartialApplyConsumingUse() const {
+    return llvm::any_of(consumingUsesNeedingCopy,
+                        [](Operand *use) {
+                          return isa<PartialApplyInst>(use->getUser());
+                        }) ||
+           llvm::any_of(finalConsumingUses, [](Operand *use) {
+             return isa<PartialApplyInst>(use->getUser());
+           });
+  }
+
+  bool hasNonPartialApplyConsumingUse() const {
+    return llvm::any_of(consumingUsesNeedingCopy,
+                        [](Operand *use) {
+                          return !isa<PartialApplyInst>(use->getUser());
+                        }) ||
+           llvm::any_of(finalConsumingUses, [](Operand *use) {
+             return !isa<PartialApplyInst>(use->getUser());
+           });
+  }
 };
 
 } // anonymous namespace
@@ -147,7 +167,12 @@ struct DiagnosticEmitter {
   }
 
 private:
-  void emitDiagnosticsForFoundUses() const;
+  /// Emit diagnostics for the final consuming uses and consuming uses needing
+  /// copy. If filter is non-null, allow for the caller to pre-process operands
+  /// and emit their own diagnostic. If filter returns true, then we assume that
+  /// the caller processed it correctly. false, then we continue to process it.
+  void emitDiagnosticsForFoundUses(bool ignorePartialApply = false) const;
+  void emitDiagnosticsForPartialApplyUses() const;
 };
 
 } // anonymous namespace
@@ -172,11 +197,24 @@ void DiagnosticEmitter::emitGuaranteedDiagnostic(
   auto &astContext = fn->getASTContext();
   StringRef varName = getVariableNameForValue(markedValue);
 
+  // See if we have any closure capture uses and emit a better diagnostic.
+  if (canonicalizer.hasPartialApplyConsumingUse()) {
+    diagnose(astContext,
+             markedValue->getDefiningInstruction()->getLoc().getSourceLoc(),
+             diag::sil_moveonlychecker_guaranteed_value_captured_by_closure,
+             varName);
+    emitDiagnosticsForPartialApplyUses();
+    valuesWithDiagnostics.insert(markedValue);
+  }
+
+  // If we do not have any non-partial apply consuming uses... just exit early.
+  if (!canonicalizer.hasNonPartialApplyConsumingUse())
+    return;
+
   diagnose(astContext,
            markedValue->getDefiningInstruction()->getLoc().getSourceLoc(),
            diag::sil_moveonlychecker_guaranteed_value_consumed, varName);
-
-  emitDiagnosticsForFoundUses();
+  emitDiagnosticsForFoundUses(true /*ignore partial apply uses*/);
   valuesWithDiagnostics.insert(markedValue);
 }
 
@@ -193,7 +231,8 @@ void DiagnosticEmitter::emitOwnedDiagnostic(MarkMustCheckInst *markedValue) {
   valuesWithDiagnostics.insert(markedValue);
 }
 
-void DiagnosticEmitter::emitDiagnosticsForFoundUses() const {
+void DiagnosticEmitter::emitDiagnosticsForFoundUses(
+    bool ignorePartialApplyUses) const {
   auto &astContext = fn->getASTContext();
 
   for (auto *consumingUse : canonicalizer.consumingUsesNeedingCopy) {
@@ -209,6 +248,9 @@ void DiagnosticEmitter::emitDiagnosticsForFoundUses() const {
       }
     }
 
+    if (ignorePartialApplyUses &&
+        isa<PartialApplyInst>(consumingUse->getUser()))
+      continue;
     diagnose(astContext, loc.getSourceLoc(),
              diag::sil_moveonlychecker_consuming_use_here);
   }
@@ -226,11 +268,58 @@ void DiagnosticEmitter::emitDiagnosticsForFoundUses() const {
       }
     }
 
+    if (ignorePartialApplyUses &&
+        isa<PartialApplyInst>(consumingUse->getUser()))
+      continue;
+
     diagnose(astContext, loc.getSourceLoc(),
              diag::sil_moveonlychecker_consuming_use_here);
   }
 }
 
+void DiagnosticEmitter::emitDiagnosticsForPartialApplyUses() const {
+  auto &astContext = fn->getASTContext();
+
+  for (auto *consumingUse : canonicalizer.consumingUsesNeedingCopy) {
+    // See if the consuming use is an owned moveonly_to_copyable whose only
+    // user is a return. In that case, use the return loc instead. We do this
+    // b/c it is illegal to put a return value location on a non-return value
+    // instruction... so we have to hack around this slightly.
+    auto *user = consumingUse->getUser();
+    auto loc = user->getLoc();
+    if (auto *mtc = dyn_cast<MoveOnlyWrapperToCopyableValueInst>(user)) {
+      if (auto *ri = mtc->getSingleUserOfType<ReturnInst>()) {
+        loc = ri->getLoc();
+      }
+    }
+
+    if (!isa<PartialApplyInst>(consumingUse->getUser()))
+      continue;
+    diagnose(astContext, loc.getSourceLoc(),
+             diag::sil_moveonlychecker_consuming_closure_use_here);
+  }
+
+  for (auto *consumingUse : canonicalizer.finalConsumingUses) {
+    // See if the consuming use is an owned moveonly_to_copyable whose only
+    // user is a return. In that case, use the return loc instead. We do this
+    // b/c it is illegal to put a return value location on a non-return value
+    // instruction... so we have to hack around this slightly.
+    auto *user = consumingUse->getUser();
+    auto loc = user->getLoc();
+    if (auto *mtc = dyn_cast<MoveOnlyWrapperToCopyableValueInst>(user)) {
+      if (auto *ri = mtc->getSingleUserOfType<ReturnInst>()) {
+        loc = ri->getLoc();
+      }
+    }
+
+    if (!isa<PartialApplyInst>(consumingUse->getUser()))
+      continue;
+
+    diagnose(astContext, loc.getSourceLoc(),
+             diag::sil_moveonlychecker_consuming_closure_use_here);
+  }
+}
+
 //===----------------------------------------------------------------------===//
 //                              MARK: Main Pass
 //===----------------------------------------------------------------------===//

test/SILGen/closures.swift

@@ -845,3 +845,22 @@ func closeOverStructDontCopyTrivial() {
   a = StructWithMutatingMethod()
   takeClosure { a.x }
 }
+
+// Make sure that we handle closure argument initialization when due to forward
+// declaration we represent a let as an lvalue
+func test() {
+    let k: SomeClass
+    let k2: SomeClass
+    var boolean: Bool { true }
+
+    if boolean {
+        k = SomeClass()
+        k2 = SomeClass()
+    } else {
+        k = SomeClass()
+        k2 = SomeClass()
+    }
+
+    assert(k.x == k2.x)
+}
+