[move-only] Teach the move checker how to correctly check an address only type used by escaping closures.

When we have a loadable type, SILGen always eagerly loads the value, so the move
checker handled checking those cases by just looking for loads and emitting the
error based off of a load. Of course for address only types, the codegen looks
slightly different (we consume the value directly rather than load it). So this
change just teaches the checker how to do that.

rdar://105106470

Author: gottesmm

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -1728,9 +1728,30 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
   // Now that we have handled or loadTakeOrCopy, we need to now track our
   // additional pure takes.
   if (::memInstMustConsume(op)) {
+    // If we don't have a consumeable and assignable check kind, then we can't
+    // consume. Emit an error.
+    //
+    // NOTE: Since SILGen eagerly loads loadable types from memory, this
+    // generally will only handle address only types.
+    if (markedValue->getCheckKind() !=
+        MarkMustCheckInst::CheckKind::ConsumableAndAssignable) {
+      auto *fArg = dyn_cast<SILFunctionArgument>(
+          stripAccessMarkers(markedValue->getOperand()));
+      if (fArg && fArg->isClosureCapture() && fArg->getType().isAddress()) {
+        moveChecker.diagnosticEmitter.emitPromotedBoxArgumentError(markedValue,
+                                                                   fArg);
+      } else {
+        moveChecker.diagnosticEmitter
+            .emitAddressEscapingClosureCaptureLoadedAndConsumed(markedValue);
+      }
+      emittedEarlyDiagnostic = true;
+      return true;
+    }
+
     auto leafRange = TypeTreeLeafTypeRange::get(op->get(), getRootAddress());
     if (!leafRange)
       return false;
+
     LLVM_DEBUG(llvm::dbgs() << "Pure consuming use: " << *user);
     useState.takeInsts.insert({user, *leafRange});
     return true;
@@ -2430,7 +2451,6 @@ bool MoveOnlyAddressCheckerPImpl::performSingleCheck(
     LLVM_DEBUG(llvm::dbgs() << "Failed access path visit: " << *markedAddress);
     return false;
   }
-  addressUseState.initializeInOutTermUsers();
 
   // If we found a load [copy] or copy_addr that requires multiple copies or an
   // exclusivity error, then we emitted an early error. Bail now and allow the
@@ -2445,9 +2465,14 @@ bool MoveOnlyAddressCheckerPImpl::performSingleCheck(
   if (diagCount != diagnosticEmitter.getDiagnosticCount())
     return true;
 
-  // Then check if we emitted an error. If we did not, return true.
-  if (diagCount != diagnosticEmitter.getDiagnosticCount())
-    return true;
+  // Now that we know that we have run our visitor and did not emit any errors
+  // and successfully visited everything, see if have any
+  // assignable_but_not_consumable of address only types that are consumed.
+  //
+  // DISCUSSION: For non address only types, this is not an issue since we
+  // eagerly load
+
+  addressUseState.initializeInOutTermUsers();
 
   //===---
   // Liveness Checking

test/SILOptimizer/moveonly_addresschecker_diagnostics.sil

@@ -59,12 +59,30 @@ public struct AggStruct {
     var pair: KlassPair
 }
 
+protocol P {
+    static var value: Self { get }
+    static var value2: any P { get }
+}
+
+@_moveOnly
+public struct AddressOnlyGeneric<T : P> {
+    var copyable: T
+    var moveOnly: NonTrivialStruct
+
+    init()
+    init(_ input1: T)
+}
+
 sil @get_aggstruct : $@convention(thin) () -> @owned AggStruct
 sil @nonConsumingUseKlass : $@convention(thin) (@guaranteed Klass) -> ()
 sil @nonConsumingUseNonTrivialStruct : $@convention(thin) (@guaranteed NonTrivialStruct) -> ()
 sil @classConsume : $@convention(thin) (@owned Klass) -> () // user: %18
 sil @copyableClassConsume : $@convention(thin) (@owned CopyableKlass) -> () // user: %24
 sil @copyableClassUseMoveOnlyWithoutEscaping : $@convention(thin) (@guaranteed CopyableKlass) -> () // user: %16
+sil @getAddressOnlyGeneric : $@convention(thin) <τ_0_0 where τ_0_0 : P> () -> @out AddressOnlyGeneric<τ_0_0>
+sil @addressOnlyGenericInOutUse : $@convention(thin) <τ_0_0 where τ_0_0 : P> (@inout AddressOnlyGeneric<τ_0_0>) -> ()
+sil @addressOnlyGenericConsume : $@convention(thin) <τ_0_0 where τ_0_0 : P> (@in AddressOnlyGeneric<τ_0_0>) -> ()
+sil @addressOnlyGenericUse : $@convention(thin) <τ_0_0 where τ_0_0 : P> (@in_guaranteed AddressOnlyGeneric<τ_0_0>) -> ()
 
 ///////////
 // Tests //
@@ -316,3 +334,34 @@ bb0(%0 : $*Klass):
   %27 = tuple ()
   return %27 : $()
 }
+
+////////////////////////
+// Address Only Tests //
+////////////////////////
+
+sil [ossa] @inoutCaptureTestAddressOnlyGenericClosure2 : $@convention(thin) <T where T : P> (@guaranteed <τ_0_0 where τ_0_0 : P> { var AddressOnlyGeneric<τ_0_0> } <T>) -> () {
+bb0(%0 : @closureCapture @guaranteed $<τ_0_0 where τ_0_0 : P> { var AddressOnlyGeneric<τ_0_0> } <T>):
+  %1 = project_box %0 : $<τ_0_0 where τ_0_0 : P> { var AddressOnlyGeneric<τ_0_0> } <T>, 0
+  debug_value %1 : $*AddressOnlyGeneric<T>, var, name "x", argno 1, expr op_deref
+  %3 = alloc_stack $AddressOnlyGeneric<T>
+  %5 = function_ref @getAddressOnlyGeneric : $@convention(thin) <τ_0_0 where τ_0_0 : P> () -> @out AddressOnlyGeneric<τ_0_0>
+  %6 = apply %5<T>(%3) : $@convention(thin) <τ_0_0 where τ_0_0 : P> () -> @out AddressOnlyGeneric<τ_0_0>
+  %7 = begin_access [modify] [dynamic] %1 : $*AddressOnlyGeneric<T>
+  %8 = mark_must_check [assignable_but_not_consumable] %7 : $*AddressOnlyGeneric<T>
+  copy_addr [take] %3 to %8 : $*AddressOnlyGeneric<T>
+  end_access %7 : $*AddressOnlyGeneric<T>
+  dealloc_stack %3 : $*AddressOnlyGeneric<T>
+  %12 = begin_access [modify] [dynamic] %1 : $*AddressOnlyGeneric<T>
+  %13 = mark_must_check [assignable_but_not_consumable] %12 : $*AddressOnlyGeneric<T>
+  %14 = function_ref @addressOnlyGenericInOutUse : $@convention(thin) <τ_0_0 where τ_0_0 : P> (@inout AddressOnlyGeneric<τ_0_0>) -> ()
+  %15 = apply %14<T>(%13) : $@convention(thin) <τ_0_0 where τ_0_0 : P> (@inout AddressOnlyGeneric<τ_0_0>) -> ()
+  end_access %12 : $*AddressOnlyGeneric<T>
+  %17 = begin_access [deinit] [dynamic] %1 : $*AddressOnlyGeneric<T>
+  %18 = mark_must_check [assignable_but_not_consumable] %17 : $*AddressOnlyGeneric<T>
+  // expected-error @-1 {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
+  %19 = function_ref @addressOnlyGenericConsume : $@convention(thin) <τ_0_0 where τ_0_0 : P> (@in AddressOnlyGeneric<τ_0_0>) -> ()
+  %20 = apply %19<T>(%18) : $@convention(thin) <τ_0_0 where τ_0_0 : P> (@in AddressOnlyGeneric<τ_0_0>) -> ()
+  end_access %17 : $*AddressOnlyGeneric<T>
+  %22 = tuple ()
+  return %22 : $()
+}

test/SILOptimizer/moveonly_addresschecker_diagnostics.swift

@@ -3609,6 +3609,29 @@ func inoutCaptureTest() -> (() -> ()) {
     return f
 }
 
+func inoutCaptureTestAddressOnlyGeneric<T : P>(_ t: T.Type) -> (() -> ()) {
+    var x = AddressOnlyGeneric<T>()
+    x = AddressOnlyGeneric<T>()
+
+    func useInOut(_ x: inout AddressOnlyGeneric<T>) {}
+    let f = {
+        useInOut(&x)
+    }
+
+    borrowVal(x)
+    consumeVal(x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
+    x = AddressOnlyGeneric<T>()
+
+    let g = {
+        x = AddressOnlyGeneric<T>()
+        useInOut(&x)
+        consumeVal(x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
+    }
+    g()
+
+    return f
+}
+
 ////////////////
 // Misc Tests //
 ////////////////