Merge pull request #78034 from nate-chandler/rdar141037060

nate-chandler · web-flow · commit 159259acff70 · 2024-12-09T08:25:41.000-08:00
[OSSACompleteLifetime] Handle scoped addresses.
diff --git a/include/swift/SIL/OSSALifetimeCompletion.h b/include/swift/SIL/OSSALifetimeCompletion.h
@@ -90,20 +90,29 @@ class OSSALifetimeCompletion {
   ///
   /// Returns true if any new instructions were created to complete the
   /// lifetime.
-  ///
-  /// TODO: We also need to complete scoped addresses (e.g. store_borrow)!
   LifetimeCompletion completeOSSALifetime(SILValue value, Boundary boundary) {
-    if (value->getOwnershipKind() == OwnershipKind::None)
-      return LifetimeCompletion::NoLifetime;
-
-    if (value->getOwnershipKind() != OwnershipKind::Owned) {
+    switch (value->getOwnershipKind()) {
+    case OwnershipKind::None: {
+      auto scopedAddress = ScopedAddressValue(value);
+      if (!scopedAddress)
+        return LifetimeCompletion::NoLifetime;
+      break;
+    }
+    case OwnershipKind::Owned:
+      break;
+    case OwnershipKind::Any:
+      llvm::report_fatal_error("value with any ownership kind!?");
+    case OwnershipKind::Guaranteed:
+    case OwnershipKind::Unowned: {
       BorrowedValue borrowedValue(value);
       if (!borrowedValue)
         return LifetimeCompletion::NoLifetime;
 
       if (!borrowedValue.isLocalScope())
         return LifetimeCompletion::AlreadyComplete;
     }
+    }
+
     if (!completedValues.insert(value))
       return LifetimeCompletion::AlreadyComplete;
 
diff --git a/include/swift/SIL/ScopedAddressUtils.h b/include/swift/SIL/ScopedAddressUtils.h
@@ -122,7 +122,8 @@ struct ScopedAddressValue {
   AddressUseKind updateTransitiveLiveness(SSAPrunedLiveness &liveness) const;
 
   /// Create appropriate scope ending instruction at \p insertPt.
-  void createScopeEnd(SILBasicBlock::iterator insertPt, SILLocation loc) const;
+  SILInstruction *createScopeEnd(SILBasicBlock::iterator insertPt,
+                                 SILLocation loc) const;
 
   /// Create scope ending instructions at \p liveness boundary.
   void endScopeAtLivenessBoundary(SSAPrunedLiveness *liveness) const;
diff --git a/lib/SIL/Utils/OSSALifetimeCompletion.cpp b/lib/SIL/Utils/OSSALifetimeCompletion.cpp
@@ -76,6 +76,9 @@ static SILInstruction *endOSSALifetime(SILValue value,
     }
     return builder.createDestroyValue(loc, value, DontPoisonRefs, isDeadEnd);
   }
+  if (auto scopedAddress = ScopedAddressValue(value)) {
+    return scopedAddress.createScopeEnd(builder.getInsertionPoint(), loc);
+  }
   return builder.createEndBorrow(loc, lookThroughBorrowedFromUser(value));
 }
 
@@ -419,36 +422,49 @@ static bool endLifetimeAtAvailabilityBoundary(SILValue value,
   return changed;
 }
 
+static bool endLifetimeAtBoundary(SILValue value,
+                                  SSAPrunedLiveness const &liveness,
+                                  OSSALifetimeCompletion::Boundary boundary,
+                                  DeadEndBlocks &deadEndBlocks) {
+  bool changed = false;
+  switch (boundary) {
+  case OSSALifetimeCompletion::Boundary::Liveness:
+    changed |= endLifetimeAtLivenessBoundary(value, liveness, deadEndBlocks);
+    break;
+  case OSSALifetimeCompletion::Boundary::Availability:
+    changed |=
+        endLifetimeAtAvailabilityBoundary(value, liveness, deadEndBlocks);
+    break;
+  }
+  return changed;
+}
+
 /// End the lifetime of \p value at unreachable instructions.
 ///
 /// Returns true if any new instructions were created to complete the lifetime.
 bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(SILValue value,
                                                       Boundary boundary) {
+  if (auto scopedAddress = ScopedAddressValue(value)) {
+    SmallVector<SILBasicBlock *, 8> discoveredBlocks;
+    SSAPrunedLiveness liveness(value->getFunction(), &discoveredBlocks);
+    scopedAddress.computeTransitiveLiveness(liveness);
+    return endLifetimeAtBoundary(value, liveness, boundary, deadEndBlocks);
+  }
+
   // Called for inner borrows, inner adjacent reborrows, inner reborrows, and
   // scoped addresses.
   auto handleInnerScope = [this, boundary](SILValue innerBorrowedValue) {
     completeOSSALifetime(innerBorrowedValue, boundary);
   };
   InteriorLiveness liveness(value);
   liveness.compute(domInfo, handleInnerScope);
-
-  bool changed = false;
-  switch (boundary) {
-  case Boundary::Liveness:
-    changed |= endLifetimeAtLivenessBoundary(value, liveness.getLiveness(),
-                                             deadEndBlocks);
-    break;
-  case Boundary::Availability:
-    changed |= endLifetimeAtAvailabilityBoundary(value, liveness.getLiveness(),
-                                                 deadEndBlocks);
-    break;
-  }
   // TODO: Rebuild outer adjacent phis on demand (SILGen does not currently
   // produce guaranteed phis). See FindEnclosingDefs &
   // findSuccessorDefsFromPredDefs. If no enclosing phi is found, we can create
   // it here and use updateSSA to recursively populate phis.
   assert(liveness.getUnenclosedPhis().empty());
-  return changed;
+  return endLifetimeAtBoundary(value, liveness.getLiveness(), boundary,
+                               deadEndBlocks);
 }
 
 namespace swift::test {
diff --git a/lib/SIL/Utils/ScopedAddressUtils.cpp b/lib/SIL/Utils/ScopedAddressUtils.cpp
@@ -151,16 +151,15 @@ AddressUseKind ScopedAddressValue::updateTransitiveLiveness(
   return addressKind;
 }
 
-void ScopedAddressValue::createScopeEnd(SILBasicBlock::iterator insertPt,
-                                        SILLocation loc) const {
+SILInstruction *
+ScopedAddressValue::createScopeEnd(SILBasicBlock::iterator insertPt,
+                                   SILLocation loc) const {
   switch (kind) {
   case ScopedAddressValueKind::StoreBorrow: {
-    SILBuilderWithScope(insertPt).createEndBorrow(loc, value);
-    return;
+    return SILBuilderWithScope(insertPt).createEndBorrow(loc, value);
   }
   case ScopedAddressValueKind::BeginAccess: {
-    SILBuilderWithScope(insertPt).createEndAccess(loc, value, false);
-    return;
+    return SILBuilderWithScope(insertPt).createEndAccess(loc, value, false);
   }
   case ScopedAddressValueKind::Invalid:
     llvm_unreachable("Using invalid case?!");
diff --git a/test/SILOptimizer/ossa_lifetime_completion.sil b/test/SILOptimizer/ossa_lifetime_completion.sil
@@ -792,3 +792,52 @@ bb3:
   return %0 : $C
 }
 
+// CHECK-LABEL: begin running test {{.*}} on store_borrow: ossa_lifetime_completion
+// CHECK-LABEL: sil [ossa] @store_borrow : {{.*}} {
+// CHECK:         [[TOKEN:%[^,]+]] = store_borrow
+// CHECK:         cond_br undef, {{bb[0-9]+}}, [[DIE:bb[0-9]+]]
+// CHECK:       [[DIE]]:
+// CHECK-NEXT:    end_borrow [[TOKEN]]
+// CHECK-NEXT:    unreachable
+// CHECK-LABEL: } // end sil function 'store_borrow'
+// CHECK-LABEL: end running test {{.*}} on store_borrow: ossa_lifetime_completion
+sil [ossa] @store_borrow : $@convention(thin) (@guaranteed C) -> () {
+entry(%instance : @guaranteed $C):
+  specify_test "ossa_lifetime_completion %token availability"
+  %addr = alloc_stack $C
+  %token = store_borrow %instance to %addr : $*C
+  cond_br undef, exit, die
+
+exit:
+  end_borrow %token : $*C
+  dealloc_stack %addr : $*C
+  %retval = tuple ()
+  return %retval : $()
+die:
+  unreachable
+}
+
+// CHECK-LABEL: begin running test {{.*}} on begin_access: ossa_lifetime_completion
+// CHECK-LABEL: sil [ossa] @begin_access : {{.*}} {
+// CHECK:         [[TOKEN:%[^,]+]] = begin_access
+// CHECK:         cond_br undef, {{bb[0-9]+}}, [[DIE:bb[0-9]+]]
+// CHECK:       [[DIE]]:
+// CHECK-NEXT:    end_access [[TOKEN]]
+// CHECK-NEXT:    unreachable
+// CHECK-LABEL: } // end sil function 'begin_access'
+// CHECK-LABEL: end running test {{.*}} on begin_access: ossa_lifetime_completion
+sil [ossa] @begin_access : $@convention(thin) (@guaranteed C) -> () {
+entry(%instance : @guaranteed $C):
+  specify_test "ossa_lifetime_completion %access availability"
+  %addr = alloc_stack $C
+  %access = begin_access [static] [read] %addr : $*C
+  cond_br undef, exit, die
+
+exit:
+  end_access %access : $*C
+  dealloc_stack %addr : $*C
+  %retval = tuple ()
+  return %retval : $()
+die:
+  unreachable
+}
diff --git a/test/SILOptimizer/silgen_cleanup_complete_ossa.sil b/test/SILOptimizer/silgen_cleanup_complete_ossa.sil
@@ -11,6 +11,8 @@ class Klass {
 }
 class SubKlass : Klass {}
 
+class C {}
+
 enum FakeOptional<T> {
 case none
 case some(T)
@@ -32,14 +34,15 @@ protocol P : AnyObject {}
 
 sil @unreachableHandler : $@convention(thin) () -> ()
 
-// CHECK-LABEL: sil [ossa] @testCompleteOSSALifetimes : $@convention(thin) (@owned FakeOptional<Klass>) -> () {
+// CHECK-LABEL: sil [ossa] @testCompleteOSSALifetimes : {{.*}} {
 // CHECK:   [[BOX:%.*]] = alloc_box ${ var FakeOptional<Klass> }, var, name "c"
 // CHECK:   [[BORROW:%.,*]] = begin_borrow [lexical] [[BOX]] : ${ var FakeOptional<Klass> }
 // CHECK: bb2:
 // CHECK:   apply
 // CHECK:   end_borrow [[BORROW]] : ${ var FakeOptional<Klass> }
 // CHECK:   dealloc_box [dead_end] [[BOX]] : ${ var FakeOptional<Klass> }
 // CHECK:   unreachable
+// CHECK-LABEL: } // end sil function 'testCompleteOSSALifetimes'
 sil [ossa] @testCompleteOSSALifetimes : $@convention(thin) (@owned FakeOptional<Klass>) -> () {
 bb0(%0 : @owned $FakeOptional<Klass>):
   %box = alloc_box ${ var FakeOptional<Klass> }, var, name "c"
@@ -73,11 +76,68 @@ bb5:
   return %36 : $()
 }
 
-// CHECK-LABEL: sil [ossa] @testExistentialLifetime : $@convention(thin) (@owned any P) -> @owned AnyObject {
+// CHECK-LABEL: sil [ossa] @testExistentialLifetime : {{.*}} {
 // CHECK-NOT: destroy
+// CHECK-LABEL: } // end sil function 'testExistentialLifetime'
 sil [ossa] @testExistentialLifetime : $@convention(thin) (@owned any P) -> @owned AnyObject {
 bb0(%0 : @owned $any P):
   %1 = open_existential_ref %0 : $any P to $@opened("34B79428-2E49-11ED-901A-8AC134504E1C", any P) Self
   %2 = init_existential_ref %1 : $@opened("34B79428-2E49-11ED-901A-8AC134504E1C", any P) Self : $@opened("34B79428-2E49-11ED-901A-8AC134504E1C", any P) Self, $AnyObject
   return %2 : $AnyObject
 }
+
+// CHECK-LABEL: sil [ossa] @store_borrow : {{.*}} {
+// CHECK:       bb0([[INSTANCE:%[^,]+]] :
+// CHECK:         [[ADDR:%[^,]+]] = alloc_stack $C
+// CHECK:         [[TOKEN:%[^,]+]] = store_borrow [[INSTANCE]] to [[ADDR]]
+// CHECK:         [[LOAD:%[^,]+]] = load_borrow [[TOKEN]]
+// CHECK:         cond_br undef, {{bb[0-9]+}}, [[DIE:bb[0-9]+]]
+// CHECK:       [[DIE]]:
+// CHECK:         end_borrow [[LOAD]]
+// CHECK:         end_borrow [[TOKEN]]
+// CHECK:         destroy_value [dead_end] [[INSTANCE]]
+// CHECK:         unreachable
+// CHECK-LABEL: } // end sil function 'store_borrow'
+sil [ossa] @store_borrow : $@convention(thin) (@owned C) -> () {
+entry(%instance : @owned $C):
+  %addr = alloc_stack $C
+  %token = store_borrow %instance to %addr : $*C
+  %load = load_borrow %token : $*C
+  cond_br undef, exit, die
+exit:
+  end_borrow %load : $C
+  end_borrow %token : $*C
+  dealloc_stack %addr : $*C
+  apply undef(%instance) : $@convention(thin) (@guaranteed C) -> ()
+  destroy_value %instance : $C
+  %retval = tuple ()
+  return %retval : $()
+die:
+  unreachable
+}
+
+// CHECK-LABEL: sil [ossa] @begin_access : {{.*}} {
+// CHECK:         [[ADDR:%[^,]+]] = alloc_stack $C
+// CHECK:         [[ACCESS:%[^,]+]] = begin_access [modify] [static] [[ADDR]]
+// CHECK:         cond_br undef, {{bb[0-9]+}}, [[DIE:bb[0-9]+]]
+// CHECK:       [[DIE]]:
+// CHECK:         end_access [[ACCESS]]
+// CHECK:         unreachable
+// CHECK-LABEL: } // end sil function 'begin_access'
+sil [ossa] @begin_access : $@convention(thin) () -> () {
+entry:
+  %addr2 = alloc_stack $C
+  %access = begin_access [static] [modify] %addr2 : $*C
+  apply undef(%access) : $@convention(thin) () -> (@out C)
+  destroy_addr %access : $*C
+  cond_br undef, exit, die
+
+exit:
+  end_access %access : $*C
+  dealloc_stack %addr2 : $*C
+  %retval = tuple ()
+  return %retval : $()
+
+die:
+  unreachable
+}
