Enable ClosureLifetimeFixup's capture copy elimination for copyable types as well.

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -41,6 +41,12 @@ llvm::cl::opt<bool> DisableConvertEscapeToNoEscapeSwitchEnumPeephole(
         "Disable the convert_escape_to_noescape switch enum peephole. "),
     llvm::cl::Hidden);
 
+llvm::cl::opt<bool> DisableCopyEliminationOfCopyableCapture(
+    "sil-disable-copy-elimination-of-copyable-closure-capture",
+    llvm::cl::init(false),
+    llvm::cl::desc("Don't eliminate copy_addr of Copyable closure captures "
+                   "inserted by SILGen"));
+
 using namespace swift;
 
 /// Given an optional diamond, return the bottom of the diamond.
@@ -734,7 +740,7 @@ static SILValue tryRewriteToPartialApplyStack(
   SmallVector<SILInstruction *, 4> lifetimeEnds;
   collectStackClosureLifetimeEnds(lifetimeEnds, closureOp);
 
-  // For noncopyable address-only captures, see if we can eliminate the copy
+  // For address-only captures, see if we can eliminate the copy
   // that SILGen emitted to allow the original partial_apply to take ownership.
   // We do this here because otherwise the move checker will see the copy as an
   // attempt to consume the value, which we don't want.
@@ -760,14 +766,14 @@ static SILValue tryRewriteToPartialApplyStack(
       LLVM_DEBUG(llvm::dbgs() << "-- not an alloc_stack\n");
       continue;
     }
-    
-    // This would be a nice optimization to attempt for all types, but for now,
-    // limit the effect to move-only types.
-    if (!copy->getType().isMoveOnly()) {
-      LLVM_DEBUG(llvm::dbgs() << "-- not move-only\n");
-      continue;
+
+    if (DisableCopyEliminationOfCopyableCapture) {
+      if (!copy->getType().isMoveOnly()) {
+        LLVM_DEBUG(llvm::dbgs() << "-- not move-only\n");
+        continue;
+      }
     }
-    
+
     // Is the capture a borrow?
 
     auto paramIndex = i + appliedArgStartIdx;
@@ -806,33 +812,33 @@ static SILValue tryRewriteToPartialApplyStack(
         if (!lookThroughMarkDependenceChainForValue(mark, newPA) ||
             mark->getBase() != stack) {
           LLVM_DEBUG(llvm::dbgs() << "-- had unexpected mark_dependence use\n";
-                     use->getUser()->print(llvm::dbgs());
-                     llvm::dbgs() << "\n");
+                     use->getUser()->print(llvm::dbgs()); llvm::dbgs() << "\n");
           initialization = nullptr;
           break;
         }
         markDep = mark;
 
         continue;
       }
-      
+
       // If we saw more than just the initialization, this isn't a pattern we
       // recognize.
       if (initialization) {
-        LLVM_DEBUG(llvm::dbgs() << "-- had non-initialization, non-partial-apply use\n";
-                   use->getUser()->print(llvm::dbgs());
-                   llvm::dbgs() << "\n");
-                   
+        LLVM_DEBUG(llvm::dbgs()
+                       << "-- had non-initialization, non-partial-apply use\n";
+                   use->getUser()->print(llvm::dbgs()); llvm::dbgs() << "\n");
+
         initialization = nullptr;
         break;
       }
       if (auto possibleInit = dyn_cast<CopyAddrInst>(use->getUser())) {
         // Should copy the source and initialize the destination.
-        if (possibleInit->isTakeOfSrc()
-            || !possibleInit->isInitializationOfDest()) {
-          LLVM_DEBUG(llvm::dbgs() << "-- had non-initialization, non-partial-apply use\n";
-                     use->getUser()->print(llvm::dbgs());
-                     llvm::dbgs() << "\n");
+        if (possibleInit->isTakeOfSrc() ||
+            !possibleInit->isInitializationOfDest()) {
+          LLVM_DEBUG(
+              llvm::dbgs()
+                  << "-- had non-initialization, non-partial-apply use\n";
+              use->getUser()->print(llvm::dbgs()); llvm::dbgs() << "\n");
 
           break;
         }
@@ -860,33 +866,36 @@ static SILValue tryRewriteToPartialApplyStack(
     LLVM_DEBUG(llvm::dbgs() << "++ found original:\n";
                orig->print(llvm::dbgs());
                llvm::dbgs() << "\n");
-
-    bool origIsUnusedDuringClosureLifetime = true;
 
-    class OrigUnusedDuringClosureLifetimeWalker
+    bool origIsUnmodifiedDuringClosureLifetime = true;
+
+    class OrigUnmodifiedDuringClosureLifetimeWalker
         : public TransitiveAddressWalker<
-              OrigUnusedDuringClosureLifetimeWalker> {
+              OrigUnmodifiedDuringClosureLifetimeWalker> {
       SSAPrunedLiveness &closureLiveness;
-      bool &origIsUnusedDuringClosureLifetime;
+      bool &origIsUnmodifiedDuringClosureLifetime;
+
     public:
-      OrigUnusedDuringClosureLifetimeWalker(SSAPrunedLiveness &closureLiveness,
-                                        bool &origIsUnusedDuringClosureLifetime)
-        : closureLiveness(closureLiveness),
-          origIsUnusedDuringClosureLifetime(origIsUnusedDuringClosureLifetime)
-      {}
+      OrigUnmodifiedDuringClosureLifetimeWalker(
+          SSAPrunedLiveness &closureLiveness,
+          bool &origIsUnmodifiedDuringClosureLifetime)
+          : closureLiveness(closureLiveness),
+            origIsUnmodifiedDuringClosureLifetime(
+                origIsUnmodifiedDuringClosureLifetime) {}
 
       bool visitUse(Operand *origUse) {
         LLVM_DEBUG(llvm::dbgs() << "looking at use\n";
                    origUse->getUser()->printInContext(llvm::dbgs());
                    llvm::dbgs() << "\n");
-        
+
         // If the user doesn't write to memory, then it's harmless.
         if (!origUse->getUser()->mayWriteToMemory()) {
           return true;
         }
         if (closureLiveness.isWithinBoundary(origUse->getUser())) {
-          origIsUnusedDuringClosureLifetime = false;
-          LLVM_DEBUG(llvm::dbgs() << "-- original has other possibly writing use during closure lifetime\n";
+          origIsUnmodifiedDuringClosureLifetime = false;
+          LLVM_DEBUG(llvm::dbgs() << "-- original has other possibly writing "
+                                     "use during closure lifetime\n";
                      origUse->getUser()->print(llvm::dbgs());
                      llvm::dbgs() << "\n");
           return false;
@@ -895,12 +904,12 @@ static SILValue tryRewriteToPartialApplyStack(
       }
     };
 
-    OrigUnusedDuringClosureLifetimeWalker origUseWalker(closureLiveness,
-                                             origIsUnusedDuringClosureLifetime);
+    OrigUnmodifiedDuringClosureLifetimeWalker origUseWalker(
+        closureLiveness, origIsUnmodifiedDuringClosureLifetime);
     auto walkResult = std::move(origUseWalker).walk(orig);
-    
-    if (walkResult == AddressUseKind::Unknown
-        || !origIsUnusedDuringClosureLifetime) {
+
+    if (walkResult == AddressUseKind::Unknown ||
+        !origIsUnmodifiedDuringClosureLifetime) {
       continue;
     }
 

test/SILOptimizer/closure-lifetime-fixup-debuginfo.sil

@@ -1,4 +1,4 @@
-// RUN: %target-sil-opt -enable-sil-verify-all -closure-lifetime-fixup -emit-verbose-sil %s | %FileCheck %s
+// RUN: %target-sil-opt -enable-sil-verify-all -closure-lifetime-fixup -emit-verbose-sil -sil-disable-copy-elimination-of-copyable-closure-capture=true %s | %FileCheck %s
 
 import Builtin
 import Swift

test/SILOptimizer/closure-lifetime-fixup.sil

@@ -267,22 +267,16 @@ bb3:
 
 // CHECK-LABEL: sil [ossa] @test_alloc_stack_arg_with_frontier_outside_try_apply_successors : $@convention(thin) <Self> (@in_guaranteed Self) -> @error any Error {
 // CHECK:       {{bb[0-9]+}}([[INSTANCE:%[^,]+]] :
-// CHECK:         [[STACK:%[^,]+]] = alloc_stack $Self
 // CHECK:         try_apply undef() {{.*}}, normal [[SUCCESS_1:bb[0-9]+]], error [[FAILURE_1:bb[0-9]+]]
 // CHECK:       [[SUCCESS_1]]
-// CHECK:         copy_addr [[INSTANCE]] to [init] [[STACK]] : $*Self, {{.*}} scope [[STACK_SCOPE:[0-9]+]]
-// CHECK:         [[CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] undef<Self>([[STACK]])
-// CHECK:         [[DEPENDENCY:%[^,]+]] = mark_dependence [nonescaping] [[CLOSURE]] {{.*}} on [[STACK]]
+// CHECK:         [[CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] undef<Self>([[INSTANCE]])
+// CHECK:         [[DEPENDENCY:%[^,]+]] = mark_dependence [nonescaping] [[CLOSURE]] {{.*}} on [[INSTANCE]]
 // CHECK:         try_apply undef([[DEPENDENCY]]) {{.*}}, normal [[SUCCESS_2:bb[0-9]+]], error [[FAILURE_2:bb[0-9]+]]
 // CHECK:       [[SUCCESS_2]]
 // CHECK:         destroy_value [[DEPENDENCY]]
-// CHECK:         destroy_addr [[STACK]]
-// CHECK:         dealloc_stack [[STACK]] : $*Self, {{.*}} scope [[STACK_SCOPE]]
 // CHECK:       [[FAILURE_1]]
-// CHECK:         dealloc_stack [[STACK]] : $*Self, {{.*}} scope [[STACK_SCOPE]]
 // CHECK:         throw
 // CHECK:       [[FAILURE_2]]
-// CHECK-NOT:     dealloc_stack
 // CHECK:         unreachable
 // CHECK-LABEL: } // end sil function 'test_alloc_stack_arg_with_frontier_outside_try_apply_successors'
 sil [ossa] @test_alloc_stack_arg_with_frontier_outside_try_apply_successors : $@convention(thin) <Self> (@in_guaranteed Self) -> (@error any Error) {

test/SILOptimizer/closure_lifetime_fixup_copyelim.swift

@@ -0,0 +1,34 @@
+// RUN: %target-swift-frontend -emit-sil %s -O| %FileCheck %s
+
+public enum Foo {
+  case bar
+  case baz(Int)
+}
+
+public protocol Frob {
+  func foo(_ x: Int) -> Foo
+}
+
+public struct Nicate {
+  public var frob: any Frob
+  public var isInitializing: Bool
+
+  // CHECK-LABEL: sil @$s31closure_lifetime_fixup_copyelim6NicateV4slowyS2iF :
+  // CHECK-NOT: copy_addr 
+  // CHECK-LABEL: } // end sil function '$s31closure_lifetime_fixup_copyelim6NicateV4slowyS2iF'
+  public func slow(_ x: Int) -> Int {
+    let foo = frob.foo(x)
+    switch foo {
+    case .bar:
+      return 10
+    case .baz(let y):
+      if y == 0 && isInitializing {
+        return foos[x]
+      } else {
+        return y
+      }
+    }
+  }
+}
+
+private let foos = [1, 2, 3]

test/SILOptimizer/moveonly_closure_lifetime_fixup_address_only_borrowed_captures.swift

@@ -8,6 +8,7 @@ struct E<T>: ~Copyable {
     var t: T
 }
 
+class Klass {}
 struct C<T> {
   var t: T
 }
@@ -19,15 +20,31 @@ func escaping(_ x: @escaping () -> ()) {  escaper = x }
 
 func borrow<T>(_: borrowing E<T>) {}
 func borrow<T>(_: borrowing C<T>) {}
+func mutate<T>(_: inout C<T>) {}
+
+// CHECK-LABEL: sil {{.*}}14testCopySafetyyyAA1CVyxGlF : 
+// CHECK: [[STK1:%.*]] = alloc_stack {{.*}} $C<T>
+// CHECK: copy_addr %0 to [init] [[STK1]]
+// CHECK: partial_apply {{.*}}[on_stack] {{.*}}([[STK1]]) : $@convention(thin) <τ_0_0> (@inout_aliasable C<τ_0_0>) -> ()
+// CHECK: [[STK2:%.*]] = alloc_stack {{.*}} $C<T>
+// CHECK: [[BA:%.*]] = begin_access [read] [static] [[STK1]]
+// CHECK: copy_addr [[BA]] to [init] [[STK2]]
+// CHECK: end_access [[BA]]
+// CHECK: partial_apply {{.*}}[on_stack] {{.*}}([[STK2]]) : $@convention(thin) <τ_0_0> (@in_guaranteed C<τ_0_0>) -> ()
+func testCopySafety<T>(_ c: C<T>) {
+  var mutC = c
+  nonescaping({ mutate(&mutC) }, and: {[mutC] in borrow(mutC) })
+}
 
-func testMultiCapture<T>(_ e: borrowing E<T>, _ c: C<T>) {
-    // CHECK: [[C_STK:%.*]] = alloc_stack $C<T>
-    // CHECK: copy_addr %1 to [init] [[C_STK]] : $*C<T>
-    // CHECK: partial_apply {{.*}}[on_stack] {{.*}}([[C_STK]], %0) :
-    nonescaping {
-        borrow(c)
-        borrow(e)
-    }
+// CHECK-LABEL: sil {{.*}}18testVarReadCaptureyyAA1CVyxGlF :
+// CHECK: [[STK:%.*]] = alloc_stack {{.*}} $C<T>
+// CHECK: copy_addr %0 to [init] [[STK]]
+// CHECK: partial_apply {{.*}}[on_stack] {{.*}}([[STK]]) : $@convention(thin) <τ_0_0> (@inout_aliasable C<τ_0_0>) -> ()
+func testVarReadCapture<T>(_ c: C<T>) {
+  var mutC = c
+  nonescaping {
+    borrow(mutC)
+  }
 }
 
 // CHECK-LABEL: sil {{.*}}16testNeverEscaped

test/SILOptimizer/opaque_values_Onone.swift

@@ -68,12 +68,10 @@ func doit<T>(_ f: () -> T) -> T {
 }
 // CHECK-LABEL: sil hidden @duplicate1 : {{.*}} {
 // CHECK:       {{bb[0-9]+}}({{%[^,]+}} : $*Value, {{%[^,]+}} : $*Value, [[INSTANCE_ADDR_IN:%[^,]+]] : $*Value):
-// CHECK:         [[INSTANCE_ADDR:%[^,]+]] = alloc_stack $Value
 // CHECK:         [[OUTPUT_TUPLE_ADDR:%[^,]+]] = alloc_stack $(Value, Value)
 // CHECK:         [[DUPLICATE_CLOSURE:%[^,]+]] = function_ref @$s19opaque_values_Onone10duplicate15valuex_xtx_tlFx_xtyXEfU_
-// CHECK:         copy_addr [[INSTANCE_ADDR_IN]] to [init] [[INSTANCE_ADDR]]
-// CHECK:         [[DUPLICATE_INSTANCE_CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[DUPLICATE_CLOSURE]]<Value>([[INSTANCE_ADDR]])
-// CHECK:         [[DEPENDENDENCY:%[^,]+]] = mark_dependence [nonescaping] [[DUPLICATE_INSTANCE_CLOSURE]] : $@noescape @callee_guaranteed () -> @out (Value, Value) on [[INSTANCE_ADDR]] : $*Value
+// CHECK:         [[DUPLICATE_INSTANCE_CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[DUPLICATE_CLOSURE]]<Value>([[INSTANCE_ADDR_IN]])
+// CHECK:         [[DEPENDENDENCY:%[^,]+]] = mark_dependence [nonescaping] [[DUPLICATE_INSTANCE_CLOSURE]] : $@noescape @callee_guaranteed () -> @out (Value, Value) on [[INSTANCE_ADDR_IN]] : $*Value
 // CHECK:         [[CONVERTED:%[^,]+]] = convert_function [[DEPENDENDENCY]]
 // CHECK:         apply {{%[^,]+}}<(Value, Value)>([[OUTPUT_TUPLE_ADDR]], [[CONVERTED]])
 // CHECK-LABEL: } // end sil function 'duplicate1'
@@ -92,12 +90,10 @@ func duplicate1<Value>(value: Value) -> (Value, Value) {
 }
 // CHECK-LABEL: sil hidden @duplicate2 : {{.*}} {
 // CHECK:       {{bb[0-9]+}}({{%[^,]+}} : $*Value, {{%[^,]+}} : $*Value, [[INSTANCE_ADDR_IN:%[^,]+]] : $*Value):
-// CHECK:         [[INSTANCE_ADDR:%[^,]+]] = alloc_stack $Value
 // CHECK:         [[OUTPUT_TUPLE_ADDR:%[^,]+]] = alloc_stack $(one: Value, two: Value)
 // CHECK:         [[DUPLICATE_CLOSURE:%[^,]+]] = function_ref @$s19opaque_values_Onone10duplicate25valuex3one_x3twotx_tlFxAD_xAEtyXEfU_
-// CHECK:         copy_addr [[INSTANCE_ADDR_IN]] to [init] [[INSTANCE_ADDR]]
-// CHECK:         [[DUPLICATE_INSTANCE_CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[DUPLICATE_CLOSURE]]<Value>([[INSTANCE_ADDR]])
-// CHECK:         [[DEPENDENDENCY:%[^,]+]] = mark_dependence [nonescaping] [[DUPLICATE_INSTANCE_CLOSURE]] : $@noescape @callee_guaranteed () -> @out (one: Value, two: Value) on [[INSTANCE_ADDR]] : $*Value
+// CHECK:         [[DUPLICATE_INSTANCE_CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[DUPLICATE_CLOSURE]]<Value>([[INSTANCE_ADDR_IN]])
+// CHECK:         [[DEPENDENDENCY:%[^,]+]] = mark_dependence [nonescaping] [[DUPLICATE_INSTANCE_CLOSURE]] : $@noescape @callee_guaranteed () -> @out (one: Value, two: Value) on [[INSTANCE_ADDR_IN]] : $*Value
 // CHECK:         [[CONVERTED:%[^,]+]] = convert_function [[DEPENDENDENCY]]
 // CHECK:         apply {{%[^,]+}}<(one: Value, two: Value)>([[OUTPUT_TUPLE_ADDR]], [[CONVERTED]])
 // CHECK-LABEL: } // end sil function 'duplicate2'
@@ -129,10 +125,8 @@ func duplicate_with_int2<Value>(value: Value) -> ((Value, Value), Int) {
 
 // CHECK-LABEL: sil hidden @duplicate_with_int3 : {{.*}} {
 // CHECK:       {{bb[0-9]+}}({{%[^,]+}} : $*Value, {{%[^,]+}} : $*Value, {{%[^,]+}} : $*Value, {{%[^,]+}} : $*Value, [[INSTANCE_ADDR_IN:%[^,]+]] : $*Value):
-// CHECK:         [[INSTANCE_ADDR:%[^,]+]] = alloc_stack $Value
 // CHECK:         [[CLOSURE:%[^,]+]] = function_ref @$s19opaque_values_Onone19duplicate_with_int35valueSi_x_x_x_SitxttSitx_tlFSi_x_x_x_SitxttSityXEfU_
-// CHECK:         copy_addr [[INSTANCE_ADDR_IN]] to [init] [[INSTANCE_ADDR]]
-// CHECK:         [[INSTANCE_CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[CLOSURE]]<Value>([[INSTANCE_ADDR]])
+// CHECK:         [[INSTANCE_CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[CLOSURE]]<Value>([[INSTANCE_ADDR_IN]])
 // CHECK:         [[DEPENDENCY:%[^,]+]] = mark_dependence [nonescaping] [[INSTANCE_CLOSURE]]
 // CHECK:         [[CONVERTED:%[^,]+]] = convert_function [[DEPENDENCY]]
 // CHECK:         apply {{%[^,]+}}<(Int, (Value, (Value, (Value, Int), Value)), Int)>({{%[^,]+}}, [[CONVERTED]])