Merge pull request #74309 from nate-chandler/rdar129622373

[ClosureSpecializer] Bail on noncopyable argument.

Author: nate-chandler

lib/IRGen/IRGenSIL.cpp

@@ -5109,6 +5109,7 @@ void IRGenSILFunction::visitCondBranchInst(swift::CondBranchInst *i) {
 }
 
 void IRGenSILFunction::visitRetainValueInst(swift::RetainValueInst *i) {
+  assert(!i->getOperand()->getType().isMoveOnly());
   Explosion in = getLoweredExplosion(i->getOperand());
   Explosion out;
   cast<LoadableTypeInfo>(getTypeInfo(i->getOperand()->getType()))
@@ -5119,6 +5120,7 @@ void IRGenSILFunction::visitRetainValueInst(swift::RetainValueInst *i) {
 
 void IRGenSILFunction::visitRetainValueAddrInst(swift::RetainValueAddrInst *i) {
   SILValue operandValue = i->getOperand();
+  assert(!operandValue->getType().isMoveOnly());
   Address addr = getLoweredAddress(operandValue);
   SILType addrTy = operandValue->getType();
   SILType objectT = addrTy.getObjectType();

lib/SILOptimizer/IPO/ClosureSpecializer.cpp

@@ -675,22 +675,30 @@ static bool isSupportedClosure(const SILInstruction *Closure) {
     return false;
 
   if (auto *PAI = dyn_cast<PartialApplyInst>(Closure)) {
-    // Bail if any of the arguments are passed by address and
-    // are not @inout.
-    // This is a temporary limitation.
+    // Check whether each argument is supported.
     auto ClosureCallee = FRI->getReferencedFunction();
     auto ClosureCalleeConv = ClosureCallee->getConventions();
-    unsigned ClosureArgIdx =
+    unsigned ClosureArgIdxBase =
         ClosureCalleeConv.getNumSILArguments() - PAI->getNumArguments();
-    for (auto Arg : PAI->getArguments()) {
+    for (auto pair : llvm::enumerate(PAI->getArguments())) {
+      auto Arg = pair.value();
+      auto ClosureArgIdx = pair.index() + ClosureArgIdxBase;
+      auto ArgConvention =
+          ClosureCalleeConv.getSILArgumentConvention(ClosureArgIdx);
+
       SILType ArgTy = Arg->getType();
+      // Specializing (currently) always produces a retain in the caller.
+      // That's not allowed for values of move-only type.
+      if (ArgTy.isMoveOnly()) {
+        return false;
+      }
+
+      // Only @inout/@inout_aliasable addresses are (currently) supported.
       // If our argument is an object, continue...
       if (ArgTy.isObject()) {
         ++ClosureArgIdx;
         continue;
       }
-      auto ArgConvention =
-          ClosureCalleeConv.getSILArgumentConvention(ClosureArgIdx);
       if (ArgConvention != SILArgumentConvention::Indirect_Inout &&
           ArgConvention != SILArgumentConvention::Indirect_InoutAliasable)
         return false;
@@ -1394,7 +1402,7 @@ bool SILClosureSpecializerTransform::gatherCallSites(
         //   foo({ c() })
         // }
         //
-        // A limit of 2 is good enough and will not be exceed in "regular"
+        // A limit of 2 is good enough and will not be exceeded in "regular"
         // optimization scenarios.
         if (getSpecializationLevel(getClosureCallee(ClosureInst))
             > SpecializationLevelLimit) {

test/SILOptimizer/closure_specialize.sil

@@ -938,3 +938,35 @@ bb0(%0 : $Int):
   %empty = tuple ()
   return %empty : $()
 }
+
+struct NC : ~Copyable {
+  deinit {}
+}
+
+sil hidden [noinline] @noncopyable_arg_closure : $@convention(thin) (@guaranteed NC) -> () {
+bb0(%0 : $NC):
+  %retval = tuple ()
+  return %retval : $()
+}
+
+sil hidden [noinline] @use_noncopyable_arg_closure : $@convention(thin) (@noescape @callee_guaranteed () -> ()) -> () {
+bb0(%0 : $@noescape @callee_guaranteed () -> ()):
+  %2 = apply %0() : $@noescape @callee_guaranteed () -> ()
+  %3 = tuple ()
+  return %3 : $()
+}
+
+// Ensure that a retain_value of a noncopyable value isn't created.
+// CHECK-LABEL: sil @dont_specialize_noncopyable_arg_closure : {{.*}} {
+// CHECK-NOT:     retain_value {{%.*}} : $NC
+// CHECK-LABEL: } // end sil function 'dont_specialize_noncopyable_arg_closure'
+sil @dont_specialize_noncopyable_arg_closure : $@convention(thin) (@guaranteed NC) -> () {
+bb0(%nc : $NC):
+  %closure_fn = function_ref @noncopyable_arg_closure : $@convention(thin) (@guaranteed NC) -> ()
+  %closure = partial_apply [callee_guaranteed] [on_stack] %closure_fn(%nc) : $@convention(thin) (@guaranteed NC) -> ()
+  %use = function_ref @use_noncopyable_arg_closure : $@convention(thin) (@noescape @callee_guaranteed () -> ()) -> ()
+  apply %use(%closure) : $@convention(thin) (@noescape @callee_guaranteed () -> ()) -> ()
+  dealloc_stack %closure : $@noescape @callee_guaranteed () -> ()
+  %11 = tuple ()
+  return %11 : $()
+}

validation-test/SILOptimizer/issue-71495.swift

@@ -0,0 +1,48 @@
+// RUN: %target-run-simple-swift(-Xfrontend -sil-verify-all)
+// RUN: %target-run-simple-swift(-O -Xfrontend -sil-verify-all)
+
+// REQUIRES: executable_test
+
+/// A unique value represented by a heap memory location.
+struct Handle: ~Copyable {
+  var address: UnsafeMutableRawPointer
+
+  init() {
+    self.address = .allocate(byteCount: 2, alignment: 2)
+  }
+
+  consuming func done() {
+    let address = self.address
+    discard self
+    address.deallocate()
+    print("deallocated handle via done()")
+  }
+
+  deinit {
+    address.deallocate()
+    print("deallocated handle via deinit")
+  }
+}
+
+func description(of pointer: UnsafeRawPointer) -> String {
+  let address = UInt(bitPattern: pointer)
+  return "0x" + String(address, radix: 16)
+}
+
+func borrowHandleAndCrashNoMore(_ handle: borrowing Handle) -> String {
+  var string = ""
+  return string.withUTF8 { _ in
+    description(of: handle.address)
+  }
+}
+
+let handleDescription: String
+
+do {
+  let handle = Handle()
+  handleDescription = borrowHandleAndCrashNoMore(handle)
+  handle.done()
+}
+
+// CHECK: result: 0x
+print("result: \(handleDescription)")