Enable copy elimination of captures for all types in ClosureLifetimeFixup

rdar://130026658

Author: meg-gupta

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -734,7 +734,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 +760,7 @@ 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;
-    }
-    
+
     // Is the capture a borrow?
 
     auto paramIndex = i + appliedArgStartIdx;

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

@@ -21,9 +21,7 @@ func borrow<T>(_: borrowing E<T>) {}
 func borrow<T>(_: borrowing C<T>) {}
 
 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) :
+    // CHECK: partial_apply {{.*}}[on_stack] {{.*}}(%1, %0) :
     nonescaping {
         borrow(c)
         borrow(e)

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]])