[move-only] Ensure that we treat captured escaping closure arguments as such even if the closure doesn't actually escape

Specifically, we already have the appropriate semantics for arguments captured
by escaping closures but in certain cases allocbox to stack is able to prove
that the closure doesn’t actually escape. This results in the capture being
converted into a non-escaping SIL form. This then causes the move checker to
emit the wrong kind of error.

The solution is to create an early allocbox to stack that doesn’t promote move
only types in boxes from heap -> stack if it is captured by an escaping closure
but does everything else normally. Then once the move checking is completed, we
run alloc box to stack an additional time to ensure that we keep the guarantee
that heap -> stack is performed in those cases.

rdar://108905586
(cherry picked from commit 224674c)

Author: gottesmm

include/swift/SILOptimizer/PassManager/Passes.def

@@ -112,6 +112,9 @@ PASS(AccessMarkerElimination, "access-marker-elim",
      "Access Marker Elimination.")
 PASS(AddressLowering, "address-lowering",
      "SIL Address Lowering")
+PASS(EarlyAllocBoxToStack, "early-allocbox-to-stack",
+     "Stack Promotion of Box Objects. Doesn't attempt to promote noncopyable "
+     "types captured by escaping closures")
 PASS(AllocBoxToStack, "allocbox-to-stack",
      "Stack Promotion of Box Objects")
 IRGEN_PASS(AllocStackHoisting, "alloc-stack-hoisting",

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -125,7 +125,7 @@ static void addMandatoryDiagnosticOptPipeline(SILPassPipelinePlan &P) {
   // This guarantees that stack-promotable boxes have [static] enforcement.
   P.addAccessEnforcementSelection();
 
-  P.addAllocBoxToStack();
+  P.addEarlyAllocBoxToStack();
   P.addNoReturnFolding();
   addDefiniteInitialization(P);
 
@@ -169,6 +169,11 @@ static void addMandatoryDiagnosticOptPipeline(SILPassPipelinePlan &P) {
   // End Ownership Optimizations
   //===---
 
+  // Now that we have finished checking noncopyable types, run later alloc box
+  // to stack, so that we promote to the stack any heap boxes that are captured
+  // by escaping closures where the closure does not actually escape.
+  P.addAllocBoxToStack();
+
 #ifndef NDEBUG
   // Add a verification pass to check our work when skipping
   // function bodies.

lib/SILOptimizer/Transforms/AllocBoxToStack.cpp

@@ -16,6 +16,7 @@
 #include "swift/AST/SemanticAttrs.h"
 #include "swift/Basic/BlotMapVector.h"
 #include "swift/Basic/GraphNodeWorklist.h"
+#include "swift/Basic/OptionSet.h"
 #include "swift/SIL/ApplySite.h"
 #include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SIL/Dominance.h"
@@ -87,6 +88,18 @@ static bool useCaptured(Operand *UI) {
   return true;
 }
 
+namespace {
+
+enum class AllocBoxToStackFlags {
+  InAppliedFunction = 0x1,
+  CanPromoteNonCopyableCapturedByEscapingClosure = 0x2,
+  HasMoveOnlyBox = 0x4,
+};
+
+using AllocBoxToStackOptions = OptionSet<AllocBoxToStackFlags>;
+
+} // namespace
+
 //===----------------------------------------------------------------------===//
 //                 Liveness for alloc_box Promotion
 //===----------------------------------------------------------------------===//
@@ -301,7 +314,7 @@ static bool partialApplyEscapes(SILValue V, bool examineApply) {
 }
 
 static SILInstruction *recursivelyFindBoxOperandsPromotableToAddress(
-    SILValue Box, bool inAppliedFunction, SmallVectorImpl<Operand *> &,
+    SILValue Box, AllocBoxToStackOptions Options, SmallVectorImpl<Operand *> &,
     SmallPtrSetImpl<SILFunction *> &, unsigned CurrentRecurDepth);
 
 /// checkLocalApplyBody - Check the body of an apply's callee to see
@@ -311,7 +324,8 @@ static SILInstruction *recursivelyFindBoxOperandsPromotableToAddress(
 static bool checkLocalApplyBody(Operand *O,
                                 SmallVectorImpl<Operand *> &PromotedOperands,
                                 SmallPtrSetImpl<SILFunction *> &VisitedCallees,
-                                unsigned CurrentRecurDepth) {
+                                unsigned CurrentRecurDepth,
+                                AllocBoxToStackOptions Options) {
   SILFunction *F = ApplySite(O->getUser()).getReferencedFunctionOrNull();
   // If we cannot examine the function body, assume the worst.
   if (!F || F->empty())
@@ -323,10 +337,11 @@ static bool checkLocalApplyBody(Operand *O,
   if (!iter.second)
     return false;
 
+  Options |= {AllocBoxToStackFlags::InAppliedFunction};
+
   auto calleeArg = F->getArgument(ApplySite(O->getUser()).getCalleeArgIndex(*O));
   auto res = !recursivelyFindBoxOperandsPromotableToAddress(
-      calleeArg,
-      /* inAppliedFunction = */ true, PromotedOperands, VisitedCallees,
+      calleeArg, Options, PromotedOperands, VisitedCallees,
       CurrentRecurDepth + 1);
   return res;
 }
@@ -370,7 +385,7 @@ static bool isOptimizableApplySite(ApplySite Apply) {
 /// box is passed don't have any unexpected uses, `PromotedOperands` will be
 /// populated with the box arguments in DFS order.
 static SILInstruction *recursivelyFindBoxOperandsPromotableToAddress(
-    SILValue Box, bool inAppliedFunction,
+    SILValue Box, AllocBoxToStackOptions Options,
     SmallVectorImpl<Operand *> &PromotedOperands,
     SmallPtrSetImpl<SILFunction *> &VisitedCallees,
     unsigned CurrentRecurDepth = 0) {
@@ -394,7 +409,8 @@ static SILInstruction *recursivelyFindBoxOperandsPromotableToAddress(
     // Projections are fine as well.
     if (isa<StrongRetainInst>(User) || isa<StrongReleaseInst>(User) ||
         isa<ProjectBoxInst>(User) || isa<DestroyValueInst>(User) ||
-        (!inAppliedFunction && isa<DeallocBoxInst>(User)) ||
+        (!Options.contains(AllocBoxToStackFlags::InAppliedFunction) &&
+         isa<DeallocBoxInst>(User)) ||
         isa<EndBorrowInst>(User))
       continue;
 
@@ -413,8 +429,17 @@ static SILInstruction *recursivelyFindBoxOperandsPromotableToAddress(
       }
       switch (Apply.getKind()) {
       case ApplySiteKind::PartialApplyInst: {
+        // If we have a noncopyable type in a box and we were passed in the
+        // option that tells us that we should not promote any boxes that are
+        // captured by an escaping closure (even if we can prove the closure
+        // does not escape), treat the partial apply use as an escape.
+        if (Box->getType().isBoxedNonCopyableType(Box->getFunction()) &&
+            !Options.contains(
+                AllocBoxToStackFlags::
+                    CanPromoteNonCopyableCapturedByEscapingClosure))
+          break;
         if (checkLocalApplyBody(Op, LocalPromotedOperands, VisitedCallees,
-                                CurrentRecurDepth) &&
+                                CurrentRecurDepth, Options) &&
             !partialApplyEscapes(cast<PartialApplyInst>(User),
                                  /* examineApply = */ true)) {
           LocalPromotedOperands.push_back(Op);
@@ -427,7 +452,7 @@ static SILInstruction *recursivelyFindBoxOperandsPromotableToAddress(
       case ApplySiteKind::TryApplyInst:
         if (isOptimizableApplySite(Apply) &&
             checkLocalApplyBody(Op, LocalPromotedOperands, VisitedCallees,
-                                CurrentRecurDepth)) {
+                                CurrentRecurDepth, Options)) {
           LocalPromotedOperands.push_back(Op);
           continue;
         }
@@ -450,13 +475,13 @@ static InFlightDiagnostic diagnose(ASTContext &Context, SourceLoc loc,
 
 /// canPromoteAllocBox - Can we promote this alloc_box to an alloc_stack?
 static bool canPromoteAllocBox(AllocBoxInst *ABI,
-                               SmallVectorImpl<Operand *> &PromotedOperands) {
+                               SmallVectorImpl<Operand *> &PromotedOperands,
+                               AllocBoxToStackOptions Options) {
   SmallPtrSet<SILFunction *, 8> VisitedCallees;
   // Scan all of the uses of the address of the box to see if any
   // disqualifies the box from being promoted to the stack.
   if (auto *User = recursivelyFindBoxOperandsPromotableToAddress(
-          ABI,
-          /* inAppliedFunction = */ false, PromotedOperands, VisitedCallees,
+          ABI, Options, PromotedOperands, VisitedCallees,
           /* CurrentRecurDepth = */ 0)) {
     (void)User;
     // Otherwise, we have an unexpected use.
@@ -1221,6 +1246,8 @@ static unsigned rewritePromotedBoxes(AllocBoxToStackState &pass) {
 }
 
 namespace {
+
+template <bool CanPromoteMoveOnlyCapturedByEscapingClosure>
 class AllocBoxToStack : public SILFunctionTransform {
   /// The entry point to the transformation.
   void run() override {
@@ -1229,10 +1256,16 @@ class AllocBoxToStack : public SILFunctionTransform {
       return;
 
     AllocBoxToStackState pass(this);
+
+    AllocBoxToStackOptions Options;
+    if (CanPromoteMoveOnlyCapturedByEscapingClosure)
+      Options |=
+          AllocBoxToStackFlags::CanPromoteNonCopyableCapturedByEscapingClosure;
+
     for (auto &BB : *getFunction()) {
       for (auto &I : BB)
         if (auto *ABI = dyn_cast<AllocBoxInst>(&I))
-          if (canPromoteAllocBox(ABI, pass.PromotedOperands))
+          if (canPromoteAllocBox(ABI, pass.PromotedOperands, Options))
             pass.Promotable.push_back(ABI);
     }
 
@@ -1251,8 +1284,13 @@ class AllocBoxToStack : public SILFunctionTransform {
     }
   }
 };
+
 } // end anonymous namespace
 
 SILTransform *swift::createAllocBoxToStack() {
-  return new AllocBoxToStack();
+  return new AllocBoxToStack<true>();
+}
+
+SILTransform *swift::createEarlyAllocBoxToStack() {
+  return new AllocBoxToStack<false>();
 }

test/SILGen/moveonly_escaping_closure.swift

@@ -125,15 +125,13 @@ func testGlobalClosureCaptureVar() {
 // CHECK:   end_access [[READ_ACCESS]]
 // CHECK: } // end sil function '$s16moveonly_closure29testLocalLetClosureCaptureVaryyFyycfU_'
 func testLocalLetClosureCaptureVar() {
-    var x = SingleElt() // expected-error {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
-    // expected-error @-1 {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
-    // expected-error @-2 {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
+    var x = SingleElt()
     x = SingleElt()
     let f = {
         borrowVal(x)
-        consumeVal(x) // expected-note {{consuming use here}}
-        consumeVal(x) // expected-note {{consuming use here}}
-        borrowConsumeVal(x, x) // expected-note {{consuming use here}}
+        consumeVal(x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
+        consumeVal(x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
+        borrowConsumeVal(x, x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
         // expected-error @-1 {{overlapping accesses, but deinitialization requires exclusive access}}
         // expected-note @-2 {{conflicting access is here}}
     }
@@ -476,14 +474,11 @@ func testGlobalClosureCaptureLet() {
 // CHECK: } // end sil function '$s16moveonly_closure026testLocalLetClosureCaptureE0yyFyycfU_'
 func testLocalLetClosureCaptureLet() {
     let x = SingleElt()
-    // expected-error @-1 {{'x' was consumed but it is illegal to consume a noncopyable immutable capture of an escaping closure. One can only read from it}}
-    // expected-error @-2 {{'x' was consumed but it is illegal to consume a noncopyable immutable capture of an escaping closure. One can only read from it}}
-    // expected-error @-3 {{'x' was consumed but it is illegal to consume a noncopyable immutable capture of an escaping closure. One can only read from it}}
     let f = {
         borrowVal(x)
-        consumeVal(x) // expected-note {{consuming use here}}
-        consumeVal(x) // expected-note {{consuming use here}}
-        borrowConsumeVal(x, x) // expected-note {{consuming use here}}
+        consumeVal(x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable immutable capture of an escaping closure. One can only read from it}}
+        consumeVal(x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable immutable capture of an escaping closure. One can only read from it}}
+        borrowConsumeVal(x, x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable immutable capture of an escaping closure. One can only read from it}}
     }
     f()
 }
@@ -1090,14 +1085,11 @@ func testGlobalClosureCaptureConsuming(_ x: consuming SingleElt) {
 // CHECK:   end_access [[READ_ACCESS]]
 // CHECK: } // end sil function '$s16moveonly_closure35testLocalLetClosureCaptureConsumingyyAA9SingleEltVnFyycfU_'
 func testLocalLetClosureCaptureConsuming(_ x: consuming SingleElt) {
-    // expected-error @-1 {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
-    // expected-error @-2 {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
-    // expected-error @-3 {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
     let f = {
         borrowVal(x)
-        consumeVal(x) // expected-note {{consuming use here}}
-        consumeVal(x) // expected-note {{consuming use here}}
-        borrowConsumeVal(x, x) // expected-note {{consuming use here}}
+        consumeVal(x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
+        consumeVal(x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
+        borrowConsumeVal(x, x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable mutable capture of an escaping closure. One can only read from it or assign over it}}
         // expected-error @-1 {{overlapping accesses, but deinitialization requires exclusive access}}
         // expected-note @-2 {{conflicting access is here}}
     }
@@ -1373,14 +1365,11 @@ func testGlobalClosureCaptureOwned(_ x: __owned SingleElt) {
 // CHECK:   apply {{%.*}}([[LOADED_READ]], [[LOADED_TAKE]])
 // CHECK: } // end sil function '$s16moveonly_closure31testLocalLetClosureCaptureOwnedyyAA9SingleEltVnFyycfU_'
 func testLocalLetClosureCaptureOwned(_ x: __owned SingleElt) {
-    // expected-error @-1 {{'x' was consumed but it is illegal to consume a noncopyable immutable capture of an escaping closure. One can only read from it}}
-    // expected-error @-2 {{'x' was consumed but it is illegal to consume a noncopyable immutable capture of an escaping closure. One can only read from it}}
-    // expected-error @-3 {{'x' was consumed but it is illegal to consume a noncopyable immutable capture of an escaping closure. One can only read from it}}
     let f = {
         borrowVal(x)
-        consumeVal(x) // expected-note {{consuming use here}}
-        consumeVal(x) // expected-note {{consuming use here}}
-        borrowConsumeVal(x, x) // expected-note {{consuming use here}}
+        consumeVal(x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable immutable capture of an escaping closure. One can only read from it}}
+        consumeVal(x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable immutable capture of an escaping closure. One can only read from it}}
+        borrowConsumeVal(x, x) // expected-error {{'x' was consumed but it is illegal to consume a noncopyable immutable capture of an escaping closure. One can only read from it}}
     }
     f()
 }

test/SILOptimizer/allocbox_to_stack_noncopyable.sil

@@ -1,4 +1,5 @@
 // RUN: %target-sil-opt -enable-sil-verify-all %s -allocbox-to-stack -enable-copy-propagation=requested-passes-only | %FileCheck %s
+// RUN: %target-sil-opt -enable-sil-verify-all %s -early-allocbox-to-stack -enable-copy-propagation=requested-passes-only | %FileCheck -check-prefix=EARLY %s
 
 sil_stage raw
 
@@ -160,3 +161,43 @@ bb3:
   %30 = tuple ()
   return %30 : $()
 }
+
+// Make sure that if our box is captured by a partial_apply, we do not process
+// it with the early allocbox to stack, but we do with the normal allocbox to
+// stack.
+// EARLY: sil hidden [ossa] @earlyallocbox_to_stack_partial_apply_test_caller : $@convention(thin) (@owned NonTrivialStruct) -> () {
+// EARLY: alloc_box ${ var NonTrivialStruct }, var, name "x"
+// EARLY: } // end sil function 'earlyallocbox_to_stack_partial_apply_test_caller'
+//
+// CHECK: sil hidden [ossa] @earlyallocbox_to_stack_partial_apply_test_caller : $@convention(thin) (@owned NonTrivialStruct) -> () {
+// CHECK: alloc_stack [lexical] $NonTrivialStruct, var, name "x"
+// CHECK: } // end sil function 'earlyallocbox_to_stack_partial_apply_test_caller'
+sil hidden [ossa] @earlyallocbox_to_stack_partial_apply_test_caller : $@convention(thin) (@owned NonTrivialStruct) -> () {
+bb0(%arg : @owned $NonTrivialStruct):
+  %0 = alloc_box ${ var NonTrivialStruct }, var, name "x"
+  %1 = begin_borrow [lexical] %0 : ${ var NonTrivialStruct }
+  %2 = project_box %1 : ${ var NonTrivialStruct }, 0
+  store %arg to [init] %2 : $*NonTrivialStruct
+  %15 = function_ref @earlyallocbox_to_stack_partial_apply_test_callee : $@convention(thin) (@guaranteed { var NonTrivialStruct }) -> ()
+  %16 = copy_value %1 : ${ var NonTrivialStruct }
+  %18 = partial_apply [callee_guaranteed] %15(%16) : $@convention(thin) (@guaranteed { var NonTrivialStruct }) -> ()
+  destroy_value %18 : $@callee_guaranteed () -> ()
+  end_borrow %1 : ${ var NonTrivialStruct }
+  destroy_value %0 : ${ var NonTrivialStruct }
+  %24 = tuple ()
+  return %24 : $()
+}
+
+sil private [ossa] @earlyallocbox_to_stack_partial_apply_test_callee : $@convention(thin) (@guaranteed { var NonTrivialStruct }) -> () {
+bb0(%0 : @closureCapture @guaranteed ${ var NonTrivialStruct }):
+  %1 = project_box %0 : ${ var NonTrivialStruct }, 0
+  debug_value %1 : $*NonTrivialStruct, var, name "x", argno 1, expr op_deref
+  %3 = begin_access [read] [unknown] %1 : $*NonTrivialStruct
+  %4 = mark_must_check [no_consume_or_assign] %3 : $*NonTrivialStruct
+  %5 = load [copy] %4 : $*NonTrivialStruct
+  end_access %3 : $*NonTrivialStruct
+  %7 = move_value %5 : $NonTrivialStruct
+  destroy_value %7 : $NonTrivialStruct
+  %9 = tuple ()
+  return %9 : $()
+}

test/SILOptimizer/allocbox_to_stack_ownership.sil

@@ -1,4 +1,5 @@
 // RUN: %target-sil-opt -enable-sil-verify-all %s -allocbox-to-stack -enable-copy-propagation=requested-passes-only -enable-lexical-borrow-scopes=false | %FileCheck %s
+// RUN: %target-sil-opt -enable-sil-verify-all %s -early-allocbox-to-stack -enable-copy-propagation=requested-passes-only -enable-lexical-borrow-scopes=false | %FileCheck %s
 
 sil_stage raw
 

test/SILOptimizer/allocboxtostack_escape.swift

@@ -9,9 +9,12 @@ class Klass {}
 func myPrint(_ x: inout Box<Klass>) -> () { print(x) }
 
 func testError() -> (() -> ()) {
+  // We emit this error twice since we run allocbox to stack twice in the pipeline for now. This is not an official feature, so it is ok for now.
   @_semantics("boxtostack.mustbeonstack")
   var x = Box<Klass>(value: Klass()) // expected-error {{Can not promote value from heap to stack due to value escaping}}
+  // expected-error @-1 {{Can not promote value from heap to stack due to value escaping}}
   let result = { // expected-note {{value escapes here}}
+      // expected-note @-1 {{value escapes here}}
     myPrint(&x)
   }
   return result