Merge pull request #85076 from atrick/lifedep-reassign

Lifetimes: handle MutableSpan reassignment and 'inout' arguments

Author: atrick

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceDiagnostics.swift

@@ -160,7 +160,7 @@ private struct DiagnoseDependence {
   func checkInScope(operand: Operand) -> WalkResult {
     if let range, !range.inclusiveRangeContains(operand.instruction) {
       log("  out-of-range error: \(operand.instruction)")
-      reportError(operand: operand, diagID: .lifetime_outside_scope_use)
+      reportError(escapingValue: operand.value, user: operand.instruction, diagID: .lifetime_outside_scope_use)
       return .abortWalk
     }
     log("  contains: \(operand.instruction)")
@@ -169,7 +169,12 @@ private struct DiagnoseDependence {
 
   func reportEscaping(operand: Operand) {
     log("  escaping error: \(operand.instruction)")
-    reportError(operand: operand, diagID: .lifetime_outside_scope_escape)
+    reportError(escapingValue: operand.value, user: operand.instruction, diagID: .lifetime_outside_scope_escape)
+  }
+
+  func reportEscaping(value: Value, user: Instruction) {
+    log("  escaping error: \(value) at \(user)")
+    reportError(escapingValue: value, user: user, diagID: .lifetime_outside_scope_escape)
   }
 
   func reportUnknown(operand: Operand) {
@@ -184,7 +189,8 @@ private struct DiagnoseDependence {
       if inoutArg == sourceArg {
         return .continueWalk
       }
-      if function.argumentConventions.getDependence(target: inoutArg.index, source: sourceArg.index) != nil {
+      if function.argumentConventions.parameterDependence(targetArgumentIndex: inoutArg.index,
+                                                          sourceArgumentIndex: sourceArg.index) != nil {
         // The inout result depends on a lifetime that is inherited or borrowed in the caller.
         log("  has dependent inout argument: \(inoutArg)")
         return .continueWalk
@@ -257,15 +263,15 @@ private struct DiagnoseDependence {
     return .abortWalk
   }
 
-  func reportError(operand: Operand, diagID: DiagID) {
+  func reportError(escapingValue: Value, user: Instruction, diagID: DiagID) {
     // If the dependent value is Escapable, then mark_dependence resolution fails, but this is not a diagnostic error.
     if dependence.dependentValue.isEscapable {
       return
     }
     onError()
 
     // Identify the escaping variable.
-    let escapingVar = LifetimeVariable(usedBy: operand, context)
+    let escapingVar = LifetimeVariable(definedBy: escapingValue, user: user, context)
     if let varDecl = escapingVar.varDecl {
       // Use the variable location, not the access location.
       // Variable names like $return_value and $implicit_value don't have source locations.
@@ -287,7 +293,7 @@ private struct DiagnoseDependence {
     diagnoseImplicitFunction()
     reportScope()
     // Identify the use point.
-    if let userSourceLoc = operand.instruction.location.sourceLoc {
+    if let userSourceLoc = user.location.sourceLoc {
       diagnose(userSourceLoc, diagID)
     }
   }
@@ -358,13 +364,12 @@ private struct LifetimeVariable {
     return varDecl?.userFacingName
   }
 
-  init(usedBy operand: Operand, _ context: some Context) {
-    self = .init(dependent: operand.value, context)
+  init(definedBy value: Value, user: Instruction, _ context: some Context) {
+    self = .init(dependent: value, context)
     // variable names like $return_value and $implicit_value don't have source locations.
-    // For @out arguments, the operand's location is the best answer.
+    // For @out arguments, the user's location is the best answer.
     // Otherwise, fall back to the function's location.
-    self.sourceLoc = self.sourceLoc ?? operand.instruction.location.sourceLoc
-      ?? operand.instruction.parentFunction.location.sourceLoc
+    self.sourceLoc = self.sourceLoc ?? user.location.sourceLoc ?? user.parentFunction.location.sourceLoc
   }
 
   init(definedBy value: Value, _ context: some Context) {
@@ -552,9 +557,9 @@ extension DiagnoseDependenceWalker : LifetimeDependenceDefUseWalker {
     return .abortWalk
   }
 
-  mutating func inoutDependence(argument: FunctionArgument, on operand: Operand) -> WalkResult {
+  mutating func inoutDependence(argument: FunctionArgument, functionExit: Instruction) -> WalkResult {
     if diagnostics.checkInoutResult(argument: argument) == .abortWalk {
-      diagnostics.reportEscaping(operand: operand)
+      diagnostics.reportEscaping(value: argument, user: functionExit)
       return .abortWalk
     }
     return .continueWalk

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceScopeFixup.swift

@@ -685,7 +685,7 @@ extension ScopeExtension {
     do {
       // The innermost scope that must be extended must dominate all uses.
       var walker = LifetimeDependentUseWalker(function, localReachabilityCache, context) {
-        inRangeUses.append($0.instruction)
+        inRangeUses.append($0)
         return .continueWalk
       }
       defer {walker.deinitialize()}
@@ -1064,15 +1064,15 @@ private extension BeginApplyInst {
 private struct LifetimeDependentUseWalker : LifetimeDependenceDefUseWalker {
   let function: Function
   let context: Context
-  let visitor: (Operand) -> WalkResult
+  let visitor: (Instruction) -> WalkResult
   let localReachabilityCache: LocalVariableReachabilityCache
   var visitedValues: ValueSet
 
   /// Set to true if the dependence is returned from the current function.
   var dependsOnCaller = false
 
   init(_ function: Function, _ localReachabilityCache: LocalVariableReachabilityCache, _ context: Context,
-       visitor: @escaping (Operand) -> WalkResult) {
+       visitor: @escaping (Instruction) -> WalkResult) {
     self.function = function
     self.context = context
     self.visitor = visitor
@@ -1091,42 +1091,42 @@ private struct LifetimeDependentUseWalker : LifetimeDependenceDefUseWalker {
   mutating func deadValue(_ value: Value, using operand: Operand?)
   -> WalkResult {
     if let operand {
-      return visitor(operand)
+      return visitor(operand.instruction)
     }
     return .continueWalk
   }
 
   mutating func leafUse(of operand: Operand) -> WalkResult {
-    return visitor(operand)
+    return visitor(operand.instruction)
   }
 
   mutating func escapingDependence(on operand: Operand) -> WalkResult {
     log(">>> Escaping dependence: \(operand)")
-    _ = visitor(operand)
+    _ = visitor(operand.instruction)
     // Make a best-effort attempt to extend the access scope regardless of escapes. It is possible that some mandatory
     // pass between scope fixup and diagnostics will make it possible for the LifetimeDependenceDefUseWalker to analyze
     // this use.
     return .continueWalk
   }
 
-  mutating func inoutDependence(argument: FunctionArgument, on operand: Operand) -> WalkResult {
+  mutating func inoutDependence(argument: FunctionArgument, functionExit: Instruction) -> WalkResult {
     dependsOnCaller = true
-    return visitor(operand)
+    return visitor(functionExit)
   }
 
   mutating func returnedDependence(result operand: Operand) -> WalkResult {
     dependsOnCaller = true
-    return visitor(operand)
+    return visitor(operand.instruction)
   }
 
   mutating func returnedDependence(address: FunctionArgument,
                                    on operand: Operand) -> WalkResult {
     dependsOnCaller = true
-    return visitor(operand)
+    return visitor(operand.instruction)
   }
 
   mutating func yieldedDependence(result: Operand) -> WalkResult {
-    return visitor(result)
+    return visitor(result.instruction)
   }
 
   mutating func storeToYieldDependence(address: Value, of operand: Operand) -> WalkResult {

SwiftCompilerSources/Sources/Optimizer/Utilities/AddressUtils.swift

@@ -279,7 +279,7 @@ extension AccessBase {
     default:
       return nil
     }
-    return AddressInitializationWalker.findSingleInitializer(ofAddress: baseAddr, context: context)
+    return AddressInitializationWalker.findSingleInitializer(ofAddress: baseAddr, requireFullyAssigned: .value, context)
   }
 }
 
@@ -311,25 +311,31 @@ extension AccessBase {
 // modification of memory.
 struct AddressInitializationWalker: AddressDefUseWalker, AddressUseVisitor {
   let baseAddress: Value
+  let requireFullyAssigned: IsFullyAssigned
+  let onRead: WalkResult
   let context: any Context
 
   var walkDownCache = WalkerCache<SmallProjectionPath>()
 
   var isProjected = false
   var initializer: AccessBase.Initializer?
 
-  static func findSingleInitializer(ofAddress baseAddr: Value, context: some Context)
+  static func findSingleInitializer(ofAddress baseAddr: Value, requireFullyAssigned: IsFullyAssigned,
+                                    allowRead: Bool = true, _ context: some Context)
     -> AccessBase.Initializer? {
 
-    var walker = AddressInitializationWalker(baseAddress: baseAddr, context)
+    var walker = AddressInitializationWalker(baseAddress: baseAddr, requireFullyAssigned, allowRead: allowRead, context)
     if walker.walkDownUses(ofAddress: baseAddr, path: SmallProjectionPath()) == .abortWalk {
       return nil
     }
     return walker.initializer
   }
 
-  private init(baseAddress: Value, _ context: some Context) {
+  private init(baseAddress: Value, _ requireFullyAssigned: IsFullyAssigned, allowRead: Bool, _ context: some Context) {
+    assert(requireFullyAssigned != .no)
     self.baseAddress = baseAddress
+    self.requireFullyAssigned = requireFullyAssigned
+    self.onRead = allowRead ? .continueWalk : .abortWalk
     self.context = context
     if let arg = baseAddress as? FunctionArgument {
       assert(!arg.convention.isIndirectIn, "@in arguments cannot be initialized")
@@ -387,12 +393,26 @@ extension AddressInitializationWalker {
     // FIXME: check mayWriteToMemory but ignore non-stores. Currently,
     // stores should all be checked my isAddressInitialization, but
     // this is not robust.
-    return .continueWalk
+    return onRead
   }
 
   mutating func appliedAddressUse(of operand: Operand, by apply: FullApplySite)
     -> WalkResult {
-    if operand.isAddressInitialization {
+    switch apply.fullyAssigns(operand: operand) {
+    case .no:
+      if onRead == .abortWalk {
+        return .abortWalk
+      }
+      break
+    case .lifetime:
+      if onRead == .abortWalk {
+        return .abortWalk
+      }
+      if requireFullyAssigned == .value  {
+        break
+      }
+      fallthrough
+    case .value:
       return setInitializer(instruction: operand.instruction)
     }
     guard let convention = apply.convention(of: operand) else {
@@ -403,26 +423,26 @@ extension AddressInitializationWalker {
 
   mutating func loadedAddressUse(of operand: Operand, intoValue value: Value)
     -> WalkResult {
-    return .continueWalk
+    return onRead
   }
 
   mutating func loadedAddressUse(of operand: Operand, intoAddress address: Operand)
     -> WalkResult {
-    return .continueWalk
+    return onRead
   }
 
   mutating func yieldedAddressUse(of operand: Operand) -> WalkResult {
     // An inout yield is a partial write. Initialization via coroutine is not supported, so we assume a prior
     // initialization must dominate the yield.
-    return .continueWalk
+    return onRead
   }
 
   mutating func dependentAddressUse(of operand: Operand, dependentValue value: Value) -> WalkResult {
-    return .continueWalk
+    return onRead
   }
 
   mutating func dependentAddressUse(of operand: Operand, dependentAddress address: Value) -> WalkResult {
-    return .continueWalk
+    return onRead
   }
 
   mutating func escapingAddressUse(of operand: Operand) -> WalkResult {
@@ -629,7 +649,7 @@ extension AddressOwnershipLiveRange {
     var reachableUses = Stack<LocalVariableAccess>(context)
     defer { reachableUses.deinitialize() }
 
-    localReachability.gatherKnownLifetimeUses(from: assignment, in: &reachableUses)
+    localReachability.gatherKnownLivenessUses(from: assignment, in: &reachableUses)
 
     let assignmentInst = assignment.instruction ?? allocation.parentFunction.entryBlock.instructions.first!
     var range = InstructionRange(begin: assignmentInst, context)

SwiftCompilerSources/Sources/Optimizer/Utilities/FunctionSignatureTransforms.swift

@@ -26,7 +26,7 @@ private extension ArgumentConventions {
 
     // Check if `argIndex` is a lifetime source in parameterDependencies
     for targetIndex in firstParameterIndex..<self.count {
-      if getDependence(target: targetIndex, source: argIndex) != nil {
+      if parameterDependence(targetArgumentIndex: targetIndex, sourceArgumentIndex: argIndex) != nil {
         return true
       }
     }