Merge pull request #35676 from atrick/fix-copyprop-reborrow

Fix CopyPropagation's reborrow handling.

Author: atrick

lib/SILOptimizer/Utils/CanonicalOSSALifetime.cpp

@@ -142,6 +142,10 @@ bool CanonicalizeOSSALifetime::computeBorrowLiveness() {
   if (!EnableRewriteBorrows) {
     return false;
   }
+  // Note that there is no need to look through any reborrows. The reborrowed
+  // value is considered a separate lifetime for canonicalization. Any copies of
+  // the reborrowed value will not be rewritten when canonicalizing the current
+  // borrow scope because they are "hidden" behind the reborrow.
   borrowedVal.visitLocalScopeEndingUses([this](Operand *use) {
     liveness.updateForUse(use->getUser(), /*lifetimeEnding*/ true);
     return true;
@@ -174,6 +178,10 @@ static CopyValueInst *createOuterCopy(BeginBorrowInst *beginBorrow) {
   return copy;
 }
 
+// If this succeeds, then all uses of the borrowed value outside the borrow
+// scope will be rewritten to use an outer copy, and all remaining uses of the
+// borrowed value will be confined to the borrow scope.
+//
 // TODO: Canonicalize multi-block borrow scopes, load_borrow scope, and phi
 // borrow scopes by adding one copy per block to persistentCopies for
 // each block that dominates an outer use.
@@ -200,7 +208,16 @@ bool CanonicalizeOSSALifetime::consolidateBorrowScope() {
     outerUses.push_back(use);
     outerUseInsts.insert(use->getUser());
   };
-
+  // getCanonicalCopiedDef ensures that if currentDef is a guaranteed value,
+  // then it is a borrow scope introducer.
+  assert(BorrowedValue(currentDef).isLocalScope());
+
+  // This def-use traversal is similar to
+  // findExtendedTransitiveGuaranteedUses(), however, to cover the canonical
+  // lifetime, it looks through copies. It also considered uses within the
+  // introduced borrow scope itself (instead of simply visiting the scope-ending
+  // uses). It does not, however, look into nested borrow scopes uses, since
+  // nested scopes are canonicalized independently.
   defUseWorklist.clear();
   defUseWorklist.insert(currentDef);
   while (!defUseWorklist.empty()) {
@@ -227,11 +244,18 @@ bool CanonicalizeOSSALifetime::consolidateBorrowScope() {
       case OperandOwnership::EndBorrow:
       case OperandOwnership::Reborrow:
         // Ignore uses that must be within the borrow scope.
+        // Rewriting does not look through reborrowed values--it considers them
+        // part of a separate lifetime.
         break;
 
       case OperandOwnership::ForwardingUnowned:
       case OperandOwnership::PointerEscape:
-        return false;
+        // Pointer escapes are only allowed if they use the guaranteed value,
+        // which means that the escaped value must be confied to the current
+        // borrow scope.
+        if (use->get().getOwnershipKind() != OwnershipKind::Guaranteed)
+          return false;
+        break;
 
       case OperandOwnership::InstantaneousUse:
       case OperandOwnership::UnownedInstantaneousUse:
@@ -242,21 +266,20 @@ bool CanonicalizeOSSALifetime::consolidateBorrowScope() {
         break;
       case OperandOwnership::Borrow:
         BorrowingOperand borrowOper(use);
-        if (borrowOper.kind == BorrowingOperandKind::Invalid) {
-          return false;
-        }
+        assert(borrowOper && "BorrowingOperand must handle OperandOwnership");
         recordOuterUse(use);
         // For borrows, record the scope-ending instructions in addition to the
-        // borrow instruction as an outer use point.
-        borrowOper.visitScopeEndingUses([&](Operand *endBorrow) {
+        // borrow instruction outer use points. The logic below to check whether
+        // a borrow scope is an outer use must visit the same set of uses.
+        borrowOper.visitExtendedScopeEndingUses([&](Operand *endBorrow) {
           if (!isUserInLiveOutBlock(endBorrow->getUser())) {
             outerUseInsts.insert(endBorrow->getUser());
           }
           return true;
         });
         break;
       }
-    }
+    } // end switch OperandOwnership
   } // end def-use traversal
 
   auto *beginBorrow = cast<BeginBorrowInst>(currentDef);
@@ -286,7 +309,7 @@ bool CanonicalizeOSSALifetime::consolidateBorrowScope() {
       }
       // For sub-borrows also check that the scope-ending instructions are
       // within the scope.
-      if (borrowOper.visitScopeEndingUses([&](Operand *endBorrow) {
+      if (borrowOper.visitExtendedScopeEndingUses([&](Operand *endBorrow) {
             return !outerUseInsts.count(endBorrow->getUser());
           })) {
         continue;
@@ -364,7 +387,7 @@ bool CanonicalizeOSSALifetime::computeCanonicalLiveness() {
         llvm_unreachable("this operand cannot handle ownership");
 
       // Conservatively treat a conversion to an unowned value as a pointer
-      // escape. Is it legal to canonicalize these?
+      // escape. Is it legal to canonicalize ForwardingUnowned?
       case OperandOwnership::ForwardingUnowned:
       case OperandOwnership::PointerEscape:
         return false;
@@ -385,12 +408,13 @@ bool CanonicalizeOSSALifetime::computeCanonicalLiveness() {
         recordConsumingUse(use);
         break;
       case OperandOwnership::Borrow:
-        // An entire borrow scope is considered a single use that occurs at the
-        // point of the end_borrow.
-        BorrowingOperand(use).visitScopeEndingUses([this](Operand *end) {
-          liveness.updateForUse(end->getUser(), /*lifetimeEnding*/ false);
-          return true;
-        });
+        // The liveness of an entire extended borrow scope is modeled as use
+        // points and the scope-ending uses.
+        BorrowingOperand(use).visitExtendedScopeEndingUses(
+          [this](Operand *end) {
+            liveness.updateForUse(end->getUser(), /*lifetimeEnding*/ false);
+            return true;
+          });
         break;
       case OperandOwnership::InteriorPointer:
       case OperandOwnership::ForwardingBorrow:

test/SIL/ownership-verifier/borrow_validate.sil

@@ -740,3 +740,31 @@ bb3(%borrow : @guaranteed $FakeOptional<Klass>, %mKlass : @owned $FakeOptional<K
   %9999 = tuple()
   return %9999 : $()
 }
+
+// Test copying a borrowed value after it has been reborrowed in a different block.
+//
+// CHECK-LABEL: Error#: 0. Begin Error in Function: 'test_copy_borrow_after_reborrow'
+// CHECK: Found outside of lifetime use!
+// CHECK: Value:   %1 = begin_borrow %0 : $Klass
+// CHECK: User:  %6 = copy_value %1 : $Klass
+// CHECK: Block: bb3
+// CHECK-LABEL: End Error in Function: 'test_copy_borrow_after_reborrow'
+sil [ossa] @test_copy_borrow_after_reborrow : $@convention(thin) () -> () {
+bb0:
+  %alloc = alloc_ref $Klass
+  %borrow = begin_borrow %alloc : $Klass
+  cond_br undef, bb1, bb2
+bb1:
+  br bb3(%borrow: $Klass)
+bb2:
+  br bb3(%borrow: $Klass)
+bb3(%borrowphi : @guaranteed $Klass):
+  %innercopy = copy_value %borrow : $Klass
+  end_borrow %borrowphi : $Klass
+  %outercopy = copy_value %innercopy : $Klass
+  destroy_value %innercopy : $Klass
+  destroy_value %outercopy : $Klass
+  destroy_value %alloc : $Klass
+  %99 = tuple ()
+  return %99 : $()
+}

test/SILOptimizer/copy_propagation.sil

@@ -1345,19 +1345,21 @@ bb3:
   return %result : $()
 }
 
-// Do not consolidate this local borrowscope because it has a PointerEscape.
+// Consolidate this local borrowscope even though it has a
+// PointerEscape. The escaping value can be assumed not to be used
+// outside the borrow scope.
 //
 // CHECK-LABEL: sil [ossa] @testBorrowEscape : $@convention(thin) (@guaranteed C) -> () {
 // CHECK:       [[BORROW:%.*]] = begin_borrow %0 : $C
-// CHECK:       [[COPY:%.*]] = copy_value [[BORROW]] : $C
-// CHECK-NEXT:  end_borrow %1 : $C
+// CHECK-NOT:   copy_value
+// CHECK:       end_borrow [[BORROW]] : $C
 // CHECK-NEXT:  cond_br undef, bb1, bb2
 // CHECK:     bb1:
+// CHECK:       [[COPY:%.*]] = copy_value %0 : $C
 // CHECK:       apply %{{.*}}([[COPY]]) : $@convention(thin) (@owned C) -> ()
 // CHECK:       br bb3
 // CHECK:     bb2:
-// CHECK-NEXT:  destroy_value [[COPY]] : $C
-// CHECK:     br bb3
+// CHECK-NEXT:  br bb3
 // CHECK:     bb3:
 // CHECK-NOT:   destroy
 // CHECK-LABEL: } // end sil function 'testBorrowEscape'
@@ -1447,3 +1449,182 @@ bb3:
   destroy_value %2 : $C
   return %0 : $C
 }
+
+// =============================================================================
+// Test reborrows
+// =============================================================================
+
+// Extend the lifetime of an owned value through a nested borrow scope
+// with cross-block reborrows. Its lifetime must be extended past the
+// end_borrow of the phi.
+//
+// CHECK-LABEL: sil [ossa] @testSubReborrowExtension : $@convention(thin) () -> () {
+// CHECK: bb0:
+// CHECK: [[ALLOC:%.*]] = alloc_ref $C
+// CHECK: [[BORROW:%.*]] = begin_borrow %0 : $C
+// CHECK:   cond_br undef, bb1, bb2
+// CHECK: bb1:
+// CHECK:   [[CP1:%.*]] = copy_value %0 : $C
+// CHECK:   br bb3([[CP1]] : $C, [[BORROW]] : $C)
+// CHECK: bb2:
+// CHECK:   [[CP2:%.*]] = copy_value %0 : $C
+// CHECK:   br bb3([[CP2]] : $C, [[BORROW]] : $C)
+// CHECK: bb3([[OWNEDPHI:%.*]] : @owned $C, [[BORROWPHI:%.*]] @guaranteed $C
+// CHECK:   end_borrow [[BORROWPHI]]
+// CHECK:   destroy_value [[OWNEDPHI]] : $C
+// CHECK:   destroy_value %0 : $C
+// CHECK-LABEL: } // end sil function 'testSubReborrowExtension'
+sil [ossa] @testSubReborrowExtension : $@convention(thin) () -> () {
+bb0:
+  %alloc = alloc_ref $C
+  %copy = copy_value %alloc : $C
+  %borrow = begin_borrow %alloc : $C
+  cond_br undef, bb1, bb2
+bb1:
+  br bb3(%copy : $C, %borrow: $C)
+bb2:
+  br bb3(%copy : $C, %borrow: $C)
+bb3(%phi : @owned $C, %borrowphi : @guaranteed $C):
+  end_borrow %borrowphi : $C
+  destroy_value %phi : $C
+  destroy_value %alloc : $C
+  %99 = tuple ()
+  return %99 : $()
+}
+
+// Test a cross-block reborrow and a live nested copy of the reborrow:
+//   def: borrow
+//   use: phi reborrow (in the same block as the borrow)
+//   copy reborrowed phi
+//   end reborrow
+//   copy outside borrow scope.
+//
+// consolidateBorrow only processes the SSA borrow scope, not the
+// extended borrow scope, and it does not process the borrow scope
+// introduced by the phi. Therefore, it should leave the in-scope copy
+// alone-- it will be treated like the definition of a separate OSSA
+// lifetime. The inner copy's lifetime will be canonicalized, removing
+// outercopy.
+//
+// CHECK-LABEL: sil [ossa] @testLiveCopyAfterReborrow : $@convention(thin) () -> () {
+// CHECK: [[ALLOC:%.*]] = alloc_ref $C
+// CHECK: bb3([[BORROWPHI:%.*]] : @guaranteed $C):
+// CHECK: [[COPY:%.*]] = copy_value [[BORROWPHI]]
+// CHECK: end_borrow [[BORROWPHI]] : $C
+// CHECK-NOT: copy_value
+// CHECK-NOT: destroy_value
+// CHECK: apply
+// CHECK: destroy_value [[COPY]]
+// CHECK: destroy_value [[ALLOC]] : $C
+// CHECK-LABEL: } // end sil function 'testLiveCopyAfterReborrow'
+sil [ossa] @testLiveCopyAfterReborrow : $@convention(thin) () -> () {
+bb0:
+  %alloc = alloc_ref $C
+  cond_br undef, bb1, bb2
+bb1:
+  %borrow1 = begin_borrow %alloc : $C
+  br bb3(%borrow1: $C)
+bb2:
+  %borrow2 = begin_borrow %alloc : $C
+  br bb3(%borrow2: $C)
+bb3(%borrowphi : @guaranteed $C):
+  %innercopy = copy_value %borrowphi : $C
+  end_borrow %borrowphi : $C
+  %outercopy = copy_value %innercopy : $C
+  destroy_value %innercopy : $C
+  %f = function_ref @takeGuaranteedC : $@convention(thin) (@guaranteed C) -> ()
+  apply %f(%outercopy) : $@convention(thin) (@guaranteed C) -> ()
+  destroy_value %outercopy : $C
+  destroy_value %alloc : $C
+  %99 = tuple ()
+  return %99 : $()
+}
+
+// Test a cross-block reborrow and a dead nested copy of the reborrow:
+//   def: borrow
+//   use: phi reborrow (in the same block as the borrow)
+//   copy reborrowed phi
+//   end reborrow
+//   use of copied reborrow outside borrow scope.
+//
+// consolidateBorrow only processes the SSA borrow scope, not the
+// extended borrow scope, and it does not process the borrow scope
+// introduced by the phi. Therefore, it should leave the in-scope copy alone--
+// it will be treated like the definition of a separate OSSA lifetime.
+// The inner copy's lifetime will be canonicalized, removing
+// outercopy.
+//
+// CHECK-LABEL: sil [ossa] @testDeadCopyAfterReborrow : $@convention(thin) () -> () {
+// CHECK: [[ALLOC:%.*]] = alloc_ref $C
+// CHECK: bb3([[BORROWPHI:%.*]] : @guaranteed $C):
+// CHECK-NOT: copy_value
+// CHECK: end_borrow [[BORROWPHI]] : $C
+// CHECK-NOT: copy_value
+// CHECK: destroy_value [[ALLOC]] : $C
+// CHECK-LABEL: } // end sil function 'testDeadCopyAfterReborrow'
+sil [ossa] @testDeadCopyAfterReborrow : $@convention(thin) () -> () {
+bb0:
+  %alloc = alloc_ref $C
+  cond_br undef, bb1, bb2
+bb1:
+  %borrow1 = begin_borrow %alloc : $C
+  br bb3(%borrow1: $C)
+bb2:
+  %borrow2 = begin_borrow %alloc : $C
+  br bb3(%borrow2: $C)
+bb3(%borrowphi : @guaranteed $C):
+  %innercopy = copy_value %borrowphi : $C
+  end_borrow %borrowphi : $C
+  %outercopy = copy_value %innercopy : $C
+  destroy_value %innercopy : $C
+  destroy_value %outercopy : $C
+  destroy_value %alloc : $C
+  %99 = tuple ()
+  return %99 : $()
+}
+
+// Test a reborrow with a nested borrow with an outside use
+//   def: borrowphi
+//   nested borrow -- consolidated borrow scope
+//   inner copy -- removed when rewriting nested borrow
+//   end nested borrow
+//   middle copy -- when consolidating borrow scope,
+//                  its operand is replace with a copy of borrowphi
+//                  then removeed when borrowphi's copy is canonicalized
+//   end borrowphi
+//   outer copy -- removeed when borrowphi's copy is canonicalized
+//
+// CHECK-LABEL: sil [ossa] @testNestedReborrowOutsideUse : $@convention(thin) () -> () {
+// CHECK:        [[ALLOC:%.*]] = alloc_ref $C
+// CHECK:      bb3([[BORROWPHI:%.*]] : @guaranteed $C):
+// CHECK-NEXT:   [[COPY:%.*]] = copy_value [[BORROWPHI]] : $C
+// CHECK-NEXT:   begin_borrow [[BORROWPHI]] : $C
+// CHECK-NEXT:   end_borrow
+// CHECK-NEXT:   end_borrow
+// CHECK-NEXT:   destroy_value [[COPY]] : $C
+// CHECK-NEXT:   destroy_value [[ALLOC]] : $C
+// CHECK-LABEL: } // end sil function 'testNestedReborrowOutsideUse'
+sil [ossa] @testNestedReborrowOutsideUse : $@convention(thin) () -> () {
+bb0:
+  %alloc = alloc_ref $C
+  cond_br undef, bb1, bb2
+bb1:
+  %borrow1 = begin_borrow %alloc : $C
+  br bb3(%borrow1: $C)
+bb2:
+  %borrow2 = begin_borrow %alloc : $C
+  br bb3(%borrow2: $C)
+bb3(%borrowphi : @guaranteed $C):
+  %nestedborrow = begin_borrow %borrowphi : $C
+  %innercopy = copy_value %nestedborrow : $C
+  end_borrow %nestedborrow : $C
+  %middlecopy = copy_value %innercopy : $C
+  end_borrow %borrowphi : $C
+  %outercopy = copy_value %middlecopy : $C
+  destroy_value %innercopy : $C
+  destroy_value %middlecopy : $C
+  destroy_value %outercopy : $C
+  destroy_value %alloc : $C
+  %99 = tuple ()
+  return %99 : $()
+}