SILGenCleanup: extend to handle trivial local var scopes

Improves OSSALifetimeCompletion to handle trivial variables when running from
SILGenCleanup. This only affects lifetime dependence diagnostics.

For the purpose of lifetime diagnostics, trivial local variables are only valid
within their lexical scope. Sadly, SILGen only know how to insert cleanup code
on normal function exits. SILGenCleanup relies on lifetime completion to fix
lifetimes on dead end paths.

  %var = move_value [var_decl]
  try_apply %f() : $..., normal bb1, error error
error:
  extend_lifetime %var <=== insert this
  unreachable

This allows Span to depend on local unsafe pointers AND be used within
throwing closures:

    _ = a.withUnsafeBufferPointer {
      let buffer = $0
      let view = Span(_unsafeElements: buffer)
      return view.withUnsafeBufferPointer(\.count)
    }

Author: atrick

include/swift/SIL/OSSALifetimeCompletion.h

@@ -39,6 +39,10 @@ class DeadEndBlocks;
 enum class LifetimeCompletion { NoLifetime, AlreadyComplete, WasCompleted };
 
 class OSSALifetimeCompletion {
+public:
+  enum HandleTrivialVariable_t { IgnoreTrivialVariable, ExtendTrivialVariable };
+
+private:
   // If domInfo is nullptr, then InteriorLiveness never assumes dominance. As a
   // result it may report extra unenclosedPhis. In that case, any attempt to
   // create a new phi would result in an immediately redundant phi.
@@ -50,11 +54,15 @@ class OSSALifetimeCompletion {
   // recomputing their lifetimes.
   ValueSet completedValues;
 
+  // Extend trivial variables for lifetime diagnostics (only in SILGenCleanup).
+  HandleTrivialVariable_t handleTrivialVariable;
+
 public:
   OSSALifetimeCompletion(SILFunction *function, const DominanceInfo *domInfo,
-                         DeadEndBlocks &deadEndBlocks)
+                         DeadEndBlocks &deadEndBlocks,
+                         HandleTrivialVariable_t handleTrivialVariable = IgnoreTrivialVariable)
       : domInfo(domInfo), deadEndBlocks(deadEndBlocks),
-        completedValues(function) {}
+        completedValues(function), handleTrivialVariable(handleTrivialVariable) {}
 
   // The kind of boundary at which to complete the lifetime.
   //
@@ -93,10 +101,15 @@ class OSSALifetimeCompletion {
   LifetimeCompletion completeOSSALifetime(SILValue value, Boundary boundary) {
     switch (value->getOwnershipKind()) {
     case OwnershipKind::None: {
-      auto scopedAddress = ScopedAddressValue(value);
-      if (!scopedAddress)
-        return LifetimeCompletion::NoLifetime;
-      break;
+      if (auto scopedAddress = ScopedAddressValue(value)) {
+        break;
+      }
+      // During SILGenCleanup, extend move_value [var_decl].
+      if (handleTrivialVariable == ExtendTrivialVariable
+          && value->isFromVarDecl()) {
+        break;
+      }
+      return LifetimeCompletion::NoLifetime;
     }
     case OwnershipKind::Owned:
       break;

include/swift/SIL/OwnershipUseVisitor.h

@@ -165,6 +165,8 @@ class OwnershipUseVisitor : OwnershipUseVisitorBase<Impl> {
 protected:
   bool visitConsumes(SILValue ssaDef);
 
+  bool visitExtends(SILValue ssaDef);
+
   bool visitOuterBorrow(SILValue borrowBegin);
 
   bool visitOuterBorrowScopeEnd(Operand *borrowEnd);
@@ -202,8 +204,12 @@ bool OwnershipUseVisitor<Impl>::visitLifetimeEndingUses(SILValue ssaDef) {
   case OwnershipKind::Guaranteed:
     return visitOuterBorrow(ssaDef);
 
-  case OwnershipKind::Any:
   case OwnershipKind::None:
+    if (ssaDef->isFromVarDecl()) {
+      return visitExtends(ssaDef);
+    }
+    LLVM_FALLTHROUGH;
+  case OwnershipKind::Any:
   case OwnershipKind::Unowned:
     llvm_unreachable("requires an owned or guaranteed orignalDef");
   }
@@ -231,6 +237,18 @@ bool OwnershipUseVisitor<Impl>::visitConsumes(SILValue ssaDef) {
   return true;
 }
 
+template <typename Impl>
+bool OwnershipUseVisitor<Impl>::visitExtends(SILValue ssaDef) {
+  for (Operand *use : ssaDef->getUses()) {
+    if (isa<ExtendLifetimeInst>(use->getUser())) {
+      if (!handleUsePoint(use, UseLifetimeConstraint::NonLifetimeEnding))
+        return false;
+      continue;
+    }
+  }
+  return true;
+}
+
 template <typename Impl>
 bool OwnershipUseVisitor<Impl>::visitOuterBorrow(SILValue borrowBegin) {
   BorrowedValue borrow(borrowBegin);

lib/SIL/Utils/OSSALifetimeCompletion.cpp

@@ -80,6 +80,9 @@ static SILInstruction *endOSSALifetime(SILValue value,
   if (auto scopedAddress = ScopedAddressValue(value)) {
     return scopedAddress.createScopeEnd(builder.getInsertionPoint(), loc);
   }
+  if (value->getOwnershipKind() == OwnershipKind::None) {
+    return builder.createExtendLifetime(loc, value);
+  }
   return builder.createEndBorrow(loc, lookThroughBorrowedFromUser(value));
 }
 
@@ -297,9 +300,16 @@ void AvailabilityBoundaryVisitor::computeRegion(
     regionWorklist.push(block);
   };
 
-  for (auto *endBlock : boundary.endBlocks) {
-    if (!consumingBlocks.contains(endBlock)) {
-      collect(endBlock);
+  // Trivial values that correspond to local variables (as opposed to
+  // ScopedAddresses) are available only up to their last extend_lifetime on
+  // non-dead-end paths. They cannot be consumed, but are only "available" up to
+  // the end of their scope.
+  if (value->getOwnershipKind() != OwnershipKind::None
+      || ScopedAddressValue(value)) {
+    for (auto *endBlock : boundary.endBlocks) {
+      if (!consumingBlocks.contains(endBlock)) {
+        collect(endBlock);
+      }
     }
   }
   for (SILBasicBlock *edge : boundary.boundaryEdges) {
@@ -497,12 +507,21 @@ bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(SILValue value,
   if (auto scopedAddress = ScopedAddressValue(value)) {
     return analyzeAndUpdateLifetime(scopedAddress, boundary);
   }
-
   // Called for inner borrows, inner adjacent reborrows, inner reborrows, and
   // scoped addresses.
   auto handleInnerScope = [this, boundary](SILValue innerBorrowedValue) {
     completeOSSALifetime(innerBorrowedValue, boundary);
   };
+  if (value->getOwnershipKind() == OwnershipKind::None) {
+    // Trivial variable lifetimes are only relevant up to the extend_lifetime
+    // instructions emitted by SILGen. Their other uses have no meaning with
+    // respect to lifetime. The only purpose of "completing" their lifetime is
+    // to insert extend_lifetime on dead-end blocks.
+    LinearLiveness liveness(value);
+    liveness.compute();
+    return endLifetimeAtBoundary(value, liveness.getLiveness(), boundary,
+                                 deadEndBlocks);
+  }
   InteriorLiveness liveness(value);
   liveness.compute(domInfo, handleInnerScope);
   // TODO: Rebuild outer adjacent phis on demand (SILGen does not currently

lib/SIL/Utils/OwnershipLiveness.cpp

@@ -79,10 +79,21 @@ struct LinearLivenessVisitor :
 LinearLiveness::LinearLiveness(SILValue def,
                                IncludeExtensions_t includeExtensions)
     : OSSALiveness(def), includeExtensions(includeExtensions) {
-  if (def->getOwnershipKind() != OwnershipKind::Owned) {
+  switch (def->getOwnershipKind()) {
+  case OwnershipKind::Owned:
+    break;
+  case OwnershipKind::Guaranteed: {
     BorrowedValue borrowedValue(def);
     assert(borrowedValue && borrowedValue.isLocalScope());
     (void)borrowedValue;
+    break;
+  }
+  case OwnershipKind::None:
+    assert(def->isFromVarDecl());
+    break;
+  case OwnershipKind::Unowned:
+  case OwnershipKind::Any:
+    llvm_unreachable("bad ownership for LinearLiveness");
   }
 }
 

lib/SILOptimizer/Mandatory/SILGenCleanup.cpp

@@ -267,7 +267,9 @@ bool SILGenCleanup::completeOSSALifetimes(SILFunction *function) {
   }
 
   bool changed = false;
-  OSSALifetimeCompletion completion(function, /*DomInfo*/ nullptr, *deba);
+  OSSALifetimeCompletion completion(
+    function, /*DomInfo*/ nullptr, *deba,
+    OSSALifetimeCompletion::ExtendTrivialVariable);
   BasicBlockSet completed(function);
   for (auto *root : roots) {
     if (root == function->getEntryBlock()) {

test/SILOptimizer/lifetime_dependence/semantics.swift

@@ -228,6 +228,19 @@ func testTrivialScope<T>(a: Array<T>) -> Span<T> {
   // expected-note  @-3{{this use causes the lifetime-dependent value to escape}}
 }
 
+extension Span {
+  public func withThrowingClosure<E: Error>(_ body: () throws(E) -> ()) throws(E) -> () {
+    try body()
+  }
+}
+
+// Test dependence on an local variable that needs to be extended into the dead-end block of a never-throwing apply.
+public func test(p: UnsafePointer<Int>) {
+  let pointer = p
+  let span = Span(base: pointer, count: 1)
+  span.withThrowingClosure {}
+}
+
 // =============================================================================
 // Scoped dependence on property access
 // =============================================================================

test/SILOptimizer/silgen_cleanup_complete_ossa.sil

@@ -6,6 +6,8 @@ sil_stage raw
 
 typealias AnyObject = Builtin.AnyObject
 
+protocol Error {}
+
 class Klass {
   var property: Builtin.Int64
 }
@@ -28,6 +30,12 @@ struct UInt8 {
 
 protocol P : AnyObject {}
 
+struct Err: Error {
+  var i: Int
+}
+
+sil @throwing : $@convention(thin) (UInt8) -> @error_indirect Err
+
 // =============================================================================
 // Test complete OSSA lifetimes
 // =============================================================================
@@ -367,3 +375,33 @@ left:
 right:
   unreachable
 }
+
+// CHECK-LABEL: sil [ossa] @testExtendTrivialToDeadEnd : $@convention(thin) (UInt8) -> () {
+// CHECK: bb0(%0 : $UInt8):
+// CHECK:   [[MV:%.*]] = move_value [var_decl] %0 : $UInt8
+// CHECK:   try_apply %{{.*}}(%{{.*}}, [[MV]]) : $@convention(thin) (UInt8) -> @error_indirect Err, normal bb1, error bb2
+// CHECK: bb1(
+// CHECK:   extend_lifetime [[MV]] : $UInt8
+// CHECK:   return
+// CHECK: bb2:
+// CHECK:   dealloc_stack
+// CHECK:   extend_lifetime [[MV]] : $UInt8
+// CHECK:   unreachable
+// CHECK-LABEL: } // end sil function 'testExtendTrivialToDeadEnd'
+sil [ossa] @testExtendTrivialToDeadEnd : $@convention(thin) (UInt8) -> () {
+bb0(%0 : $UInt8):
+  %mv = move_value [var_decl] %0
+  %e = alloc_stack $Err
+  %f = function_ref @throwing : $@convention(thin) (UInt8) -> @error_indirect Err
+  try_apply %f(%e, %mv) : $@convention(thin) (UInt8) -> @error_indirect Err, normal bb1, error bb2
+
+bb1(%ret : $()):
+  extend_lifetime %mv
+  dealloc_stack %e : $*Err
+  %99 = tuple ()
+  return %99
+
+bb2:
+  dealloc_stack %e : $*Err
+  unreachable
+}