LifetimeDependence: diagnose yield and store-to-yield.

Author: atrick

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceDiagnostics.swift

@@ -131,7 +131,7 @@ private struct DiagnoseDependence {
     reportEscaping(operand: operand)
   }
 
-  func checkInoutResult(argument inoutArg: FunctionArgument, result operand: Operand) -> WalkResult {
+  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.
@@ -140,13 +140,28 @@ private struct DiagnoseDependence {
       }
       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
       }
     }
-    reportEscaping(operand: operand)
     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 {
@@ -174,9 +189,9 @@ private struct DiagnoseDependence {
       // 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
   }
 
@@ -328,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
@@ -380,20 +439,44 @@ extension DiagnoseDependenceWalker : LifetimeDependenceDefUseWalker {
   }
 
   mutating func inoutDependence(argument: FunctionArgument, on operand: Operand) -> WalkResult {
-    return diagnostics.checkInoutResult(argument: argument, result: operand)
+    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,
                                    on operand: Operand) -> WalkResult {
-    return diagnostics.checkFunctionResult(operand: operand)
+    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/LifetimeDependenceScopeFixup.swift

@@ -253,5 +253,9 @@ private struct LifetimeDependenceScopeFixupWalker : LifetimeDependenceDefUseWalk
   mutating func yieldedDependence(result: Operand) -> WalkResult {
     return .continueWalk
   }
+
+  mutating func storeToYieldDependence(address: Value, of operand: Operand) -> WalkResult {
+    return .continueWalk
+  }
 }
 

SwiftCompilerSources/Sources/Optimizer/Utilities/AddressUtils.swift

@@ -64,6 +64,10 @@ protocol AddressUseVisitor {
   mutating func loadedAddressUse(of operand: Operand, into address: Operand)
     -> WalkResult
 
+  /// Yielding an address may modify the value at the yield, but not past the yield. The yielded value may escape, but
+  /// only if its type is escapable.
+  mutating func yieldedAddressUse(of operand: Operand) -> WalkResult
+
   /// A non-address owned `value` whose ownership depends on the in-memory
   /// value at `address`, such as `mark_dependence %value on %address`.
   mutating func dependentAddressUse(of operand: Operand, into value: Value)
@@ -119,7 +123,7 @@ extension AddressUseVisitor {
     case let pai as PartialApplyInst where pai.mayEscape:
       return escapingAddressUse(of: operand)
 
-    case is ReturnInst, is ThrowInst, is YieldInst, is AddressToPointerInst:
+    case is ThrowInst, is AddressToPointerInst:
       return escapingAddressUse(of: operand)
 
     case is StructElementAddrInst, is TupleElementAddrInst,
@@ -149,12 +153,14 @@ extension AddressUseVisitor {
          is PackElementSetInst:
       return leafAddressUse(of: operand)
 
-    case is LoadInst, is LoadUnownedInst,  is LoadWeakInst, 
-         is ValueMetatypeInst, is ExistentialMetatypeInst,
+    case is LoadInst, is LoadUnownedInst,  is LoadWeakInst, is ValueMetatypeInst, is ExistentialMetatypeInst,
          is PackElementGetInst:
       let svi = operand.instruction as! SingleValueInstruction
       return loadedAddressUse(of: operand, into: svi)
 
+    case is YieldInst:
+      return yieldedAddressUse(of: operand)
+
     case let sdai as SourceDestAddrInstruction
            where sdai.sourceOperand == operand:
       return loadedAddressUse(of: operand, into: sdai.destinationOperand)
@@ -334,6 +340,11 @@ extension AddressInitializationWalker {
     return .continueWalk
   }
 
+  mutating func yieldedAddressUse(of operand: Operand) -> WalkResult {
+    // An inout yield is a partial write.
+    return .abortWalk
+  }
+
   mutating func dependentAddressUse(of operand: Operand, into value: Value)
     -> WalkResult {
     return .continueWalk