Merge pull request #67218 from nate-chandler/rdar112078094

[AllocBoxToStack] Computing partial_apply frontier looks thru copies+moves.

Author: nate-chandler

lib/SILOptimizer/Transforms/AllocBoxToStack.cpp

@@ -20,6 +20,7 @@
 #include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SIL/Dominance.h"
 #include "swift/SIL/MemAccessUtils.h"
+#include "swift/SIL/NodeDatastructures.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILCloner.h"
@@ -1101,7 +1102,31 @@ specializeApplySite(SILOptFunctionBuilder &FuncBuilder, ApplySite Apply,
       // borrows its captures, and we don't need to adjust capture lifetimes.
       if (!PAI->isOnStack()) {
         if (PAFrontier.empty()) {
-          ValueLifetimeAnalysis VLA(PAI, PAI->getUses());
+          SmallVector<SILInstruction *, 8> users;
+          InstructionWorklist worklist(PAI->getFunction());
+          worklist.push(PAI);
+          while (auto *inst = worklist.pop()) {
+            auto *svi = cast<SingleValueInstruction>(inst);
+            for (auto *use : svi->getUses()) {
+              auto *user = use->getUser();
+              SingleValueInstruction *svi;
+              // A copy_value produces a value with a new lifetime on which the
+              // captured alloc_box's lifetime depends.  If the transformation
+              // were only to create a destroy_value of the alloc_box (and to
+              // rewrite the closure not to consume it), the alloc_box would be
+              // kept alive by the copy_value.  The transformation does more,
+              // however: it rewrites the alloc_box as an alloc_stack, creating
+              // the alloc_stack/dealloc_stack instructions where the alloc_box/
+              // destroy_value instructions are respectively.  The copy_value
+              // can't keep the alloc_stack alive.
+              if ((svi = dyn_cast<CopyValueInst>(user)) ||
+                  (svi = dyn_cast<MoveValueInst>(user))) {
+                worklist.push(svi);
+              }
+              users.push_back(user);
+            }
+          }
+          ValueLifetimeAnalysis VLA(PAI, users);
           pass.CFGChanged |= !VLA.computeFrontier(
               PAFrontier, ValueLifetimeAnalysis::AllowToModifyCFG);
           assert(!PAFrontier.empty() &&

test/SILOptimizer/allocboxtostack_localapply_ossa.sil

@@ -488,3 +488,63 @@ bb0(%0 : @guaranteed ${ var Int }, %1 : @guaranteed ${ var Int }):
   %8 = apply %7(%0, %1) : $@convention(thin) (@guaranteed { var Int }, @guaranteed { var Int }) -> Int
   return %8 : $Int
 }
+
+class C {}
+
+sil @getC : $@convention(thin) () -> (@owned C)
+
+sil [ossa] @borrow_c_box : $@convention(thin) (@guaranteed { var C }) -> () {
+entry(%box : @guaranteed ${var C}):
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_copy_applied : {{.*}} {
+// CHECK:         [[CLOSURE:%[^,]+]] = partial_apply
+// CHECK:         [[COPY:%[^,]+]] = copy_value [[CLOSURE]]
+// CHECK:         destroy_value [[CLOSURE]]
+// CHECK:         apply [[COPY]]()
+// CHECK:         destroy_addr
+// CHECK:         dealloc_stack
+// CHECK-LABEL: } // end sil function 'test_copy_applied'
+sil [ossa] @test_copy_applied : $@convention(thin) () -> () {
+bb0:
+  %box = alloc_box ${ var C }, var, name "x"
+  %addr = project_box %box : ${ var C }, 0
+  %getC = function_ref @getC : $@convention(thin) () -> (@owned C)
+  %c = apply %getC() : $@convention(thin) () -> (@owned C)
+  store %c to [init] %addr : $*C
+
+  %borrow_int_box = function_ref @borrow_c_box : $@convention(thin) (@guaranteed { var C }) -> ()
+  %closure = partial_apply [callee_guaranteed] %borrow_int_box(%box) : $@convention(thin) (@guaranteed { var C }) -> ()
+  %copy = copy_value %closure : $@callee_guaranteed () -> ()
+  destroy_value %closure : $@callee_guaranteed () -> ()
+  apply %copy() : $@callee_guaranteed () -> ()
+  destroy_value %copy : $@callee_guaranteed () -> ()
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_move_applied : {{.*}} {
+// CHECK:         [[CLOSURE:%[^,]+]] = partial_apply
+// CHECK:         [[MOVE:%[^,]+]] = move_value [[CLOSURE]]
+// CHECK:         apply [[MOVE]]()
+// CHECK:         destroy_addr
+// CHECK:         dealloc_stack
+// CHECK-LABEL: } // end sil function 'test_move_applied'
+sil [ossa] @test_move_applied : $@convention(thin) () -> () {
+bb0:
+  %box = alloc_box ${ var C }, var, name "x"
+  %addr = project_box %box : ${ var C }, 0
+  %getC = function_ref @getC : $@convention(thin) () -> (@owned C)
+  %c = apply %getC() : $@convention(thin) () -> (@owned C)
+  store %c to [init] %addr : $*C
+
+  %borrow_int_box = function_ref @borrow_c_box : $@convention(thin) (@guaranteed { var C }) -> ()
+  %closure = partial_apply [callee_guaranteed] %borrow_int_box(%box) : $@convention(thin) (@guaranteed { var C }) -> ()
+  %move = move_value %closure : $@callee_guaranteed () -> ()
+  apply %move() : $@callee_guaranteed () -> ()
+  destroy_value %move : $@callee_guaranteed () -> ()
+  %retval = tuple ()
+  return %retval : $()
+}