Redesign LifetimeDependenceInsertion for coroutines

Generalize the design. Extend it to handle different kinds of coroutine
dependencies: address/value, scoped/inherited.

Author: atrick

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceInsertion.swift

@@ -98,107 +98,133 @@ extension LifetimeDependentApply {
 
   /// A lifetime argument that either inherits or creates a new scope for the lifetime of the argument value.
   struct LifetimeSource {
+    let targetKind: TargetKind
     let convention: LifetimeDependenceConvention
     let value: Value
   }
 
   /// List of lifetime dependencies for a single target.
-  struct LifetimeSources {
-    let targetKind: TargetKind
+  struct LifetimeSourceInfo {
     var sources = SingleInlineArray<LifetimeSource>()
+    var bases = [Value]()
   }
 
-  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)
+  func getResultDependenceSources() -> LifetimeSourceInfo? {
+    guard applySite.hasResultDependence else {
+      return nil
+    }
+    var info = LifetimeSourceInfo()
+    if let beginApply = applySite as? BeginApplyInst {
+      return getYieldDependenceSources(beginApply: beginApply)
+    }
+    for operand in applySite.parameterOperands {
+      guard let dep = applySite.resultDependence(on: operand) else {
+        continue
       }
-    default:
-      sources = LifetimeSources(targetKind: .result)
+      info.sources.push(LifetimeSource(targetKind: .result, convention: dep, value: operand.value))
+    }
+    return info
+  }
+
+  func getYieldDependenceSources(beginApply: BeginApplyInst) -> LifetimeSourceInfo? {
+    var info = LifetimeSourceInfo()
+    let hasScopedYield = applySite.parameterOperands.contains {
+      if let dep = applySite.resultDependence(on: $0) {
+        return dep == .scope
+      }
+      return false
+    }
+    if hasScopedYield {
+      // for consistency, we you yieldAddress if any yielded value is an address.
+      let targetKind = beginApply.yieldedValues.contains(where: { $0.type.isAddress })
+        ? TargetKind.yieldAddress : TargetKind.yield
+      info.sources.push(LifetimeSource(targetKind: targetKind, convention: .scope, value: beginApply.token))
     }
     for operand in applySite.parameterOperands {
       guard let dep = applySite.resultDependence(on: operand) else {
         continue
       }
-      sources.sources.push(LifetimeSource(convention: dep, value: operand.value))
+      switch dep {
+      case .inherit:
+        continue
+      case .scope:
+        for yieldedValue in beginApply.yieldedValues {
+          let targetKind = yieldedValue.type.isAddress ? TargetKind.yieldAddress : TargetKind.yield
+          info.sources.push(LifetimeSource(targetKind: targetKind, convention: .inherit, value: operand.value))
+        }
+      }
     }
-    return sources
+    return info
   }
 
-  func getParameterDependenceSources(target: Operand) -> LifetimeSources? {
+  func getParameterDependenceSources(target: Operand) -> LifetimeSourceInfo? {
     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")
-    }
+    var info = LifetimeSourceInfo()
+    let targetKind = {
+      let convention = applySite.convention(of: target)!
+      switch convention {
+      case .indirectInout, .indirectInoutAliasable, .packInout:
+        return TargetKind.inoutParameter
+      case .indirectIn, .indirectInGuaranteed, .indirectInCXX, .directOwned, .directUnowned, .directGuaranteed,
+           .packOwned, .packGuaranteed:
+        return 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))
+      info.sources.push(LifetimeSource(targetKind: targetKind, convention: dep, value: operand.value))
     }
-    return sources
+    return info
   }
+}
 
-  // 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 {
+private extension LifetimeDependentApply.LifetimeSourceInfo {
+  mutating func initializeBases(_ context: FunctionPassContext) {
+    for source in sources {
+      // 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.
       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
-          }
-        }
+        break
       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)
-          }
+        initializeScopedBases(source: source, context)
+      }
+    }
+  }
+
+  // Scoped dependencies require a mark_dependence for every variable that introduces this scope.
+  mutating func initializeScopedBases(source: LifetimeDependentApply.LifetimeSource, _ context: FunctionPassContext) {
+    switch source.targetKind {
+    case .yield, .yieldAddress:
+      // A coroutine creates its own borrow scope, nested within its borrowed operand.
+      bases.append(source.value)
+    case .result, .inParameter, .inoutParameter:
+      // 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 bases
   }
 }
 
@@ -207,16 +233,17 @@ extension LifetimeDependentApply {
 /// result on each argument so that the result is recognized as a
 /// dependent value within each scope.
 private func insertResultDependencies(for apply: LifetimeDependentApply, _ context: FunctionPassContext ) {
-  guard let sources = apply.getResultDependenceSources() else {
+  guard var sources = apply.getResultDependenceSources() else {
     return
   }
   log("Creating dependencies for \(apply.applySite)")
 
-  let bases = LifetimeDependentApply.findDependenceBases(sources: sources, context)
+  // Find the dependence base for each source.
+  sources.initializeBases(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: sources.bases, builder: builder, context)
   }
   for resultOper in apply.applySite.indirectResultOperands {
     let accessBase = resultOper.value.accessBase
@@ -231,23 +258,23 @@ private func insertResultDependencies(for apply: LifetimeDependentApply, _ conte
     }
     assert(initializingStore == resultOper.instruction, "an indirect result is a store")
     Builder.insert(after: apply.applySite, context) { builder in
-      insertMarkDependencies(value: initialAddress, initializer: initializingStore, bases: bases, builder: builder,
-                             context)
+      insertMarkDependencies(value: initialAddress, initializer: initializingStore, bases: sources.bases,
+                             builder: builder, context)
     }
   }
 }
 
 private func insertParameterDependencies(apply: LifetimeDependentApply, target: Operand,
                                          _ context: FunctionPassContext ) {
-  guard let sources = apply.getParameterDependenceSources(target: target) else {
+  guard var sources = apply.getParameterDependenceSources(target: target) else {
     return
   }
   log("Creating dependencies for \(apply.applySite)")
 
-  let bases = LifetimeDependentApply.findDependenceBases(sources: sources, context)
+  sources.initializeBases(context)
 
   Builder.insert(after: apply.applySite, context) {
-    insertMarkDependencies(value: target.value, initializer: nil, bases: bases, builder: $0, context)    
+    insertMarkDependencies(value: target.value, initializer: nil, bases: sources.bases, builder: $0, context)
   }
 }
 
@@ -259,11 +286,18 @@ private func insertMarkDependencies(value: Value, initializer: Instruction?,
     let markDep = builder.createMarkDependence(
       value: currentValue, base: base, kind: .Unresolved)
 
-    let uses = currentValue.uses.lazy.filter {
-      let inst = $0.instruction
-      return inst != markDep && inst != initializer && !(inst is Deallocation)
+    // Address dependencies cannot be represented as SSA values, so it doesn not make sense to replace any uses of the
+    // dependent address. TODO: consider a separate mark_dependence_addr instruction since the semantics are different.
+    if !value.type.isAddress {
+      let uses = currentValue.uses.lazy.filter {
+        if $0.isScopeEndingUse {
+          return false
+        }
+        let inst = $0.instruction
+        return inst != markDep && inst != initializer && !(inst is Deallocation)
+      }
+      uses.replaceAll(with: markDep, context)
     }
-    uses.replaceAll(with: markDep, context)
     currentValue = markDep
   }
 }