Merge pull request #82049 from atrick/lifedep-fix-yield

LifetimeDependenceScopeFixup: handle yielded value copies.

Author: atrick

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceScopeFixup.swift

@@ -329,7 +329,7 @@ private struct ScopeExtension {
 // violation, and that subsequent optimizations do not shrink the inner access `%a1`.
 extension ScopeExtension {
   mutating func extendScopes(dependence: LifetimeDependence) -> Bool {
-    log("Scope fixup for lifetime dependent instructions: \(dependence)")
+    log("Scope fixup for lifetime dependent instructions:\n\(dependence)")
 
     gatherExtensions(dependence: dependence)
 
@@ -1076,7 +1076,7 @@ private struct LifetimeDependentUseWalker : LifetimeDependenceDefUseWalker {
   }
 
   mutating func yieldedDependence(result: Operand) -> WalkResult {
-    return .continueWalk
+    return visitor(result)
   }
 
   mutating func storeToYieldDependence(address: Value, of operand: Operand) -> WalkResult {

SwiftCompilerSources/Sources/Optimizer/Utilities/BorrowUtils.swift

@@ -211,6 +211,9 @@ enum BorrowingInstruction : CustomStringConvertible, Hashable {
     return scopedValue
   }
 
+  /// Returns non-nil if this borrowing instruction produces an guaranteed dependent value and does not have immediate
+  /// scope-ending uses. Finding the borrow scope in such cases requires recursively following uses of the guaranteed
+  /// value.
   var dependentValue: Value? {
     switch self {
     case .borrowedFrom(let bfi):

SwiftCompilerSources/Sources/Optimizer/Utilities/LifetimeDependenceUtils.swift

@@ -834,9 +834,15 @@ extension LifetimeDependenceDefUseWalker {
       return walkDownUses(of: bfi, using: operand)
     case let .storeBorrow(sbi):
       return walkDownAddressUses(of: sbi)
-    case .beginApply:
-      // Skip the borrow scope; the type system enforces non-escapable
-      // arguments.
+    case let .beginApply(bai):
+      // First, visit the uses of any non-Escapable yields. The yielded value may be copied and used outside the
+      // coroutine scope.  Now, visit the uses of the begin_apply token. This adds the coroutine scope itself to the
+      for yield in bai.yieldedValues {
+        if walkDownUses(of: yield, using: operand) == .abortWalk {
+          return .abortWalk
+        }
+      }
+      // lifetime to account for the scope of any arguments.
       return visitInnerBorrowUses(of: borrowInst, operand: operand)
     case .partialApply, .markDependence:
       fatalError("OwnershipUseVisitor should bypass partial_apply [on_stack] "

test/SILOptimizer/lifetime_dependence/lifetime_dependence_util.sil

@@ -252,6 +252,7 @@ bb0(%0 : @guaranteed $NE):
 
 // CHECK-LABEL: dependence_coroutine: lifetime_dependence_use with: %0
 // CHECK: LifetimeDependence uses of: %0 = argument of bb0 : $NE
+// CHECK-NEXT: Leaf use: operand #0 of   destroy_value %{{.*}} : $NE
 // CHECK-NEXT: Leaf use: operand #0 of   %{{.*}} = end_apply %{{.*}} as $()
 // CHECK-NEXT: dependence_coroutine: lifetime_dependence_use with: %0
 sil [ossa] @dependence_coroutine : $@convention(thin) (@guaranteed NE) -> () {

test/SILOptimizer/lifetime_dependence/scope_fixup.sil

@@ -72,6 +72,9 @@ sil @useNE : $@convention(thin) (@guaranteed NE) -> ()
 
 sil [ossa] @Wrapper_init : $@convention(method) (@owned NE, @thin Wrapper.Type) -> @lifetime(copy 0) @owned Wrapper
 
+sil @yieldWrapper : $@yield_once @convention(method) (@guaranteed NCContainer) -> @lifetime(borrow 0) @yields @guaranteed Wrapper
+sil @yieldNE : $@yield_once @convention(method) (@guaranteed Wrapper) -> @lifetime(copy 0) @yields @guaranteed NE
+
 sil [ossa] @NCContainer_ne_read : $@yield_once @convention(method) (@guaranteed NCContainer) -> @lifetime(borrow 0) @yields @guaranteed NE
 
 sil @yieldRawSpan : $@yield_once @convention(method) (@in_guaranteed A) -> @lifetime(borrow address_for_deps 0) @yields @guaranteed RawSpan
@@ -554,3 +557,59 @@ bb0(%0 : $*Holder):
   destroy_value %holder
   return %md
 }
+
+// Test a nested dependent coroutine. The inner coroutine yields a
+// value (NE) that inherits its dependency from the outer yield
+// (Wrapper). All copies of the inner yield must be used within the
+// scope of the outer coroutine. The outer end_apply must be extended
+// below the call to 'useNE'.
+//
+// CHECK-LABEL: sil hidden [ossa] @testCopiedYield : $@convention(thin) (@owned NCContainer) -> () {
+// CHECK: bb0(%0 : @owned $NCContainer):
+// CHECK:   [[BOX:%[0-9]+]] = alloc_box ${ var NCContainer }, var, name "nc"
+// CHECK:   [[ACCESS:%[0-9]+]] = begin_access [read] [unknown]
+// CHECK:   [[LB:%[0-9]+]] = load_borrow [unchecked]
+// CHECK:   ([[YIELD:%[0-9]+]], [[TOKEN:%[0-9]+]]) = begin_apply %7(%6) : $@yield_once @convention(method) (@guaranteed NCContainer) -> @lifetime(borrow 0) @yields @guaranteed Wrapper
+// CHECK:   mark_dependence [unresolved] [[YIELD]] on [[TOKEN]]
+// CHECK:   ([[NE:%[0-9]+]], %{{.*}}) = begin_apply %13(%12) : $@yield_once @convention(method) (@guaranteed Wrapper) -> @lifetime(copy 0) @yields @guaranteed NE
+// CHECK:   destroy_value
+// CHECK:   apply %{{.*}}(%{{.*}}) : $@convention(thin) (@guaranteed NE) -> ()
+// CHECK:   destroy_value %{{.*}}
+// CHECK:   end_apply [[TOKEN]] as $()
+// CHECK:   end_borrow [[LB]]
+// CHECK:   end_access [[ACCESS]]
+// CHECK:   destroy_value [[BOX]]
+// CHECK-LABEL: } // end sil function 'testCopiedYield'
+sil hidden [ossa] @testCopiedYield : $@convention(thin) (@owned NCContainer) -> () {
+bb0(%0 : @owned $NCContainer):
+  %1 = alloc_box ${ var NCContainer }, var, name "nc"
+  %2 = begin_borrow [lexical] [var_decl] %1
+  %3 = project_box %2, 0
+  store %0 to [init] %3
+  %5 = begin_access [read] [unknown] %3
+  %7 = load_borrow [unchecked] %5
+
+  %8 = function_ref @yieldWrapper : $@yield_once @convention(method) (@guaranteed NCContainer) -> @lifetime(borrow 0) @yields @guaranteed Wrapper
+  (%9, %10) = begin_apply %8(%7) : $@yield_once @convention(method) (@guaranteed NCContainer) -> @lifetime(borrow 0) @yields @guaranteed Wrapper
+  %11 = mark_dependence [unresolved] %9 on %10
+  %12 = copy_value %11
+  %13 = end_apply %10 as $()
+  end_borrow %7
+  end_access %5
+  %16 = move_value [var_decl] %12
+
+  %18 = function_ref @yieldNE : $@yield_once @convention(method) (@guaranteed Wrapper) -> @lifetime(copy 0) @yields @guaranteed NE
+  (%19, %20) = begin_apply %18(%16) : $@yield_once @convention(method) (@guaranteed Wrapper) -> @lifetime(copy 0) @yields @guaranteed NE
+  %21 = copy_value %19
+  %22 = end_apply %20 as $()
+  destroy_value %16
+  %23 = move_value [var_decl] %21
+
+  %25 = function_ref @useNE : $@convention(thin) (@guaranteed NE) -> ()
+  %26 = apply %25(%23) : $@convention(thin) (@guaranteed NE) -> ()
+  destroy_value %23
+  end_borrow %2
+  destroy_value %1
+  %31 = tuple ()
+  return %31
+}