Update PrunedLiveness

- To detect instructions occuring before definition
- To handle reborrows

Author: meg-gupta

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>
@@ -334,14 +342,26 @@ class PrunedLiveness {
   /// client already knows that inst occurs after the start of liveness.
   bool isWithinBoundary(SILInstruction *inst) const;
 
+  /// Return true if \p inst occurs in between the definition \p def and the
+  /// liveness boundary.
+  bool isWithinBoundaryOfDef(SILInstruction *inst, SILValue def) const;
+
   /// \p deadEndBlocks is optional.
   bool areUsesWithinBoundary(ArrayRef<Operand *> uses,
                              DeadEndBlocks *deadEndBlocks) const;
 
+  /// \p deadEndBlocks is optional.
+  bool areUsesWithinBoundaryOfDef(ArrayRef<Operand *> uses, SILValue def,
+                                  DeadEndBlocks *deadEndBlocks) const;
+
   /// \p deadEndBlocks is optional.
   bool areUsesOutsideBoundary(ArrayRef<Operand *> uses,
                               DeadEndBlocks *deadEndBlocks) const;
 
+  /// \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,8 +689,15 @@ 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;
   });
 }

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