AllocBoxToStack: handle begin_borrow when checking if a partial_apply escapes.

Author: eeckstein

lib/SILOptimizer/Transforms/AllocBoxToStack.cpp

@@ -70,7 +70,8 @@ static bool useCaptured(Operand *UI) {
   // These instructions do not cause the address to escape.
   if (isa<DebugValueInst>(User)
       || isa<StrongReleaseInst>(User) || isa<StrongRetainInst>(User)
-      || isa<DestroyValueInst>(User))
+      || isa<DestroyValueInst>(User)
+      || isa<EndBorrowInst>(User))
     return false;
 
   if (auto *Store = dyn_cast<StoreInst>(User)) {
@@ -254,8 +255,9 @@ static bool partialApplyEscapes(SILValue V, bool examineApply) {
     // If we have a copy_value, the copy value does not cause an escape, but its
     // uses might do so... so add the copy_value's uses to the worklist and
     // continue.
-    if (auto CVI = dyn_cast<CopyValueInst>(User)) {
-      llvm::copy(CVI->getUses(), std::back_inserter(Worklist));
+    if (isa<CopyValueInst>(User) || isa<BeginBorrowInst>(User)) {
+      llvm::copy(cast<SingleValueInstruction>(User)->getUses(),
+                 std::back_inserter(Worklist));
       continue;
     }
 

test/IRGen/alloc_box.swift

@@ -2,13 +2,17 @@
 
 func f() -> Bool? { return nil }
 
+var gb = false
+var gc: () -> () = {}
+
 ({
   guard var b = f() else { return }
   let c = { b = true }
-  _ = (b, c)
+  gb = b
+  gc = c
 })()
 
-// CHECK-LABEL: @"$s9alloc_boxyyXEfU_"
+// CHECK-LABEL: @"$s9alloc_boxyyXEfU0_"
 // CHECK-NOT: call void @swift_setDeallocating
 // CHECK: call void @swift_deallocUninitializedObject
 // CHECK-NOT: call void @swift_setDeallocating

test/SILOptimizer/allocbox_to_stack_ownership.sil

@@ -469,6 +469,38 @@ bb0(%0 : $Int):
   unreachable
 }
 
+// CHECK-LABEL: sil [ossa] @test_copy_and_borrow_of_closure
+// CHECK-NOT:     alloc_box
+// CHECK:         alloc_stack
+// CHECK-NOT:     alloc_box
+// CHECK:       } // end sil function 'test_copy_and_borrow_of_closure'
+sil [ossa] @test_copy_and_borrow_of_closure : $@convention(thin) (Int) -> () {
+bb0(%0 : $Int):
+  // This is not destroyed, but the unreachable makes the verifier not trip.
+  %1 = alloc_box $<τ_0_0> { var τ_0_0 } <Int>
+  %1a = project_box %1 : $<τ_0_0> { var τ_0_0 } <Int>, 0
+  store %0 to [trivial] %1a : $*Int
+  %3 = function_ref @closure3 : $@convention(thin) (@owned <τ_0_0> { var τ_0_0 } <Int>) -> ()
+  %4 = copy_value %1 : $<τ_0_0> { var τ_0_0 } <Int>
+  %5 = partial_apply %3(%4) : $@convention(thin) (@owned <τ_0_0> { var τ_0_0 } <Int>) -> ()
+  %6 = begin_borrow %5 : $@callee_owned () -> ()
+  %7 = copy_value %6 : $@callee_owned () -> ()
+  apply %7() : $@callee_owned () -> ()
+  end_borrow %6 : $@callee_owned () -> ()
+  destroy_value %5 : $@callee_owned () -> ()
+  destroy_value %1 : $<τ_0_0> { var τ_0_0 } <Int>
+  %r = tuple ()
+  return %r : $()
+}
+
+sil [ossa] @closure3 : $@convention(thin) (@owned <τ_0_0> { var τ_0_0 } <Int>) -> () {
+bb0(%0 : @owned $<τ_0_0> { var τ_0_0 } <Int>):
+  %1 = project_box %0 : $<τ_0_0> { var τ_0_0 } <Int>, 0
+  destroy_value %0 : $<τ_0_0> { var τ_0_0 } <Int>      // id: %7
+  %r = tuple ()
+  return %r : $()
+}
+
 sil [ossa] @closureWithBoxArg : $@convention(thin) (@guaranteed { var SomeClass }) -> () {
 bb0(%0 : @guaranteed ${ var SomeClass }):
   %r = tuple ()

test/SILOptimizer/move_function_kills_copyable_addressonly_vars.swift

@@ -574,15 +574,25 @@ public func castAccess2<T : P>(_ x : __owned T) {
 // Partial Apply Tests //
 /////////////////////////
 
+public func nonEscapingpartialApplyTest<T : P>(_ x: __owned T) {
+    var x2 = x // expected-error {{'x2' used after being moved}}
+    x2 = x
+    let _ = _move x2 // expected-note {{move here}}
+    let f = { // expected-note {{use here}}
+        print(x2)
+    }
+    f()
+}
+
 // This makes sure we always fail if we are asked to check in a partial apply.
-public func partialApplyTest<T : P>(_ x: __owned T) {
+public func partialApplyTest<T : P>(_ x: __owned T) -> () -> () {
     var x2 = x
     x2 = x
     let _ = _move x2 // expected-error {{move applied to value that the compiler does not support checking}}
     let f = {
         print(x2)
     }
-    f()
+    return f
 }
 
 ////////////////////////