LifetimeDependenceScopeFixup: extend store_borrow scopes

This is require to handle '@'_addressable arguments on the caller side.

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceScopeFixup.swift

@@ -249,7 +249,7 @@ private struct ScopeExtension {
 
 private extension LifetimeDependence.Scope {
   /// The instruction that introduces an extendable scope. This returns a non-nil scope introducer for
-  /// Extendable.nestedScopes.
+  /// ScopeExtension.nestedScopes.
   var extendableBegin: Instruction? {
     switch self {
     case let .access(beginAccess):
@@ -258,6 +258,17 @@ private extension LifetimeDependence.Scope {
       return beginBorrow.value.definingInstruction!
     case let .yield(yieldedValue):
       return yieldedValue.definingInstruction!
+    case let .initialized(initializer):
+      switch initializer {
+      case let .store(initializingStore: store, initialAddress: _):
+        if let sb = store as? StoreBorrowInst {
+          return sb
+        }
+        return nil
+      case .argument, .yield:
+        // TODO: extend indirectly yielded scopes.
+        return nil
+      }
     default:
       return nil
     }
@@ -277,17 +288,8 @@ private extension LifetimeDependence.Scope {
       let accessExtension = gatherAccessExtension(beginAccess: beginAccess, innerScopes: &innerScopes)
       return SingleInlineArray(element: accessExtension)
     case let .borrowed(beginBorrow):
-      let borrowedValue = beginBorrow.baseOperand!.value
+      return gatherBorrowExtension(borrowedValue: beginBorrow.baseOperand!.value, innerScopes: &innerScopes, context)
-      let enclosingScope = LifetimeDependence.Scope(base: borrowedValue, context)
+
-      innerScopes.push(self)
-      var innerBorrowScopes = innerScopes
-      innerBorrowScopes.push(enclosingScope)
-      if let extensions = enclosingScope.gatherExtensions(innerScopes: innerBorrowScopes, context) {
-        return extensions
-      }
-      // This is the outermost scope to be extended because gatherExtensions did not find an enclosing scope.
-      return SingleInlineArray(element: getOuterExtension(owner: enclosingScope.parentValue, nestedScopes: innerScopes,
-                                                          context))
     case let .yield(yieldedValue):
       innerScopes.push(self)
       var extensions = SingleInlineArray<ScopeExtension>()
@@ -300,6 +302,17 @@ private extension LifetimeDependence.Scope {
         extensions.append(contentsOf: gatherOperandExtension(on: operand, innerScopes: innerScopes, context))
       }
       return extensions
+    case let .initialized(initializer):
+      switch initializer {
+      case let .store(initializingStore: store, initialAddress: _):
+        if let sb = store as? StoreBorrowInst {
+          return gatherBorrowExtension(borrowedValue: sb.source, innerScopes: &innerScopes, context)
+        }
+        return nil
+      case .argument, .yield:
+        // TODO: extend indirectly yielded scopes.
+        return nil
+      }
     default:
       return nil
     }
@@ -360,6 +373,23 @@ private extension LifetimeDependence.Scope {
     }
     return ScopeExtension(owner: outerBeginAccess, nestedScopes: innerScopes, dependsOnArg: nil)
   }
+
+  func gatherBorrowExtension(borrowedValue: Value,
+                             innerScopes: inout SingleInlineArray<LifetimeDependence.Scope>,
+                             _ context: FunctionPassContext)
+    -> SingleInlineArray<ScopeExtension> {
+
+    let enclosingScope = LifetimeDependence.Scope(base: borrowedValue, context)
+    innerScopes.push(self)
+    var innerBorrowScopes = innerScopes
+    innerBorrowScopes.push(enclosingScope)
+    if let extensions = enclosingScope.gatherExtensions(innerScopes: innerBorrowScopes, context) {
+      return extensions
+    }
+    // This is the outermost scope to be extended because gatherExtensions did not find an enclosing scope.
+    return SingleInlineArray(element: getOuterExtension(owner: enclosingScope.parentValue, nestedScopes: innerScopes,
+                                                        context))
+  }
 }
 
 /// Compute the range of the a scope owner. Nested scopes must stay within this range.
@@ -584,6 +614,17 @@ private extension LifetimeDependence.Scope {
     case let .yield(yieldedValue):
       let beginApply = yieldedValue.definingInstruction as! BeginApplyInst
       return beginApply.createEnd(builder, context)
+    case let .initialized(initializer):
+      switch initializer {
+      case let .store(initializingStore: store, initialAddress: _):
+        if let sb = store as? StoreBorrowInst {
+          return builder.createEndBorrow(of: sb)
+        }
+        return nil
+      case .argument, .yield:
+        // TODO: extend indirectly yielded scopes.
+        return nil
+      }
     default:
       return nil
     }

test/SILOptimizer/lifetime_dependence/scope_fixup.sil

@@ -1,11 +1,15 @@
-// RUN: %target-sil-opt %s \
+// RUN: %target-sil-opt \
-// RUN:   --lifetime-dependence-scope-fixup \
+// RUN:   -lifetime-dependence-scope-fixup \
 // RUN:   -sil-verify-all \
 // RUN:   -enable-experimental-feature LifetimeDependence \
-// RUN:   2>&1 | %FileCheck %s
+// RUN:   -enable-experimental-feature AddressableParameters \
+// RUN:   -enable-experimental-feature AddressableTypes \
+// RUN:   %s | %FileCheck %s
 
 // REQUIRES: swift_in_compiler
 // REQUIRES: swift_feature_LifetimeDependence
+// REQUIRES: swift_feature_AddressableParameters
+// REQUIRES: swift_feature_AddressableTypes
 
 // Test the SIL representation for lifetime dependence scope fixup.
 
@@ -15,6 +19,10 @@ import Builtin
 import Swift
 
 struct NE : ~Escapable {
+  var p: UnsafeRawPointer
+
+  @lifetime(immortal)
+  init()
 }
 
 struct Wrapper : ~Escapable {
@@ -31,6 +39,22 @@ struct NCContainer : ~Copyable {
   var wrapper: Wrapper { get } // _read
 }
 
+struct TrivialHolder {
+  var pointer: UnsafeRawPointer
+}
+
+struct Holder {
+  var object: AnyObject
+}
+
+@_addressableForDependencies
+struct AddressableForDeps {}
+
+sil @getPtr : $@convention(thin) () -> @out UnsafeRawPointer
+sil @getSpan : $@convention(thin) (@in_guaranteed AnyObject) -> @lifetime(borrow 0) @out NE
+
+sil @useNE : $@convention(thin) (@guaranteed NE) -> ()
+
 sil [ossa] @Wrapper_init : $@convention(method) (@owned NE, @thin Wrapper.Type) -> @lifetime(copy 0) @owned Wrapper
 
 sil [ossa] @NCContainer_ne_read : $@yield_once @convention(method) (@guaranteed NCContainer) -> @lifetime(borrow 0) @yields @guaranteed NE
@@ -120,3 +144,40 @@ bb0(%0 : $UnsafePointer<Int64>):
   %19 = tuple ()
   return %19
 }
+
+// =============================================================================
+// @_addressable
+// =============================================================================
+
+sil [ossa] @addressableArg : $@convention(thin) (@in_guaranteed Holder) -> @lifetime(borrow address 0) @owned NE
+
+// CHECK-LABEL: sil hidden [ossa] @testAddressableArg : $@convention(thin) (@guaranteed Holder) -> () {
+// CHECK: bb0(%0 : @guaranteed $Holder):
+// CHECK: [[ALLOC:%.*]] = alloc_stack $Holder
+// CHECK: [[SB:%.*]] = store_borrow %0 to [[ALLOC]]
+// CHECK: [[APPLY:%.*]] = apply %{{.*}}([[SB]]) : $@convention(thin) (@in_guaranteed Holder) -> @lifetime(borrow address 0) @owned NE
+// CHECK: [[MD:%.*]] = mark_dependence [unresolved] [[APPLY]] on [[SB]]
+// CHECK: [[MV:%.*]] = move_value [var_decl] [[MD]]
+// CHECK: apply %{{.*}}([[MV]]) : $@convention(thin) (@guaranteed NE) -> ()
+// CHECK: destroy_value [[MV]]
+// CHECK: end_borrow [[SB]]
+// CHECK: dealloc_stack [[ALLOC]]
+// CHECK-LABEL: } // end sil function 'testAddressableArg'
+sil hidden [ossa] @testAddressableArg : $@convention(thin) (@guaranteed Holder) -> () {
+bb0(%0 : @guaranteed $Holder):
+  debug_value %0, let, name "arg", argno 1
+  %2 = alloc_stack $Holder
+  %3 = store_borrow %0 to %2
+  %4 = function_ref @addressableArg : $@convention(thin) (@in_guaranteed Holder) -> @lifetime(borrow address 0) @owned NE
+  %5 = apply %4(%3) : $@convention(thin) (@in_guaranteed Holder) -> @lifetime(borrow address 0) @owned NE
+  %6 = mark_dependence [unresolved] %5 on %3
+  end_borrow %3
+  %8 = move_value [var_decl] %6
+  debug_value %8, let, name "ne"
+  %useNE = function_ref @useNE : $@convention(thin) (@guaranteed NE) -> ()
+  %18 = apply %useNE(%8) : $@convention(thin) (@guaranteed NE) -> ()
+  destroy_value %8
+  dealloc_stack %2
+  %99 = tuple ()
+  return %99 : $()
+}