[move-only] Some fixes around closures.

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/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -1526,6 +1526,13 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
 
     if (markedValue->getCheckKind() ==
         MarkMustCheckInst::CheckKind::NoConsumeOrAssign) {
+      if (isa<ProjectBoxInst>(stripAccessMarkers(markedValue->getOperand()))) {
+        LLVM_DEBUG(llvm::dbgs()
+                   << "Found mark must check [nocopy] use of escaping box: " << *user);
+        diagnosticEmitter.emitAddressEscapingClosureCaptureLoadedAndConsumed(markedValue);
+        emittedEarlyDiagnostic = true;
+        return true;
+      }
       LLVM_DEBUG(llvm::dbgs()
                  << "Found mark must check [nocopy] error: " << *user);
       diagnosticEmitter.emitAddressDiagnosticNoCopy(markedValue, copyAddr);

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