Merge pull request #39883 from nate-chandler/lexical_lifetimes/avoid-incorrect-elimination

[SILOptimizer] Keep lexical lifetime markers.

Author: nate-chandler

lib/SILOptimizer/SemanticARC/BorrowScopeOpts.cpp

@@ -29,6 +29,11 @@ bool SemanticARCOptVisitor::visitBeginBorrowInst(BeginBorrowInst *bbi) {
   if (!ctx.shouldPerform(ARCTransformKind::RedundantBorrowScopeElimPeephole))
     return false;
 
+  // Lexical borrow scopes must remain in order to ensure that value lifetimes
+  // are not observably shortened.
+  if (bbi->isLexical())
+    return false;
+
   auto kind = bbi->getOperand().getOwnershipKind();
   SmallVector<EndBorrowInst *, 16> endBorrows;
   for (auto *op : bbi->getUses()) {

lib/SILOptimizer/Transforms/TempRValueElimination.cpp

@@ -495,6 +495,12 @@ void TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
   if (!tempObj)
     return;
 
+  // If the storage corresponds to a source-level var, do not optimize here.
+  // Mem2Reg will transform the lexical alloc_stack into a lexical begin_borrow
+  // which will ensure that the value's lifetime isn't observably shortened.
+  if (tempObj->isLexical())
+    return;
+
   bool isOSSA = copyInst->getFunction()->hasOwnership();
 
   SILValue copySrc = copyInst->getSrc();
@@ -633,6 +639,13 @@ TempRValueOptPass::tryOptimizeStoreIntoTemp(StoreInst *si) {
     return std::next(si->getIterator());
   }
 
+  // If the storage corresponds to a source-level var, do not optimize here.
+  // Mem2Reg will transform the lexical alloc_stack into a lexical begin_borrow
+  // which will ensure that the value's lifetime isn't observably shortened.
+  if (tempObj->isLexical()) {
+    return std::next(si->getIterator());
+  }
+
   // If our tempObj has a dynamic lifetime (meaning it is conditionally
   // initialized, conditionally taken, etc), we can not convert its uses to SSA
   // while eliminating it simply. So bail.

lib/SILOptimizer/Utils/CanonicalizeInstruction.cpp

@@ -450,6 +450,11 @@ static SILBasicBlock::iterator
 eliminateSimpleBorrows(BeginBorrowInst *bbi, CanonicalizeInstruction &pass) {
   auto next = std::next(bbi->getIterator());
 
+  // Never eliminate lexical borrow scopes.  They must be kept to ensure that
+  // value lifetimes aren't observably shortened.
+  if (bbi->isLexical())
+    return next;
+
   // We know that our borrow is completely within the lifetime of its base value
   // if the borrow is never reborrowed. We check for reborrows and do not
   // optimize such cases. Otherwise, we can eliminate our borrow and instead use

validation-test/SILOptimizer/lexical-lifetimes.swift

@@ -0,0 +1,54 @@
+// RUN: %target-run-simple-swift -Xfrontend -enable-experimental-lexical-lifetimes -O -Xfrontend -enable-copy-propagation | %FileCheck %s
+
+// =============================================================================
+// = Declarations                                                           {{ =
+// =============================================================================
+
+class C {
+  init(_ d: D) {
+    self.d = d
+  }
+  weak var d: D?
+  func foo(_ string: String) {
+    d?.cWillFoo(self, string)
+  }
+}
+class D {
+  func cWillFoo(_ c: C, _ string: String) {
+    print(#function, string)
+  }
+}
+
+// =============================================================================
+// = Declarations                                                           }} =
+// =============================================================================
+
+// =============================================================================
+// = Tests                                                                  {{ =
+// =============================================================================
+
+func test_localLetKeepsObjectAliveBeyondCallToClassWithWeakReference() {
+  let d = D()
+  let c = C(d)
+  // CHECK: cWillFoo{{.*}} test_localLetKeepsObjectAliveBeyondCallToClassWithWeakReference
+  c.foo(#function)
+}
+
+func test_localVarKeepsObjectAliveBeyondCallToClassWithWeakReference() {
+  var d = D()
+  let c = C(d)
+  // CHECK: cWillFoo{{.*}} test_localVarKeepsObjectAliveBeyondCallToClassWithWeakReference
+  c.foo(#function)
+}
+
+// =============================================================================
+// = Tests                                                                  }} =
+// =============================================================================
+
+func run() {
+  test_localLetKeepsObjectAliveBeyondCallToClassWithWeakReference()
+  test_localVarKeepsObjectAliveBeyondCallToClassWithWeakReference()
+}
+
+run()
+