Merge pull request #83728 from atrick/rvalue-extension

LifetimeDependenceScopeFixup: extend store_borrow allocations

Author: atrick

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/AllocBoxToStack.swift

@@ -449,7 +449,7 @@ private func hoistMarkUnresolvedInsts(stackAddress: Value,
     builder = Builder(atBeginOf: stackAddress.parentBlock, context)
   }
   let mu = builder.createMarkUnresolvedNonCopyableValue(value: stackAddress, checkKind: checkKind,  isStrict: false)
-  stackAddress.uses.ignore(user: mu).ignoreDebugUses.ignoreUsers(ofType: DeallocStackInst.self)
+  stackAddress.uses.ignore(user: mu).ignoreDebugUses.ignoreUses(ofType: DeallocStackInst.self)
     .replaceAll(with: mu, context)
 }
 

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/DestroyHoisting.swift

@@ -83,7 +83,7 @@ let destroyHoisting = FunctionPass(name: "destroy-hoisting") {
 private func optimize(value: Value, _ context: FunctionPassContext) {
   guard value.ownership == .owned,
         // Avoid all the analysis effort if there are no destroys to hoist.
-        !value.uses.filterUsers(ofType: DestroyValueInst.self).isEmpty
+        !value.uses.filterUses(ofType: DestroyValueInst.self).isEmpty
   else {
     return
   }

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,27 +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
-    var unusedEnds = InstructionSet(context)
-    for end in endsToErase {
+    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 originalScopeEnds {
       assert(range.inclusiveRangeContains(end))
-      unusedEnds.insert(end)
+      unreusedEnds.insert(end)
     }
-    defer { unusedEnds.deinitialize() }
+    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 {
@@ -856,60 +882,64 @@ 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 unusedEnds.contains(end) {
-        unusedEnds.erase(end)
-        assert(!unusedEnds.contains(end))
+      // 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 { unusedEnds.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:
         // TODO: extend indirectly yielded scopes.
         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/FunctionPasses/TempRValueElimination.swift

@@ -245,7 +245,7 @@ private extension AllocStackInst {
     var liferange = InstructionRange(begin: self, context)
     defer { liferange.deinitialize() }
 
-    liferange.insert(contentsOf: uses.ignoreUsers(ofType: DeallocStackInst.self).lazy.map { $0.instruction })
+    liferange.insert(contentsOf: uses.ignoreUses(ofType: DeallocStackInst.self).lazy.map { $0.instruction })
 
     for use in uses {
       switch use.instruction {

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyAllocStack.swift

@@ -206,7 +206,7 @@ private extension AllocStackInst {
           iea.replace(with: newAlloc, context)
         }
       case let oea as OpenExistentialAddrInst:
-        assert(oea.uses.ignoreUsers(ofType: DestroyAddrInst.self).isEmpty)
+        assert(oea.uses.ignoreUses(ofType: DestroyAddrInst.self).isEmpty)
         oea.replace(with: newAlloc, context)
       case let cab as CheckedCastAddrBranchInst:
         let builder = Builder(before: cab, context)
@@ -246,7 +246,7 @@ private extension AllocStackInst {
            is DebugValueInst:
         break
       case let oea as OpenExistentialAddrInst:
-        if !oea.uses.ignoreUsers(ofType: DestroyAddrInst.self).isEmpty {
+        if !oea.uses.ignoreUses(ofType: DestroyAddrInst.self).isEmpty {
           return nil
         }
       case let iea as InitExistentialAddrInst:

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyBeginAndLoadBorrow.swift

@@ -29,7 +29,7 @@ extension BeginBorrowInst : OnoneSimplifiable, SILCombineSimplifiable {
 
 extension LoadBorrowInst : Simplifiable, SILCombineSimplifiable {
   func simplify(_ context: SimplifyContext) {
-    if uses.ignoreDebugUses.ignoreUsers(ofType: EndBorrowInst.self).isEmpty {
+    if uses.ignoreDebugUses.ignoreUses(ofType: EndBorrowInst.self).isEmpty {
       context.erase(instructionIncludingAllUsers: self)
       return
     }
@@ -52,7 +52,7 @@ extension LoadBorrowInst : Simplifiable, SILCombineSimplifiable {
 
   private func tryCombineWithCopy(_ context: SimplifyContext) {
     let forwardedValue = lookThroughSingleForwardingUses()
-    guard let singleUser = forwardedValue.uses.ignoreUsers(ofType: EndBorrowInst.self).singleUse?.instruction,
+    guard let singleUser = forwardedValue.uses.ignoreUses(ofType: EndBorrowInst.self).singleUse?.instruction,
           let copy = singleUser as? CopyValueInst,
           copy.parentBlock == self.parentBlock else {
       return
@@ -81,13 +81,13 @@ private func tryReplaceBorrowWithOwnedOperand(beginBorrow: BeginBorrowInst, _ co
 private func removeBorrowOfThinFunction(beginBorrow: BeginBorrowInst, _ context: SimplifyContext) {
   guard let thin2thickFn = beginBorrow.borrowedValue as? ThinToThickFunctionInst,
         // For simplicity don't go into the trouble of removing reborrow phi arguments.
-        beginBorrow.uses.filterUsers(ofType: BranchInst.self).isEmpty else
+        beginBorrow.uses.filterUses(ofType: BranchInst.self).isEmpty else
   {
     return
   }
   // `thin_to_thick_function` has "none" ownership and is compatible with guaranteed values.
   // Therefore the `begin_borrow` is not needed.
-  beginBorrow.uses.ignoreUsers(ofType: EndBorrowInst.self).replaceAll(with: thin2thickFn, context)
+  beginBorrow.uses.ignoreUses(ofType: EndBorrowInst.self).replaceAll(with: thin2thickFn, context)
   context.erase(instructionIncludingAllUsers: beginBorrow)
 }
 
@@ -110,7 +110,7 @@ private func tryReplaceCopy(
   withCopiedOperandOf beginBorrow: BeginBorrowInst,
   _ context: SimplifyContext
 ) -> Bool {
-  guard let singleUser = forwardedValue.uses.ignoreUsers(ofType: EndBorrowInst.self).singleUse?.instruction,
+  guard let singleUser = forwardedValue.uses.ignoreUses(ofType: EndBorrowInst.self).singleUse?.instruction,
         let copy = singleUser as? CopyValueInst,
         copy.parentBlock == beginBorrow.parentBlock else {
     return false
@@ -153,7 +153,7 @@ private extension Value {
   /// ```
   /// Returns self if this value has no uses which are ForwardingInstructions.
   func lookThroughSingleForwardingUses() -> Value {
-    if let singleUse = uses.ignoreUsers(ofType: EndBorrowInst.self).singleUse,
+    if let singleUse = uses.ignoreUses(ofType: EndBorrowInst.self).singleUse,
        let fwdInst = singleUse.instruction as? (SingleValueInstruction & ForwardingInstruction),
        fwdInst.canConvertToOwned,
        fwdInst.isSingleForwardedOperand(singleUse),
@@ -165,7 +165,7 @@ private extension Value {
   }
 
   var allUsesCanBeConvertedToOwned: Bool {
-    let relevantUses = uses.ignoreUsers(ofType: EndBorrowInst.self)
+    let relevantUses = uses.ignoreUses(ofType: EndBorrowInst.self)
     return relevantUses.allSatisfy { $0.canAccept(ownership: .owned) }
   }
 
@@ -175,7 +175,7 @@ private extension Value {
   }
 
   func replaceAllDestroys(with replacement: Value, _ context: SimplifyContext) {
-    uses.filterUsers(ofType: DestroyValueInst.self).replaceAll(with: replacement, context)
+    uses.filterUses(ofType: DestroyValueInst.self).replaceAll(with: replacement, context)
   }
 }
 
@@ -187,7 +187,7 @@ private extension Instruction {
       // The value and instruction are in the same block. All uses are dominated by both.
       return true
     }
-    let destroys = value.uses.filterUsers(ofType: DestroyValueInst.self)
+    let destroys = value.uses.filterUses(ofType: DestroyValueInst.self)
     return destroys.allSatisfy({ $0.instruction.parentBlock == parentBlock})
   }
 }

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -400,7 +400,7 @@ private extension Value {
       //   var p = Point(x: 10, y: 20)
       //   let o = UnsafePointer(&p)
       // Therefore ignore the `end_access` use of a `begin_access`.
-      let relevantUses = singleUseValue.uses.ignoreDebugUses.ignoreUsers(ofType: EndAccessInst.self)
+      let relevantUses = singleUseValue.uses.ignoreDebugUses.ignoreUses(ofType: EndAccessInst.self)
 
       guard let use = relevantUses.singleUse else {
         return nil

SwiftCompilerSources/Sources/Optimizer/ModulePasses/StackProtection.swift

@@ -365,7 +365,7 @@ private struct StackProtectionOptimization {
     let builder = Builder(after: beginAccess, location: beginAccess.location.asAutoGenerated, context)
     let temporary = builder.createAllocStack(beginAccess.type)
 
-    beginAccess.uses.ignoreUsers(ofType: EndAccessInst.self).replaceAll(with: temporary, context)
+    beginAccess.uses.ignoreUses(ofType: EndAccessInst.self).replaceAll(with: temporary, context)
 
     for endAccess in beginAccess.endInstructions {
       let endBuilder = Builder(before: endAccess, location: endAccess.location.asAutoGenerated, context)

SwiftCompilerSources/Sources/Optimizer/Utilities/OwnershipLiveness.swift

@@ -717,7 +717,7 @@ extension InteriorUseWalker: AddressUseVisitor {
       if handleInner(borrowed: sb) == .abortWalk {
         return .abortWalk
       }
-      return sb.uses.filterUsers(ofType: EndBorrowInst.self).walk {
+      return sb.uses.filterUses(ofType: EndBorrowInst.self).walk {
         useVisitor($0)
       }
     case let load as LoadBorrowInst:

SwiftCompilerSources/Sources/SIL/Operand.swift

@@ -164,11 +164,15 @@ extension Sequence where Element == Operand {
     self.lazy.filter { !($0.instruction is DebugValueInst) }
   }
 
-  public func filterUsers<I: Instruction>(ofType: I.Type) -> LazyFilterSequence<Self> {
+  public func filterUses<I: Instruction>(ofType: I.Type) -> LazyFilterSequence<Self> {
     self.lazy.filter { $0.instruction is I }
   }
 
-  public func ignoreUsers<I: Instruction>(ofType: I.Type) -> LazyFilterSequence<Self> {
+  public func filterUsers<I: Instruction>(ofType: I.Type) -> LazyMapSequence<LazyFilterSequence<Self>, I> {
+    self.lazy.filter { $0.instruction is I }.lazy.map { $0.instruction as! I }
+  }
+
+  public func ignoreUses<I: Instruction>(ofType: I.Type) -> LazyFilterSequence<Self> {
     self.lazy.filter { !($0.instruction is I) }
   }
 
@@ -177,11 +181,11 @@ extension Sequence where Element == Operand {
   }
 
   public func getSingleUser<I: Instruction>(ofType: I.Type) -> I? {
-    filterUsers(ofType: I.self).singleUse?.instruction as? I
+    filterUses(ofType: I.self).singleUse?.instruction as? I
   }
 
   public func getSingleUser<I: Instruction>(notOfType: I.Type) -> Instruction? {
-    ignoreUsers(ofType: I.self).singleUse?.instruction
+    ignoreUses(ofType: I.self).singleUse?.instruction
   }
 
   public var endingLifetime: LazyFilterSequence<Self> {