Update LoadBorrowImmutability and PrunedLiveness to handle reborrows

include/swift/SIL/OwnershipUtils.h

@@ -578,15 +578,20 @@ struct BorrowedValue {
   ///
   /// \p deadEndBlocks is optional during transition. It will be completely
   /// removed in an upcoming commit.
-  bool areUsesWithinLocalScope(ArrayRef<Operand *> uses,
-                               DeadEndBlocks *deadEndBlocks) const;
+  bool areUsesWithinTransitiveScope(ArrayRef<Operand *> uses,
+                                    DeadEndBlocks *deadEndBlocks) const;
 
   /// Given a local borrow scope introducer, visit all non-forwarding consuming
   /// users. This means that this looks through guaranteed block arguments. \p
   /// visitor is *not* called on Reborrows, only on final scope ending uses.
   bool
   visitExtendedScopeEndingUses(function_ref<bool(Operand *)> visitor) const;
 
+  /// Visit all lifetime ending operands of the entire borrow scope including
+  /// reborrows
+  bool
+  visitTransitiveLifetimeEndingUses(function_ref<bool(Operand *)> func) const;
+
   void print(llvm::raw_ostream &os) const;
   SWIFT_DEBUG_DUMP { print(llvm::dbgs()); }
 

include/swift/SIL/PrunedLiveness.h

@@ -164,7 +164,6 @@ class PrunedLiveBlocks {
   }
 
   void initializeDefBlock(SILBasicBlock *defBB) {
-    assert(!seenUse && "cannot initialize more defs with partial liveness");
     markBlockLive(defBB, LiveWithin);
   }
 
@@ -229,6 +228,15 @@ class PrunedLiveness {
   /// ending uses that extend liveness into a loop body.
   SmallSetVector<SILInstruction *, 8> *nonLifetimeEndingUsesInLiveOut;
 
+private:
+  bool isWithinBoundaryHelper(SILInstruction *inst, SILValue def) const;
+
+  bool areUsesWithinBoundaryHelper(ArrayRef<Operand *> uses, SILValue def,
+                                   DeadEndBlocks *deadEndBlocks) const;
+
+  bool areUsesOutsideBoundaryHelper(ArrayRef<Operand *> uses, SILValue def,
+                                    DeadEndBlocks *deadEndBlocks) const;
+
 public:
   PrunedLiveness(SmallVectorImpl<SILBasicBlock *> *discoveredBlocks = nullptr,
                  SmallSetVector<SILInstruction *, 8>
@@ -342,6 +350,25 @@ class PrunedLiveness {
   bool areUsesOutsideBoundary(ArrayRef<Operand *> uses,
                               DeadEndBlocks *deadEndBlocks) const;
 
+  /// PrunedLiveness utilities can be used with multiple defs. This api can be
+  /// used to check if \p inst occurs in between the definition \p def and the
+  /// liveness boundary.
+  // This api varies from isWithinBoundary(SILInstruction *inst) which cannot
+  // distinguish when \p inst is a use before definition in the same block as
+  // the definition.
+  bool isWithinBoundaryOfDef(SILInstruction *inst, SILValue def) const;
+
+  /// Returns true when all \p uses are between \p def and the liveness boundary
+  /// \p deadEndBlocks is optional.
+  bool areUsesWithinBoundaryOfDef(ArrayRef<Operand *> uses, SILValue def,
+                                  DeadEndBlocks *deadEndBlocks) const;
+
+  /// Returns true if any of the \p uses are before the \p def or after the
+  /// liveness boundary
+  /// \p deadEndBlocks is optional.
+  bool areUsesOutsideBoundaryOfDef(ArrayRef<Operand *> uses, SILValue def,
+                                   DeadEndBlocks *deadEndBlocks) const;
+
   /// Compute liveness for a single SSA definition.
   void computeSSALiveness(SILValue def);
 };

lib/SIL/Utils/OwnershipUtils.cpp

@@ -689,13 +689,20 @@ llvm::raw_ostream &swift::operator<<(llvm::raw_ostream &os,
 /// Add this scopes live blocks into the PrunedLiveness result.
 void BorrowedValue::computeLiveness(PrunedLiveness &liveness) const {
   liveness.initializeDefBlock(value->getParentBlock());
-  visitLocalScopeEndingUses([&](Operand *endOp) {
-    liveness.updateForUse(endOp->getUser(), true);
+  visitTransitiveLifetimeEndingUses([&](Operand *endOp) {
+    if (endOp->getOperandOwnership() == OperandOwnership::EndBorrow) {
+      liveness.updateForUse(endOp->getUser(), /*lifetimeEnding*/ true);
+      return true;
+    }
+    assert(endOp->getOperandOwnership() == OperandOwnership::Reborrow);
+    auto *succBlock = cast<BranchInst>(endOp->getUser())->getDestBB();
+    liveness.initializeDefBlock(succBlock);
+    liveness.updateForUse(endOp->getUser(), /*lifetimeEnding*/ false);
     return true;
   });
 }
 
-bool BorrowedValue::areUsesWithinLocalScope(
+bool BorrowedValue::areUsesWithinTransitiveScope(
     ArrayRef<Operand *> uses, DeadEndBlocks *deadEndBlocks) const {
   // First make sure that we actually have a local scope. If we have a non-local
   // scope, then we have something (like a SILFunctionArgument) where a larger
@@ -741,6 +748,38 @@ bool BorrowedValue::visitExtendedScopeEndingUses(
   return true;
 }
 
+bool BorrowedValue::visitTransitiveLifetimeEndingUses(
+    function_ref<bool(Operand *)> visitor) const {
+  assert(isLocalScope());
+
+  SmallPtrSetVector<SILValue, 4> reborrows;
+
+  auto visitEnd = [&](Operand *scopeEndingUse) {
+    if (scopeEndingUse->getOperandOwnership() == OperandOwnership::Reborrow) {
+      BorrowingOperand(scopeEndingUse)
+          .visitBorrowIntroducingUserResults([&](BorrowedValue borrowedValue) {
+            reborrows.insert(borrowedValue.value);
+            return true;
+          });
+      // visitor on the reborrow
+      return visitor(scopeEndingUse);
+    }
+    // visitor on the end_borrow
+    return visitor(scopeEndingUse);
+  };
+
+  if (!visitLocalScopeEndingUses(visitEnd))
+    return false;
+
+  // reborrows grows in this loop.
+  for (unsigned idx = 0; idx < reborrows.size(); ++idx) {
+    if (!BorrowedValue(reborrows[idx]).visitLocalScopeEndingUses(visitEnd))
+      return false;
+  }
+
+  return true;
+}
+
 bool BorrowedValue::visitInteriorPointerOperandHelper(
     function_ref<void(InteriorPointerOperand)> func,
     BorrowedValue::InteriorPointerOperandVisitorKind kind) const {
@@ -975,7 +1014,7 @@ bool AddressOwnership::areUsesWithinLifetime(
   SILValue root = base.getOwnershipReferenceRoot();
   BorrowedValue borrow(root);
   if (borrow)
-    return borrow.areUsesWithinLocalScope(uses, &deadEndBlocks);
+    return borrow.areUsesWithinTransitiveScope(uses, &deadEndBlocks);
 
   // --- A reference no borrow scope. Currently happens for project_box.
 

lib/SIL/Utils/PrunedLiveness.cpp

@@ -127,67 +127,124 @@ bool PrunedLiveness::updateForBorrowingOperand(Operand *op) {
 }
 
 void PrunedLiveness::extendAcrossLiveness(PrunedLiveness &otherLivesness) {
-  // update this liveness for all the interesting users in otherLivesness.
+  // update this liveness for all the interesting users in otherLiveness.
   for (std::pair<SILInstruction *, bool> userAndEnd : otherLivesness.users) {
     updateForUse(userAndEnd.first, userAndEnd.second);
   }
 }
 
-bool PrunedLiveness::isWithinBoundary(SILInstruction *inst) const {
+bool PrunedLiveness::isWithinBoundaryHelper(SILInstruction *inst,
+                                            SILValue def) const {
   SILBasicBlock *block = inst->getParent();
+
+  /// Returns true if \p inst is before \p def in this block.
+  auto foundInBlockBeforeDef = [](SILInstruction *inst, SILBasicBlock *block,
+                                  SILValue def) {
+    if (!def || def->getParentBlock() != block) {
+      return false;
+    }
+    auto *defInst = def->getDefiningInstruction();
+    if (!defInst) {
+      return false;
+    }
+    // Check if instruction is before the definition
+    for (SILInstruction &it :
+         make_range(block->begin(), defInst->getIterator())) {
+      if (&it == inst) {
+        return true;
+      }
+    }
+    return false;
+  };
+
   switch (getBlockLiveness(block)) {
   case PrunedLiveBlocks::Dead:
     return false;
-  case PrunedLiveBlocks::LiveWithin:
-    break;
   case PrunedLiveBlocks::LiveOut:
-    return true;
-  }
-  // The boundary is within this block. This instruction is before the boundary
-  // iff any interesting uses occur after it.
-  for (SILInstruction &inst :
+    return !foundInBlockBeforeDef(inst, block, def);
+  case PrunedLiveBlocks::LiveWithin:
+    if (foundInBlockBeforeDef(inst, block, def)) {
+      return false;
+    }
+    // The boundary is within this block. This instruction is before the
+    // boundary iff any interesting uses occur after it.
+    for (SILInstruction &it :
          make_range(std::next(inst->getIterator()), block->end())) {
-    switch (isInterestingUser(&inst)) {
-    case PrunedLiveness::NonUser:
-      break;
-    case PrunedLiveness::NonLifetimeEndingUse:
-    case PrunedLiveness::LifetimeEndingUse:
-      return true;
+      switch (isInterestingUser(&it)) {
+      case PrunedLiveness::NonUser:
+        break;
+      case PrunedLiveness::NonLifetimeEndingUse:
+      case PrunedLiveness::LifetimeEndingUse:
+        return true;
+      }
     }
+    return false;
   }
-  return false;
 }
 
-bool PrunedLiveness::areUsesWithinBoundary(ArrayRef<Operand *> uses,
-                                           DeadEndBlocks *deadEndBlocks) const {
+bool PrunedLiveness::isWithinBoundary(SILInstruction *inst) const {
+  return isWithinBoundaryHelper(inst, /*def*/ SILValue());
+}
+
+bool PrunedLiveness::isWithinBoundaryOfDef(SILInstruction *inst,
+                                           SILValue def) const {
+  return isWithinBoundaryHelper(inst, def);
+}
+
+bool PrunedLiveness::areUsesWithinBoundaryHelper(
+    ArrayRef<Operand *> uses, SILValue def,
+    DeadEndBlocks *deadEndBlocks) const {
   auto checkDeadEnd = [deadEndBlocks](SILInstruction *inst) {
     return deadEndBlocks && deadEndBlocks->isDeadEnd(inst->getParent());
   };
 
   for (auto *use : uses) {
     auto *user = use->getUser();
-    if (!isWithinBoundary(user) && !checkDeadEnd(user))
+    if (!isWithinBoundaryHelper(user, def) && !checkDeadEnd(user))
       return false;
   }
   return true;
 }
 
-bool PrunedLiveness::areUsesOutsideBoundary(
-    ArrayRef<Operand *> uses, DeadEndBlocks *deadEndBlocks) const {
+bool PrunedLiveness::areUsesWithinBoundary(ArrayRef<Operand *> uses,
+                                           DeadEndBlocks *deadEndBlocks) const {
+  return areUsesWithinBoundaryHelper(uses, SILValue(), deadEndBlocks);
+}
+
+bool PrunedLiveness::areUsesWithinBoundaryOfDef(
+    ArrayRef<Operand *> uses, SILValue def,
+    DeadEndBlocks *deadEndBlocks) const {
+  return areUsesWithinBoundaryHelper(uses, def, deadEndBlocks);
+}
+
+bool PrunedLiveness::areUsesOutsideBoundaryHelper(
+    ArrayRef<Operand *> uses, SILValue def,
+    DeadEndBlocks *deadEndBlocks) const {
   auto checkDeadEnd = [deadEndBlocks](SILInstruction *inst) {
     return deadEndBlocks && deadEndBlocks->isDeadEnd(inst->getParent());
   };
 
   for (auto *use : uses) {
     auto *user = use->getUser();
-    if (isWithinBoundary(user) || checkDeadEnd(user))
+    if (isWithinBoundaryHelper(user, def) || checkDeadEnd(user))
       return false;
   }
   return true;
 }
 
-// An SSA def meets all the criteria for pruned liveness--def dominates all uses
-// with no holes in the liverange. The lifetime-ending uses are also
+bool PrunedLiveness::areUsesOutsideBoundary(
+    ArrayRef<Operand *> uses, DeadEndBlocks *deadEndBlocks) const {
+  return areUsesOutsideBoundaryHelper(uses, SILValue(), deadEndBlocks);
+}
+
+bool PrunedLiveness::areUsesOutsideBoundaryOfDef(
+    ArrayRef<Operand *> uses, SILValue def,
+    DeadEndBlocks *deadEndBlocks) const {
+  return areUsesOutsideBoundaryHelper(uses, def, deadEndBlocks);
+}
+
+// An SSA def meets all the criteria for pruned liveness--def dominates all
+// uses with no holes in the liverange. The lifetime-ending uses are also
 // recorded--destroy_value or end_borrow. However destroy_values may not
 // jointly-post dominate if dead-end blocks are present.
 void PrunedLiveness::computeSSALiveness(SILValue def) {
@@ -278,8 +335,8 @@ void PrunedLivenessBoundary::compute(const PrunedLiveness &liveness,
     // Process each block that has not been visited and is not LiveOut.
     switch (liveness.getBlockLiveness(bb)) {
     case PrunedLiveBlocks::LiveOut:
-      // A lifetimeEndBlock may be determined to be LiveOut after analyzing the
-      // extended liveness. It is irrelevant for finding the boundary.
+      // A lifetimeEndBlock may be determined to be LiveOut after analyzing
+      // the extended liveness. It is irrelevant for finding the boundary.
       break;
     case PrunedLiveBlocks::LiveWithin: {
       // The liveness boundary is inside this block. Insert a final destroy