EscapeInfo: improve handling of begin_applys which yield an address value

Instead of just aborting if a `begin_apply` gets an address argument, follow the yielded values and see if they actually escape

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/Utilities/EscapeUtils.swift

@@ -615,10 +615,14 @@ fileprivate struct EscapeWalker<V: EscapeVisitor> : ValueDefUseWalker,
     }
 
     // Indirect arguments cannot escape the function, but loaded values from such can.
-    if !followLoads(at: path) &&
-       // Except for begin_apply: it can yield an address value.
-       !apply.isBeginApplyWithIndirectResults {
-      return .continueWalk
+    if !followLoads(at: path) {
+      guard let beginApply = apply as? BeginApplyInst else {
+        return .continueWalk
+      }
+      // Except for begin_apply: it can yield an address value.
+      if !indirectResultEscapes(of: beginApply, path: path) {
+        return .continueWalk
+      }
     }
 
     if argOp.value.type.isNoEscapeFunction {
@@ -658,7 +662,16 @@ fileprivate struct EscapeWalker<V: EscapeVisitor> : ValueDefUseWalker,
     }
     return .continueWalk
   }
-  
+
+  private mutating func indirectResultEscapes(of beginApply: BeginApplyInst, path: Path) -> Bool {
+    for result in beginApply.yieldedValues where result.type.isAddress {
+      if walkDownUses(ofAddress: result, path: path) == .abortWalk {
+        return true
+      }
+    }
+    return false
+  }
+
   /// Handle `.escaping` effects for an apply argument during the walk-down.
   private mutating
   func walkDownArgument(calleeArgIdx: Int, argPath: Path,
@@ -928,13 +941,3 @@ private extension SmallProjectionPath {
     EscapeUtilityTypes.EscapePath(projectionPath: self, followStores: false, addressIsStored: false, knownType: nil)
   }
 }
-
-private extension ApplySite {
-  var isBeginApplyWithIndirectResults: Bool {
-    guard let ba = self as? BeginApplyInst else {
-      return false
-    }
-    // Note that the token result is always a non-address type.
-    return ba.results.contains { $0.type.isAddress }
-  }
-}

test/SILOptimizer/mem-behavior.sil

@@ -380,6 +380,24 @@ bb2:
   unreachable
 }
 
+// CHECK-LABEL:  @test_indirect_yield
+// CHECK:        PAIR #3.
+// CHECK-NEXT:       store %0 to %4 : $*Int32
+// CHECK-NEXT:       %1 = alloc_stack $Int32
+// CHECK-NEXT:     r=0,w=1
+sil @test_indirect_yield : $@convention(thin) (Int32) -> () {
+bb0(%0 : $Int32):
+  %1 = alloc_stack $Int32
+  store %0 to %1 : $*Int32
+  %3 = function_ref @indirect_yield_coroutine : $@yield_once @convention(thin) (@inout Int32) -> @yields @inout Int32
+  (%4, %5) = begin_apply %3(%1) : $@yield_once @convention(thin) (@inout Int32) -> @yields @inout Int32
+  store %0 to %4 : $*Int32
+  end_apply %5 as $()
+  dealloc_stack %1 : $*Int32
+  %9 = tuple ()
+  return %9 : $()
+}
+
 // CHECK-LABEL: @refelementaddr_and_reference
 // CHECK:       PAIR #0.
 // CHECK-NEXT:      %3 = apply %2(%0) : $@convention(thin) (@guaranteed X) -> Int32