Merge pull request #75162 from atrick/inout-dep

LifetimeDependenceDiagnostics support for inout reassignment.

Author: atrick

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceDiagnostics.swift

@@ -131,6 +131,37 @@ private struct DiagnoseDependence {
     reportEscaping(operand: operand)
   }
 
+  func checkInoutResult(argument inoutArg: FunctionArgument) -> WalkResult {
+    // Check that the parameter dependence for this inout argument is the same as the current dependence scope.
+    if let sourceArg = dependence.scope.parentValue as? FunctionArgument {
+      // If the inout result is also the inout source, then it's always ok.
+      if inoutArg == sourceArg {
+        return .continueWalk
+      }
+      if function.argumentConventions.getDependence(target: inoutArg.index, source: 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
+      }
+    }
+    return .abortWalk
+  }
+
+  func checkStoreToYield(address: Value) -> WalkResult {
+    var walker = DependentAddressUseDefWalker(context: context, diagnostics: self)
+    return walker.walkUp(address: address)
+  }
+
+  func checkYield(operand: Operand) -> WalkResult {
+    switch dependence.scope {
+    case .caller:
+      return checkFunctionResult(operand: operand)
+    default:
+      // local scopes can be yielded without escaping.
+      return .continueWalk
+    }
+  }
+
   func checkFunctionResult(operand: Operand) -> WalkResult {
 
     if function.hasUnsafeNonEscapableResult {
@@ -151,16 +182,16 @@ private struct DiagnoseDependence {
     if dependence.isUnsafeApplyResult, function.hasResultDependence {
       return .continueWalk
     }
-    // Check that the argument dependence for this result is the same
+    // Check that the parameter dependence for this result is the same
     // as the current dependence scope.
     if let arg = dependence.scope.parentValue as? FunctionArgument,
        function.argumentConventions[resultDependsOn: arg.index] != nil {
       // The returned value depends on a lifetime that is inherited or
       // borrowed in the caller. The lifetime of the argument value
       // itself is irrelevant here.
+      log("  has dependent function result")
       return .continueWalk
     }
-    reportEscaping(operand: operand)
     return .abortWalk
   }
 
@@ -312,6 +343,50 @@ private struct LifetimeVariable {
   }
 }
 
+/// Walk up an address into which a dependent value has been stored. If any address in the use-def chain is a
+/// mark_dependence, follow the depenedence base rather than the forwarded value. If any of the dependence bases in
+/// within the current scope is with (either local checkInoutResult), then storing a value into that address is
+/// nonescaping.
+///
+/// This supports store-to-yield. Storing to a yield is an escape unless the yielded memory location depends on another
+/// lifetime that already depends on the current scope. When setter depends on 'newValue', 'newValue' is stored to the
+/// yielded address, and the yielded addrses depends on the lifetime of 'self'. A mark_dependence should have already
+/// been inserted for that lifetime depenence:
+///
+///   (%a, %t) = begin_apply %f(%self)
+///              : $@yield_once @convention(method) (@inout Self) -> _inherit(0) @yields @inout Self.field
+///   %dep = mark_dependence [nonescaping] %yield_addr on %self
+///   store %newValue to [assign] %dep : $*Self.field
+///
+private struct DependentAddressUseDefWalker {
+  let context: Context
+  var diagnostics: DiagnoseDependence
+}
+
+extension DependentAddressUseDefWalker: AddressUseDefWalker {
+  // Follow the dependence base, not the forwarded value. Similar to the way LifetimeDependenceUseDefWalker handles
+  // MarkDependenceInst.
+  mutating func walkUp(address: Value, path: UnusedWalkingPath = UnusedWalkingPath()) -> WalkResult {
+    if let markDep = address as? MarkDependenceInst, let addressDep = LifetimeDependence(markDep, context) {
+      switch addressDep.scope {
+      case let .caller(arg):
+        return diagnostics.checkInoutResult(argument: arg)
+      case .owned, .initialized:
+        // Storing a nonescaping value to local memory cannot escape.
+        return .abortWalk
+      default:
+        break
+      }
+    }
+    return walkUpDefault(address: address, path: UnusedWalkingPath())
+  }
+
+  mutating func rootDef(address: Value, path: UnusedWalkingPath) -> WalkResult {
+    // This only searches for mark_dependence scopes.
+    return .continueWalk
+  }
+}
+
 /// Walk down lifetime depenence uses. For each check that all dependent
 /// leaf uses are non-escaping and within the dependence scope. The walk
 /// starts with add address for .access dependencies. The walk can
@@ -363,17 +438,45 @@ extension DiagnoseDependenceWalker : LifetimeDependenceDefUseWalker {
     return .abortWalk
   }
 
+  mutating func inoutDependence(argument: FunctionArgument, on operand: Operand) -> WalkResult {
+    if diagnostics.checkInoutResult(argument: argument) == .abortWalk {
+      diagnostics.reportEscaping(operand: operand)
+      return .abortWalk
+    }
+    return .continueWalk
+  }
+
   mutating func returnedDependence(result: Operand) -> WalkResult {
-    return diagnostics.checkFunctionResult(operand: result)
+    if diagnostics.checkFunctionResult(operand: result) == .abortWalk {
+      diagnostics.reportEscaping(operand: result)
+      return .abortWalk
+    }
+    return .continueWalk
   }
 
   mutating func returnedDependence(address: FunctionArgument,
-                                   using operand: Operand) -> WalkResult {
-    return diagnostics.checkFunctionResult(operand: operand)
+                                   on operand: Operand) -> WalkResult {
+    if diagnostics.checkFunctionResult(operand: operand) == .abortWalk {
+      diagnostics.reportEscaping(operand: operand)
+      return .abortWalk
+    }
+    return .continueWalk
   }
 
   mutating func yieldedDependence(result: Operand) -> WalkResult {
-    return diagnostics.checkFunctionResult(operand: result)
+    if diagnostics.checkYield(operand: result) == .abortWalk {
+      diagnostics.reportEscaping(operand: result)
+      return .abortWalk
+    }
+    return .continueWalk
+  }
+
+  mutating func storeToYieldDependence(address: Value, of operand: Operand) -> WalkResult {
+    if diagnostics.checkStoreToYield(address: address) == .abortWalk {
+      diagnostics.reportEscaping(operand: operand)
+      return .abortWalk
+    }
+    return .continueWalk
   }
 
   // Override AddressUseVisitor here because LifetimeDependenceDefUseWalker

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceInsertion.swift

@@ -24,7 +24,7 @@ private let verbose = false
 
 private func log(prefix: Bool = true, _ message: @autoclosure () -> String) {
   if verbose {
-    print((prefix ? "### " : "") + message())
+    debugLog(prefix: prefix, message())
   }
 }
 
@@ -40,7 +40,10 @@ let lifetimeDependenceInsertionPass = FunctionPass(
 
   for instruction in function.instructions {
     if let dependentApply = LifetimeDependentApply(instruction) {
-      insertDependencies(for: dependentApply, context)
+      for operand in dependentApply.applySite.parameterOperands {
+        insertParameterDependencies(apply: dependentApply, target: operand, context)
+      }
+      insertResultDependencies(for: dependentApply, context)
     }
   }
 }
@@ -53,7 +56,7 @@ private struct LifetimeDependentApply {
     guard let apply = instruction as? FullApplySite else {
       return nil
     }
-    if !apply.hasResultDependence {
+    if !apply.hasLifetimeDependence {
       return nil
     }
     self.applySite = apply
@@ -85,50 +88,145 @@ private struct LifetimeDependentApply {
 }
 
 extension LifetimeDependentApply {
-  /// A lifetime argument copies, borrows, or mutatably borrows the
-  /// lifetime of the argument value.
-  struct LifetimeArgument {
+  enum TargetKind {
+    case result
+    case inParameter
+    case inoutParameter
+    case yield
+    case yieldAddress
+  }
+  
+  /// A lifetime argument that either inherits or creates a new scope for the lifetime of the argument value.
+  struct LifetimeSource {
     let convention: LifetimeDependenceConvention
     let value: Value
   }
 
-  func getLifetimeArguments() -> SingleInlineArray<LifetimeArgument> {
-    var args = SingleInlineArray<LifetimeArgument>()
+  /// List of lifetime dependencies for a single target.
+  struct LifetimeSources {
+    let targetKind: TargetKind
+    var sources = SingleInlineArray<LifetimeSource>()
+  }
+
+  func getResultDependenceSources() -> LifetimeSources? {
+    guard applySite.hasResultDependence else { return nil }
+    var sources: LifetimeSources
+    switch applySite {
+    case let beginApply as BeginApplyInst:
+      if beginApply.yieldedValues.contains(where: { $0.type.isAddress }) {
+        sources = LifetimeSources(targetKind: .yieldAddress)
+      } else {
+        sources = LifetimeSources(targetKind: .yield)
+      }
+    default:
+      sources = LifetimeSources(targetKind: .result)
+    }
     for operand in applySite.parameterOperands {
       guard let dep = applySite.resultDependence(on: operand) else {
         continue
       }
-      args.push(LifetimeArgument(convention: dep, value: operand.value))
+      sources.sources.push(LifetimeSource(convention: dep, value: operand.value))
+    }
+    return sources
+  }
+
+  func getParameterDependenceSources(target: Operand) -> LifetimeSources? {
+    guard let deps = applySite.parameterDependencies(target: target) else {
+      return nil
+    }
+    var sources: LifetimeSources
+    let convention = applySite.convention(of: target)!
+    switch convention {
+    case .indirectInout, .indirectInoutAliasable, .packInout:
+      sources = LifetimeSources(targetKind: .inoutParameter)
+    case .indirectIn, .indirectInGuaranteed, .indirectInCXX, .directOwned, .directUnowned, .directGuaranteed,
+         .packOwned, .packGuaranteed:
+      sources = LifetimeSources(targetKind: .inParameter)
+    case .indirectOut, .packOut:
+      debugLog("\(applySite)")
+      fatalError("Lifetime dependencies cannot target \(convention) parameter")
+    }
+    for (dep, operand) in zip(deps, applySite.parameterOperands) {
+      guard let dep = dep else {
+        continue
+      }
+      sources.sources.push(LifetimeSource(convention: dep, value: operand.value))
+    }
+    return sources
+  }
+
+  // Scoped dependencies require a mark_dependence for every variable that introduces this scope.
+  //
+  // Inherited dependencies do not require a mark_dependence if the target is a result or yielded value. The inherited
+  // lifetime is nonescapable, so either
+  //
+  // (a) the result or yield is never returned from this function
+  //
+  // (b) the inherited lifetime has a dependence root within this function (it comes from a dependent function argument
+  // or scoped dependence). In this case, when that depedence root is diagnosed, the analysis will find transtive uses
+  // of this apply's result.
+  //
+  // (c) the dependent value is passed to another call with a dependent inout argument, or it is stored to a yielded
+  // address of a coroutine that has a dependent inout argument. In this case, a mark_dependence will already be created
+  // for that inout argument.
+  //
+  // Parameter dependencies and yielded addresses always require a mark_dependence.
+  static func findDependenceBases(sources: LifetimeSources, _ context: FunctionPassContext) -> [Value] {
+    var bases: [Value] = []
+    for source in sources.sources {
+      switch source.convention {
+      case .inherit:
+        switch sources.targetKind {
+        case .result, .yield:
+          continue
+        case .inParameter, .inoutParameter, .yieldAddress:
+          _ = LifetimeDependence.visitDependenceRoots(enclosing: source.value, context) { scope in
+            log("Inherited lifetime from \(source.value)")
+            log("  scope: \(scope)")
+            bases.append(scope.parentValue)
+            return .continueWalk
+          }
+        }
+      case .scope:
+        // Create a new dependence on the apply's access to the argument.
+        for varIntoducer in gatherVariableIntroducers(for: source.value, context) {
+          if let scope = LifetimeDependence.Scope(base: varIntoducer, context) {
+            log("Scoped lifetime from \(source.value)")
+            log("  scope: \(scope)")
+            bases.append(scope.parentValue)
+          }
+        }
+      }
     }
-    return args
+    return bases
   }
 }
 
 /// If the result of this apply depends on the scope of one or more
 /// arguments, then insert a mark_dependence [unresolved] from the
 /// result on each argument so that the result is recognized as a
 /// dependent value within each scope.
-private func insertDependencies(for apply: LifetimeDependentApply,
-  _ context: FunctionPassContext ) {
-  let bases = findDependenceBases(of: apply, context)
+private func insertResultDependencies(for apply: LifetimeDependentApply, _ context: FunctionPassContext ) {
+  guard let sources = apply.getResultDependenceSources() else {
+    return
+  }
+  log("Creating dependencies for \(apply.applySite)")
+
+  let bases = LifetimeDependentApply.findDependenceBases(sources: sources, context)
+
   for dependentValue in apply.applySite.resultOrYields {
     let builder = Builder(before: dependentValue.nextInstruction, context)
-    insertMarkDependencies(value: dependentValue, initializer: nil,
-                           bases: bases, builder: builder, context)
+    insertMarkDependencies(value: dependentValue, initializer: nil, bases: bases, builder: builder, context)
   }
   for resultOper in apply.applySite.indirectResultOperands {
     let accessBase = resultOper.value.accessBase
-    guard let (initialAddress, initializingStore) =
-            accessBase.findSingleInitializer(context) else {
+    guard let (initialAddress, initializingStore) = accessBase.findSingleInitializer(context) else {
       continue
     }
-    // TODO: This is currently too strict for a diagnostic pass. We
-    // should handle/cleanup projections and casts that occur before
-    // the initializingStore. Or check in the SIL verifier that all
-    // stores without an access scope follow this form. Then convert
-    // this bail-out to an assert.
-    guard initialAddress.usesOccurOnOrAfter(instruction: initializingStore,
-                                            context) else {
+    // TODO: This might bail-out on SIL that should be diagnosed. We should handle/cleanup projections and casts that
+    // occur before the initializingStore. Or check in the SIL verifier that all stores without an access scope follow
+    // this form. Then convert this bail-out to an assert.
+    guard initialAddress.usesOccurOnOrAfter(instruction: initializingStore, context) else {
       continue
     }
     assert(initializingStore == resultOper.instruction, "an indirect result is a store")
@@ -139,29 +237,18 @@ private func insertDependencies(for apply: LifetimeDependentApply,
   }
 }
 
-private func findDependenceBases(of apply: LifetimeDependentApply,
-                                 _ context: FunctionPassContext)
-  -> [Value] {
+private func insertParameterDependencies(apply: LifetimeDependentApply, target: Operand,
+                                         _ context: FunctionPassContext ) {
+  guard let sources = apply.getParameterDependenceSources(target: target) else {
+    return
+  }
   log("Creating dependencies for \(apply.applySite)")
-  var bases: [Value] = []
-  for lifetimeArg in apply.getLifetimeArguments() {
-    switch lifetimeArg.convention {
-    case .inherit:
-      continue
-    case .scope:
-      // Create a new dependence on the apply's access to the argument.
-      for varIntoducer in gatherVariableIntroducers(for: lifetimeArg.value,
-                                                    context) {
-        if let scope =
-             LifetimeDependence.Scope(base: varIntoducer, context) {
-          log("Scoped lifetime from \(lifetimeArg.value)")
-          log("  scope: \(scope)")
-          bases.append(scope.parentValue)
-        }
-      }
-    }
+
+  let bases = LifetimeDependentApply.findDependenceBases(sources: sources, context)
+
+  Builder.insert(after: apply.applySite, context) {
+    insertMarkDependencies(value: target.value, initializer: nil, bases: bases, builder: $0, context)    
   }
-  return bases
 }
 
 private func insertMarkDependencies(value: Value, initializer: Instruction?,