[OwnershipLiveness] End at extend_lifetime.

When visiting consumes, also visit `extend_lifetime` instructions.
These instructions are not lifetime ending, but together with the
consumes, they enclose the users of a value.

Add a flag to LinearLiveness to control whether these instructions are
added so that the verifier can use verify that all such instructions
appear outside the linear lifetime boundary (not including them).

Author: nate-chandler

SwiftCompilerSources/Sources/Optimizer/Utilities/OwnershipLiveness.swift

@@ -64,8 +64,11 @@ func computeLinearLiveness(for definingValue: Value, _ context: Context)
   var range = InstructionRange(for: definingValue, context)
 
   // Compute liveness.
-  definingValue.lookThroughBorrowedFromUser.uses.endingLifetime.forEach {
-    range.insert($0.instruction)
+ for use in definingValue.lookThroughBorrowedFromUser.uses {
+    let instruction = use.instruction
+    if use.endsLifetime || instruction is ExtendLifetimeInst {
+      range.insert(instruction)
+    }
   }
   return range
 }

include/swift/SIL/OwnershipLiveness.h

@@ -206,9 +206,19 @@ class LinearLiveness : public OSSALiveness {
   friend LinearLivenessVisitor;
 
 public:
-  LinearLiveness(SILValue def);
+  /// Whether extend_lifetime instructions should be added to the boundary.
+  /// Used to verify extend_lifetime instructions.
+  enum IncludeExtensions_t {
+    DoNotIncludeExtensions = false,
+    IncludeExtensions = true,
+  };
+  LinearLiveness(SILValue def,
+                 IncludeExtensions_t includeExtensions = IncludeExtensions);
 
   void compute();
+
+private:
+  const IncludeExtensions_t includeExtensions;
 };
 
 // Internal implementation

include/swift/SIL/OwnershipUseVisitor.h

@@ -213,6 +213,13 @@ bool OwnershipUseVisitor<Impl>::visitLifetimeEndingUses(SILValue ssaDef) {
 template <typename Impl>
 bool OwnershipUseVisitor<Impl>::visitConsumes(SILValue ssaDef) {
   for (Operand *use : ssaDef->getUses()) {
+    // extend_lifetime instructions are non-consuming but need to be visited
+    // because together with consuming uses they enclose all users of the value.
+    if (isa<ExtendLifetimeInst>(use->getUser())) {
+      if (!handleUsePoint(use, UseLifetimeConstraint::NonLifetimeEnding))
+        return false;
+      continue;
+    }
     if (use->isConsuming()) {
       if (PhiOperand(use) && !asImpl().handleOwnedPhi(use))
         return false;
@@ -245,6 +252,9 @@ bool OwnershipUseVisitor<Impl>::visitOuterBorrowScopeEnd(Operand *borrowEnd) {
 
     return handleUsePoint(borrowEnd, UseLifetimeConstraint::LifetimeEnding);
 
+  case OperandOwnership::InstantaneousUse:
+    assert(isa<ExtendLifetimeInst>(borrowEnd->getUser()));
+    return handleUsePoint(borrowEnd, UseLifetimeConstraint::NonLifetimeEnding);
   default:
     llvm_unreachable("expected borrow scope end");
   }

lib/SIL/Utils/OwnershipLiveness.cpp

@@ -38,6 +38,10 @@ struct LinearLivenessVisitor :
     linearLiveness(linearLiveness){}
 
   bool handleUsePoint(Operand *use, UseLifetimeConstraint useConstraint) {
+    if (!linearLiveness.includeExtensions &&
+        isa<ExtendLifetimeInst>(use->getUser())) {
+      return true;
+    }
     linearLiveness.liveness.updateForUse(
       use->getUser(), useConstraint == UseLifetimeConstraint::LifetimeEnding);
     return true;
@@ -71,7 +75,9 @@ struct LinearLivenessVisitor :
   }
 };
 
-LinearLiveness::LinearLiveness(SILValue def): OSSALiveness(def) {
+LinearLiveness::LinearLiveness(SILValue def,
+                               IncludeExtensions_t includeExtensions)
+    : OSSALiveness(def), includeExtensions(includeExtensions) {
   if (def->getOwnershipKind() != OwnershipKind::Owned) {
     BorrowedValue borrowedValue(def);
     assert(borrowedValue && borrowedValue.isLocalScope());
@@ -118,7 +124,7 @@ struct InteriorLivenessVisitor :
 
   bool handleUsePoint(Operand *use, UseLifetimeConstraint useConstraint) {
     interiorLiveness.liveness.updateForUse(
-      use->getUser(), useConstraint == UseLifetimeConstraint::LifetimeEnding);
+        use->getUser(), useConstraint == UseLifetimeConstraint::LifetimeEnding);
     return true;
   }
 
@@ -302,7 +308,7 @@ struct ExtendedLinearLivenessVisitor
 
   bool handleUsePoint(Operand *use, UseLifetimeConstraint useConstraint) {
     extendedLiveness.liveness.updateForUse(
-      use->getUser(), useConstraint == UseLifetimeConstraint::LifetimeEnding);
+        use->getUser(), useConstraint == UseLifetimeConstraint::LifetimeEnding);
     return true;
   }
 

lib/SIL/Utils/OwnershipUtils.cpp

@@ -950,6 +950,10 @@ bool BorrowedValue::visitLocalScopeEndingUses(
   case BorrowedValueKind::Phi:
   case BorrowedValueKind::BeginApplyToken:
     for (auto *use : lookThroughBorrowedFromUser(value)->getUses()) {
+      if (isa<ExtendLifetimeInst>(use->getUser())) {
+        if (!visitor(use))
+          return false;
+      }
       if (use->isLifetimeEnding()) {
         if (!visitor(use))
           return false;