Add and improve some "mandatory" peephole optimization to fix a performance diagnostic

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;
     }
 

lib/SILOptimizer/Transforms/GenericSpecializer.cpp

@@ -28,6 +28,7 @@
 #include "swift/SILOptimizer/Utils/InstructionDeleter.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
 #include "swift/SILOptimizer/Utils/SILInliner.h"
+#include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "llvm/ADT/SmallVector.h"
 
 using namespace swift;
@@ -195,10 +196,12 @@ class MandatoryGenericSpecializer : public SILModuleTransform {
 
   void run() override;
 
-  bool optimize(SILFunction *func, ClassHierarchyAnalysis *cha);
+  bool optimize(SILFunction *func, ClassHierarchyAnalysis *cha,
+                bool &invalidatedStackNesting);
 
   bool optimizeInst(SILInstruction *inst, SILOptFunctionBuilder &funcBuilder,
-                    InstructionDeleter &deleter, ClassHierarchyAnalysis *cha);
+                    InstructionDeleter &deleter, ClassHierarchyAnalysis *cha,
+                    bool &invalidatedStackNesting);
 };
 
 
@@ -220,7 +223,7 @@ void MandatoryGenericSpecializer::run() {
       visited.insert(&function);
     }
   }
-  
+
   while (!workList.empty()) {
     SILFunction *func = workList.pop_back_val();
     module->linkFunction(func, SILModule::LinkingMode::LinkAll);
@@ -229,20 +232,26 @@ void MandatoryGenericSpecializer::run() {
 
     // Perform generic specialization and other related optimization.
 
+    bool invalidatedStackNesting = false;
+
     // To avoid phase ordering problems of the involved optimizations, iterate
     // until we reach a fixed point.
     // This should always happen, but to be on the safe side, limit the number
     // of iterations to 10 (which is more than enough - usually the loop runs
     // 1 to 3 times).
     for (int i = 0; i < 10; i++) {
-      bool changed = optimize(func, cha);
+      bool changed = optimize(func, cha, invalidatedStackNesting);
       if (changed) {
         invalidateAnalysis(func, SILAnalysis::InvalidationKind::FunctionBody);
       } else {
         break;
       }
     }
 
+    if (invalidatedStackNesting) {
+      StackNesting::fixNesting(func);
+    }
+
     // Continue specializing called functions.
     for (SILBasicBlock &block : *func) {
       for (SILInstruction &inst : block) {
@@ -260,7 +269,8 @@ void MandatoryGenericSpecializer::run() {
 /// Specialize generic calls in \p func and do some other related optimizations:
 /// devirtualization and constant-folding of the Builtin.canBeClass.
 bool MandatoryGenericSpecializer::optimize(SILFunction *func,
-                                           ClassHierarchyAnalysis *cha) {
+                                           ClassHierarchyAnalysis *cha,
+                                           bool &invalidatedStackNesting) {
   bool changed = false;
   SILOptFunctionBuilder funcBuilder(*this);
   InstructionDeleter deleter;
@@ -282,7 +292,7 @@ bool MandatoryGenericSpecializer::optimize(SILFunction *func,
       continue;
 
     for (SILInstruction *inst : deleter.updatingReverseRange(&block)) {
-      changed |= optimizeInst(inst, funcBuilder, deleter, cha);
+      changed |= optimizeInst(inst, funcBuilder, deleter, cha, invalidatedStackNesting);
     }
   }
   deleter.cleanupDeadInstructions();
@@ -295,7 +305,8 @@ bool MandatoryGenericSpecializer::optimize(SILFunction *func,
 
 bool MandatoryGenericSpecializer::
 optimizeInst(SILInstruction *inst, SILOptFunctionBuilder &funcBuilder,
-             InstructionDeleter &deleter, ClassHierarchyAnalysis *cha) {
+             InstructionDeleter &deleter, ClassHierarchyAnalysis *cha,
+             bool &invalidatedStackNesting) {
   if (auto as = ApplySite::isa(inst)) {
 
     bool changed = false;
@@ -307,10 +318,22 @@ optimizeInst(SILInstruction *inst, SILOptFunctionBuilder &funcBuilder,
       as = newAS;
     }
 
-    auto fas = FullApplySite::isa(as.getInstruction());
-    if (!fas)
+    if (auto *pai = dyn_cast<PartialApplyInst>(as)) {
+      SILBuilderContext builderCtxt(funcBuilder.getModule());
+      if (tryOptimizeApplyOfPartialApply(pai, builderCtxt, deleter.getCallbacks())) {
+        // Try to delete the partial_apply.
+        // We don't need to copy all arguments again (to extend their lifetimes),
+        // because it was already done in tryOptimizeApplyOfPartialApply.
+        tryDeleteDeadClosure(pai, deleter.getCallbacks(), /*needKeepArgsAlive=*/ false);
+        invalidatedStackNesting = true;
+        return true;
+      }
       return changed;
-
+    }
+
+    auto fas = FullApplySite::isa(as.getInstruction());
+    assert(fas);
+
     SILFunction *callee = fas.getReferencedFunctionOrNull();
     if (!callee)
       return changed;

lib/SILOptimizer/Utils/PartialApplyCombiner.cpp

@@ -240,8 +240,8 @@ bool PartialApplyCombiner::combine() {
     auto *user = use->getUser();
 
     // Recurse through copy_value
-    if (auto *cvi = dyn_cast<CopyValueInst>(user)) {
-      for (auto *copyUse : cvi->getUses())
+    if (isa<CopyValueInst>(user) || isa<BeginBorrowInst>(user)) {
+      for (auto *copyUse : cast<SingleValueInstruction>(user)->getUses())
         worklist.push_back(copyUse);
       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
 }
 
 ////////////////////////

test/SILOptimizer/performance-annotations.swift

@@ -213,3 +213,13 @@ func goodClosure2() {
     })
 }
 
+@_noAllocation
+func closueWhichModifiesLocalVar() -> Int {
+  var x = 42
+  let localNonEscapingClosure = {
+    x += 1
+  }
+  localNonEscapingClosure()
+  return x
+}
+

test/SILOptimizer/sil_combine_apply_ossa.sil

@@ -1152,3 +1152,22 @@ bb3:
   %rv = tuple()
   return %rv : $()
 }
+
+// CHECK-LABEL: sil [ossa] @test_copy_and_borrow_of_closure
+// CHECK-NOT: partial_apply
+// CHECK-LABEL: } // end sil function 'test_copy_and_borrow_of_closure'
+sil [ossa] @test_copy_and_borrow_of_closure : $@convention(thin) (Int) -> () {
+bb0(%0 : $Int):
+  %3 = function_ref @closure2 : $@convention(thin) (Int) -> ()
+  %5 = partial_apply %3(%0) : $@convention(thin) (Int) -> ()
+  %6 = begin_borrow [lexical] %5 : $@callee_owned () -> ()
+  %7 = copy_value %6 : $@callee_owned () -> ()
+  apply %7() : $@callee_owned () -> ()
+  end_borrow %6 : $@callee_owned () -> ()
+  destroy_value %5 : $@callee_owned () -> ()
+  %r = tuple ()
+  return %r : $()
+}
+
+sil [ossa] @closure2 : $@convention(thin) (Int) -> ()
+