Merge pull request #84091 from atrick/lifedep-tempalloc

LifetimeDependenceDiagnostics: check for on-stack trivial copies

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
+    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.scopeEndingOperands.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.deallocations
   }
 
-  // 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

SwiftCompilerSources/Sources/Optimizer/Utilities/LifetimeDependenceUtils.swift

@@ -381,7 +381,29 @@ extension LifetimeDependence.Scope {
   }
 }
 
+// Scope helpers.
 extension LifetimeDependence.Scope {
+  // 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 {
@@ -456,7 +478,7 @@ extension LifetimeDependence.Scope {
       }
       let address = initializer.initialAddress
       if address.type.objectType.isTrivial(in: address.parentFunction) {
-        return nil
+        return LifetimeDependence.Scope.computeAddressableRange(initializer: initializer, context)
       }
       return LifetimeDependence.Scope.computeInitializedRange(initializer: initializer, context)
     case let .unknown(value):
@@ -503,6 +525,46 @@ extension LifetimeDependence.Scope {
     }
     return range
   }
+
+  // For trivial values, compute the range in which the value may have its address taken.
+  // Returns nil unless the address has both an addressable parameter use and a dealloc_stack use.
+  //
+  // This extended range is convervative. In theory, it can lead to a "false positive" diagnostic if this scope was
+  // computed for a apply site that takes this address as *non-addressable*, as opposed to the apply site discovered
+  // here that takes the alloc_stack as an *addressable* argument. This won't normally happen because addressability
+  // depends on the value's type (via @_addressableForDepenencies), so all other lifetime-dependent applies should be
+  // addressable. Furthermore, SILGen only creates temporary stack locations for addressable arguments for a single
+  // argument.
+  private static func computeAddressableRange(initializer: AccessBase.Initializer, _ context: Context)
+    -> InstructionRange? {
+    switch initializer {
+    case .argument, .yield:
+      return nil
+    case let .store(initializingStore, initialAddr):
+      guard initialAddr is AllocStackInst else {
+        return nil
+      }
+      var isAddressable = false
+      var deallocInsts = SingleInlineArray<DeallocStackInst>()
+      for use in initialAddr.uses {
+        let inst = use.instruction
+        switch inst {
+        case let apply as ApplySite:
+          isAddressable = isAddressable || apply.isAddressable(operand: use)
+        case let dealloc as DeallocStackInst:
+          deallocInsts.append(dealloc)
+        default:
+          break
+        }
+      }
+      guard isAddressable, !deallocInsts.isEmpty else {
+        return nil
+      }
+      var range = InstructionRange(begin: initializingStore, context)
+      range.insert(contentsOf: deallocInsts)
+      return range
+    }
+  }
 }
 
 // =============================================================================

test/SILOptimizer/lifetime_dependence/scope_fixup.sil

@@ -87,6 +87,11 @@ public struct AThing : ~Copyable {
   init(somethings: UnsafeMutableRawBufferPointer)
 }
 
+@_addressableForDependencies
+public struct InlineInt {
+  var i: Int
+}
+
 sil @getContainer : $@convention(thin) (@inout AThing) -> @lifetime(borrow address_for_deps 0) @out Container
 sil @getMutableView : $@convention(thin) (@in_guaranteed Container) -> @lifetime(copy 0) @out Optional<MutableView>
 sil @modAThing : $@convention(thin) (@inout AThing) -> ()
@@ -138,6 +143,10 @@ sil @getAddressableObj : $@convention(method) (@guaranteed Holder) -> @owned Add
 
 sil @getAddressableObj_NE : $@convention(thin) (@in_guaranteed AddressableForDepsObj) -> @owned NE
 
+sil @globalAddress : $@convention(thin) () -> Builtin.RawPointer
+sil @addressInt : $@convention(thin) (@in_guaranteed InlineInt) -> @lifetime(borrow address_for_deps 0) @owned Span<Int>
+sil @useSpan : $@convention(thin) (@guaranteed Span<Int>) -> ()
+
 // NCContainer.wrapper._read:
 //   var wrapper: Wrapper {
 //     _read {
@@ -914,3 +923,50 @@ bb2:
   dealloc_box [dead_end] %1
   unreachable
 }
+
+// Test trivial alloc_stack extension when used by an addressable dependency.
+//
+// CHECK-LABEL: sil hidden [noinline] [ossa] @testTempTrivialStackCopy : $@convention(thin) () -> () {
+// CHECK:   [[STK:%[0-9]+]] = alloc_stack $InlineInt
+// CHECK:   [[MD:%[0-9]+]] = mark_dependence [unresolved] %{{.*}} on [[STK]]
+// CHECK:   [[VAR:%[0-9]+]] = move_value [var_decl] [[MD]]
+// CHECK: bb3:                                              // Preds: bb2 bb1
+// CHECK:   [[BORROW:%[0-9]+]] = begin_borrow [[VAR]]
+// CHECK:   apply %{{.*}}([[BORROW]]) : $@convention(thin) (@guaranteed Span<Int>) -> ()
+// CHECK:   destroy_value [[VAR]]
+// CHECK:   dealloc_stack [[STK]]
+// CHECK-LABEL: } // end sil function 'testTempTrivialStackCopy'
+sil hidden [noinline] [ossa] @testTempTrivialStackCopy : $@convention(thin) () -> () {
+bb0:
+  %0 = function_ref @globalAddress : $@convention(thin) () -> Builtin.RawPointer
+  %1 = apply %0() : $@convention(thin) () -> Builtin.RawPointer
+  %2 = pointer_to_address %1 to [strict] $*InlineInt
+  %3 = begin_access [read] [dynamic] %2
+  %4 = load [trivial] %3
+  end_access %3
+  %6 = alloc_stack $InlineInt
+  store %4 to [trivial] %6
+
+  %8 = function_ref @addressInt : $@convention(thin) (@in_guaranteed InlineInt) -> @lifetime(borrow address_for_deps 0) @owned Span<Int>
+  %9 = apply %8(%6) : $@convention(thin) (@in_guaranteed InlineInt) -> @lifetime(borrow address_for_deps 0) @owned Span<Int>
+  %10 = mark_dependence [unresolved] %9 on %6
+  %11 = move_value [var_decl] %10
+  cond_br undef, bb1, bb2
+
+bb1:
+  dealloc_stack %6
+  br bb3
+bb2:
+  dealloc_stack %6
+  br bb3
+
+bb3:
+  %13 = begin_borrow %11
+
+  %14 = function_ref @useSpan : $@convention(thin) (@guaranteed Span<Int>) -> ()
+  %15 = apply %14(%13) : $@convention(thin) (@guaranteed Span<Int>) -> ()
+  end_borrow %13
+  destroy_value %11
+  %18 = tuple ()
+  return %18
+}