LifetimeDependenceScopeFixup: extend store_borrow allocations

Extend temporary allocations (sink dealloc_stacks) initialized by a store_borrow
across lifetime dependent uses.

Fixes rdar://143159873 ([nonescapable] extend rvalue lifetimes when they are the source of a dependency)

Author: atrick

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceScopeFixup.swift

@@ -106,6 +106,7 @@ let lifetimeDependenceScopeFixupPass = FunctionPass(
 
   let localReachabilityCache = LocalVariableReachabilityCache()
 
+  var mustFixStackNesting = false
   for instruction in function.instructions {
     guard let markDep = instruction as? MarkDependenceInstruction else {
       continue
@@ -122,6 +123,7 @@ let lifetimeDependenceScopeFixupPass = FunctionPass(
     guard scopeExtension.extendScopes(dependence: newLifetimeDep) else {
       continue
     }
+    mustFixStackNesting = mustFixStackNesting || scopeExtension.mustFixStackNesting
     let args = scopeExtension.findArgumentDependencies()
 
     // If the scope cannot be extended to the caller, this must be the outermost dependency level.
@@ -133,6 +135,9 @@ let lifetimeDependenceScopeFixupPass = FunctionPass(
     // Redirect the dependence base to the function arguments. This may create additional mark_dependence instructions.
     markDep.redirectFunctionReturn(to: args, context)
   }
+  if mustFixStackNesting {
+    context.fixStackNesting(in: function)
+  }
 }
 
 private extension Type {
@@ -285,6 +290,9 @@ private struct ScopeExtension {
   // Initialized after walking dependent uses. True if the scope can be extended into the caller.
   var dependsOnCaller: Bool?
 
+  // Does scope extension potentially invalidate stack nesting?
+  var mustFixStackNesting = false
+
   // Scopes listed in RPO over an upward walk. The outermost scope is first.
   var scopes = SingleInlineArray<ExtendableScope>()
 
@@ -383,6 +391,23 @@ private struct ExtendableScope {
     }
   }
 
+  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
+    }
+    return nil
+  }
+
   // Allow scope extension as long as `beginInst` is scoped instruction and does not define a variable scope.
   init?(_ scope: LifetimeDependence.Scope, beginInst: Instruction?) {
     self.scope = scope
@@ -742,9 +767,9 @@ extension ScopeExtension {
   // Extend the scopes that actually required extension.
   //
   // Consumes 'useRange'
-  private func extend(scopesToExtend: SingleInlineArray<ExtendableScope>,
-                      over useRange: inout InstructionRange,
-                      _ context: some MutatingContext) {
+  private mutating func extend(scopesToExtend: SingleInlineArray<ExtendableScope>,
+                               over useRange: inout InstructionRange,
+                               _ context: some MutatingContext) {
     var deadInsts = [Instruction]()
     for extScope in scopesToExtend {
       // Extend 'useRange' to to cover this scope's end instructions. 'useRange' cannot be extended until the
@@ -772,6 +797,9 @@ extension ScopeExtension {
 
     // Delete original end instructions.
     for deadInst in deadInsts {
+      if deadInst is DeallocStackInst {
+        mustFixStackNesting = true
+      }
       context.erase(instruction: deadInst)
     }
   }
@@ -790,8 +818,8 @@ extension ExtendableScope {
         // A yield is already considered nested within the coroutine.
         break
       case let .store(initializingStore, _):
-        if let sb = initializingStore as? StoreBorrowInst {
-          return canExtend(storeBorrow: sb, over: &range)
+        if initializingStore is StoreBorrowInst {
+          return true
         }
       }
       return true
@@ -823,28 +851,25 @@ extension ExtendableScope {
     return true
   }
 
-  /// A store borrow is considered to be nested within the scope of its stored values. It is, however, also
-  /// restricted to the range of its allocation.
-  ///
-  /// TODO: consider rewriting the dealloc_stack instructions if we ever find that SILGen emits them sooner that
-  /// we need for lifetime dependencies.
-  func canExtend(storeBorrow: StoreBorrowInst, over range: inout InstructionRange) -> Bool {
-    // store_borrow can be extended if all deallocations occur after the use range.
-    return storeBorrow.allocStack.deallocations.allSatisfy({ !range.contains($0) })
-  }
-
   /// Extend this scope over the 'range' boundary. Return the old scope ending instructions to be deleted.
   func extend(over range: inout InstructionRange, _ context: some MutatingContext) -> [Instruction] {
     // Collect the original end instructions and extend the range to to cover them. The resulting access scope
     // must cover the original scope because it may protect other memory operations.
-    let endsToErase = self.endInstructions
-    // Track range ending instructions that have not been reused for later deletion.
+    let originalScopeEnds = [Instruction](self.endInstructions)
+    // Track scope-ending instructions that have not yet been reused as range-ending instructions.
     var unreusedEnds = InstructionSet(context)
-    for end in endsToErase {
+    for end in originalScopeEnds {
       assert(range.inclusiveRangeContains(end))
       unreusedEnds.insert(end)
     }
     defer { unreusedEnds.deinitialize() }
+
+    // Never reuse dealloc_stack to avoid running data flow.
+    var endsToErase = [Instruction]()
+    if let deallocs = self.deallocs {
+      endsToErase.append(contentsOf: deallocs.map { $0 })
+    }
+
     for end in range.ends {
       let location = end.location.asAutoGenerated
       switch end {
@@ -857,45 +882,49 @@ extension ExtendableScope {
         // function argument.
         let builder = Builder(before: end, location: location, context)
         // Insert newEnd so that this scope will be nested in any outer scopes.
-        range.insert(createEndInstruction(builder, context))
+        range.insert(contentsOf: createEndInstructions(builder, context))
         continue
       default:
         break
       }
+      // If this range ending instruction was also scope-ending, then mark it as reused by removing it from the set.
       if unreusedEnds.contains(end) {
         unreusedEnds.erase(end)
         assert(!unreusedEnds.contains(end))
         continue
       }
       Builder.insert(after: end, location: location, context) {
-        range.insert(createEndInstruction($0, context))
+        range.insert(contentsOf: createEndInstructions($0, context))
       }
     }
     for exitInst in range.exits {
       let location = exitInst.location.asAutoGenerated
       let builder = Builder(before: exitInst, location: location, context)
-      range.insert(createEndInstruction(builder, context))
+      range.insert(contentsOf: createEndInstructions(builder, context))
     }
-    return endsToErase.filter { unreusedEnds.contains($0) }
+    endsToErase.append(contentsOf: originalScopeEnds.filter { unreusedEnds.contains($0) })
+    return endsToErase
   }
 
   /// Create a scope-ending instruction at 'builder's insertion point.
-  func createEndInstruction(_ builder: Builder, _ context: some Context) -> Instruction {
+  func createEndInstructions(_ builder: Builder, _ context: some Context) -> SingleInlineArray<Instruction> {
     switch self.scope {
     case let .access(beginAccess):
-      return builder.createEndAccess(beginAccess: beginAccess)
+      return SingleInlineArray(element: builder.createEndAccess(beginAccess: beginAccess))
     case let .borrowed(beginBorrow):
-      return builder.createEndBorrow(of: beginBorrow.value)
+      return SingleInlineArray(element: builder.createEndBorrow(of: beginBorrow.value))
     case let .yield(yieldedValue):
       let beginApply = yieldedValue.definingInstruction as! BeginApplyInst
       // createEnd() returns non-nil because beginApply.endReaches() was checked by canExtend()
-      return beginApply.createEnd(builder, context)!
+      return SingleInlineArray(element: beginApply.createEnd(builder, context)!)
     case let .initialized(initializer):
       switch initializer {
       case let .store(initializingStore: store, initialAddress: _):
         if let sb = store as? StoreBorrowInst {
-          // FIXME: we may need to rewrite the dealloc_stack.
-          return builder.createEndBorrow(of: sb)
+          var endInsts = SingleInlineArray<Instruction>()
+          endInsts.append(builder.createEndBorrow(of: sb))
+          endInsts.append(builder.createDeallocStack(sb.allocStack))
+          return endInsts
         }
         break
       case .argument, .yield:

test/SILOptimizer/lifetime_dependence/lifetime_dependence_scope_fixup.swift

@@ -327,3 +327,10 @@ func testWrite(_ w: inout ArrayWrapper) {
     span[i] = 0
   }
 }
+
+// Test store_borrow extension which is required when the addressable UTF8View is extended over the Span's uses.
+@available(SwiftStdlib 6.2, *)
+func testSpanOfBorrowByAddress(_ i: Int) -> Int {
+    let span = String(i).utf8.span
+    return span.count
+}