LifetimeDependenceScopeFixup: extend temporary stack allocations

When the source of a lifetime dependency is a stack-allocated address, extend
the stack allocation to cover all dependent uses.

This avoids miscompilations for "addressable" dependencies which arise in code
built with -enable-experimental-feature AddressableTypes or
AddressableParameters. It is always an error for SILGen to emit the alloc_stack
in such cases. Nonetheless, we want to handle these unexpected cases gracefully
in SIL as a diagnostic error rather than allowing a miscompile.

Fixes rdar://159680262 ([nonescapable] diagnose dependence on a
temporary copy of a global array)

(cherry picked from commit 7dfd105)

Author: atrick

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceScopeFixup.swift

@@ -365,59 +365,70 @@ extension ScopeExtension {
 private struct ExtendableScope {
   enum Introducer {
     case scoped(ScopedInstruction)
+    case stack(Instruction)
     case owned(Value)
   }
 
+  // scope.allocStackInstruction is always valid for Introducer.allocStack and is valid for Introducer.scoped when
+  // ScopedInstruction is a store_borrow.
   let scope: LifetimeDependence.Scope
   let introducer: Introducer
 
   var firstInstruction: Instruction {
     switch introducer {
     case let .scoped(scopedInst):
       return scopedInst.instruction
+    case let .stack(initializingStore):
+      return initializingStore
     case let .owned(value):
       if let definingInst = value.definingInstructionOrTerminator {
         return definingInst
       }
       return value.parentBlock.instructions.first!
     }
   }
+
   var endInstructions: LazyMapSequence<LazyFilterSequence<UseList>, Instruction> {
     switch introducer {
     case let .scoped(scopedInst):
       return scopedInst.endOperands.users
+    case .stack:
+      // For alloc_stack without a store-borrow scope, include the deallocs in its scope to ensure that we never shorten
+      // the original allocation. It's possible that some other use depends on the address.
+      //
+      // Same as 'AllocStackInst.deallocations' but as an Instruction list...
+      return scope.allocStackInstruction!.uses.lazy.filter
+      { $0.instruction is DeallocStackInst }.lazy.map { $0.instruction }
+
     case let .owned(value):
       return value.uses.endingLifetime.users
     }
   }
 
   var deallocs: LazyMapSequence<LazyFilterSequence<UseList>, DeallocStackInst>? {
-    switch self.scope {
-    case let .initialized(initializer):
-      switch initializer {
-      case let .store(initializingStore: store, initialAddress: _):
-        if let sb = store as? StoreBorrowInst {
-          return sb.allocStack.uses.filterUsers(ofType: DeallocStackInst.self)
-        }
-      default:
-        break
-      }
-    default:
-      break
+    guard let allocStack = scope.allocStackInstruction else {
+      return nil
     }
-    return nil
+    return allocStack.uses.users(ofType: DeallocStackInst.self)
   }
 
-  // Allow scope extension as long as `beginInst` is scoped instruction and does not define a variable scope.
+  // Allow scope extension as long as `beginInst` does not define a variable scope and is either a scoped instruction or
+  // a store to a singly-initialized temporary.
   init?(_ scope: LifetimeDependence.Scope, beginInst: Instruction?) {
     self.scope = scope
     guard let beginInst = beginInst, VariableScopeInstruction(beginInst) == nil else {
       return nil
     }
-    guard let scopedInst = beginInst as? ScopedInstruction else {
-      return nil
+    // Check for "scoped" store_borrow extension before checking allocStackInstruction.
+    if let scopedInst = beginInst as? ScopedInstruction {
+      self.introducer = .scoped(scopedInst)
+      return
     }
-    self.introducer = .scoped(scopedInst)
+    if scope.allocStackInstruction != nil {
+      self.introducer = .stack(beginInst)
+      return
+    }
+    return nil
   }
 
   // Allow extension of owned temporaries that
@@ -484,11 +495,14 @@ extension ScopeExtension {
       switch initializer {
       case let .store(initializingStore: store, initialAddress: _):
         if let sb = store as? StoreBorrowInst {
-          // Follow the source for nested scopes.
+          // Follow the stored value since the owner of the borrowed value needs to cover this allocation.
           gatherExtensions(valueOrAddress: sb.source)
-          scopes.push(ExtendableScope(scope, beginInst: sb)!)
+        }
+        if scope.allocStackInstruction != nil {
+          scopes.push(ExtendableScope(scope, beginInst: store)!)
           return
         }
+        break
       case .argument, .yield:
         // TODO: extend indirectly yielded scopes.
         break
@@ -816,13 +830,10 @@ extension ExtendableScope {
       switch initializer {
       case .argument, .yield:
         // A yield is already considered nested within the coroutine.
-        break
-      case let .store(initializingStore, _):
-        if initializingStore is StoreBorrowInst {
-          return true
-        }
+        return true
+      case .store:
+        return self.scope.allocStackInstruction != nil
       }
-      return true
     default:
       // non-yield scopes can always be ended at any point.
       return true
@@ -868,6 +879,9 @@ extension ExtendableScope {
     var endsToErase = [Instruction]()
     if let deallocs = self.deallocs {
       endsToErase.append(contentsOf: deallocs.map { $0 })
+      for dealloc in deallocs {
+        unreusedEnds.erase(dealloc)
+      }
     }
 
     for end in range.ends {
@@ -920,10 +934,12 @@ extension ExtendableScope {
     case let .initialized(initializer):
       switch initializer {
       case let .store(initializingStore: store, initialAddress: _):
+        var endInsts = SingleInlineArray<Instruction>()
         if let sb = store as? StoreBorrowInst {
-          var endInsts = SingleInlineArray<Instruction>()
           endInsts.append(builder.createEndBorrow(of: sb))
-          endInsts.append(builder.createDeallocStack(sb.allocStack))
+        }
+        if let allocStack = self.scope.allocStackInstruction {
+          endInsts.append(builder.createDeallocStack(allocStack))
           return endInsts
         }
         break
@@ -932,14 +948,14 @@ extension ExtendableScope {
         break
       }
     case let .owned(value):
-      return builder.createDestroyValue(operand: value)
+      return SingleInlineArray(element: builder.createDestroyValue(operand: value))
     case let .local(varInst):
       switch varInst {
       case let .beginBorrow(beginBorrow):
         // FIXME: we may need to rewrite the dealloc_stack.
-        return builder.createEndBorrow(of: beginBorrow)
+        return SingleInlineArray(element: builder.createEndBorrow(of: beginBorrow))
       case let .moveValue(moveValue):
-        return builder.createDestroyValue(operand: moveValue)
+        return SingleInlineArray(element: builder.createDestroyValue(operand: moveValue))
       }
     default:
       break

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -444,6 +444,10 @@ struct FunctionPassContext : MutatingContext {
       notifyInstructionsChanged()
     }
   }
+
+  func fixStackNesting(in function: Function) {
+    _bridged.fixStackNesting(function.bridged)
+  }
 }
 
 struct SimplifyContext : MutatingContext {

SwiftCompilerSources/Sources/Optimizer/Utilities/LifetimeDependenceUtils.swift

@@ -381,8 +381,30 @@ extension LifetimeDependence.Scope {
   }
 }
 
+// Scope helpers.
 extension LifetimeDependence.Scope {
-  /// Ignore "irrelevent" borrow scopes: load_borrow or begin_borrow without [var_decl]
+  // If the LifetimeDependenceScope is .initialized, then return the alloc_stack.
+  var allocStackInstruction: AllocStackInst? {
+    switch self {
+    case let .initialized(initializer):
+      switch initializer {
+      case let .store(initializingStore: store, initialAddress):
+        if let allocStackInst = initialAddress as? AllocStackInst {
+          return allocStackInst
+        }
+        if let sb = store as? StoreBorrowInst {
+          return sb.allocStack
+        }
+        return nil
+      default:
+        return nil
+      }
+      default:
+        return nil
+    }
+  }
+
+  /// Ignore "irrelevant" borrow scopes: load_borrow or begin_borrow without [var_decl]
   func ignoreBorrowScope(_ context: some Context) -> LifetimeDependence.Scope? {
     guard case let .borrowed(beginBorrowVal) = self else {
       return nil

test/SILOptimizer/lifetime_dependence/scope_fixup.sil

@@ -732,14 +732,16 @@ bb0:
 // Test that initializing LifetimeDependence.Scope from an unidentified enclosing access scope does not
 // unwrap nil.
 //
+// FIXME: the dependence should be on [[AS]] not [[SE]. This is fixed on the main branch.
+//
 // CHECK-LABEL: sil hidden [ossa] @testUnidentified : $@convention(method) (@owned B) -> () {
 // CHECK: [[AS:%.*]] = alloc_stack $B
 // CHECK: [[MU:%.*]] = mark_unresolved_non_copyable_value [consumable_and_assignable] [[AS]]
 // CHECK: [[DD:%.*]] = drop_deinit [[MU]]
 // CHECK: [[SE:%.*]] = struct_element_addr [[DD]], #B.ne
 // CHECK: [[LB:%.*]] = load_borrow
 // CHECK: apply %{{.*}} : $@convention(thin) (@guaranteed NE) -> UnsafePointer<A>
-// CHECK: mark_dependence [unresolved] %{{.*}} on [[AS]]
+// CHECK: mark_dependence [unresolved] %{{.*}} on [[SE]]
 // CHECK-LABEL: } // end sil function 'testUnidentified'
 sil hidden [ossa] @testUnidentified : $@convention(method) (@owned B) -> () {
 bb0(%0 : @owned $B):

test/SILOptimizer/lifetime_dependence/verify_diagnostics.swift

@@ -326,8 +326,16 @@ var values: InlineArray<_, Int> = [0, 1, 2]
 
 // rdar://159680262 ([nonescapable] diagnose dependence on a temporary copy of a global array)
 @available(Span 0.1, *)
-func test() -> Int {
+func readGlobalSpan() -> Int {
+  let span = values.span
+  return span[3]
+}
+
+// TODO: rdar://151320168 ([nonescapable] support immortal span over global array)
+@available(Span 0.1, *)
+@_lifetime(immortal)
+func returnGlobalSpan() -> Span<Int> {
   let span = values.span // expected-error{{lifetime-dependent variable 'span' escapes its scope}}
   // expected-note@-1{{it depends on this scoped access to variable 'values'}}
-  return span[3] // expected-note{{this use of the lifetime-dependent value is out of scope}}
+  return span // expected-note{{this use causes the lifetime-dependent value to escape}}
 }