ComputeSideEffects: correctly handle escaping arguments

If an argument escapes in a called function, we don't know anything about the argument's side effects.
For example, it could escape to the return value and effects might occur in the caller.

Fixes a miscompile
#73477
rdar://127691335

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ComputeSideEffects.swift

@@ -215,7 +215,7 @@ private struct CollectedEffects {
   }
 
   mutating func addEffectsForEscapingArgument(argument: FunctionArgument) {
-    var escapeWalker = ArgumentEscapingWalker()
+    var escapeWalker = ArgumentEscapingWalker(context)
 
     if escapeWalker.hasUnknownUses(argument: argument) {
       // Worst case: we don't know anything about how the argument escapes.
@@ -414,13 +414,18 @@ private func defaultPath(for value: Value) -> SmallProjectionPath {
 /// Checks if an argument escapes to some unknown user.
 private struct ArgumentEscapingWalker : ValueDefUseWalker, AddressDefUseWalker {
   var walkDownCache = WalkerCache<UnusedWalkingPath>()
+  private let calleeAnalysis: CalleeAnalysis
 
   /// True if the argument escapes to a load which (potentially) "takes" the memory location.
   private(set) var foundTakingLoad = false
 
   /// True, if the argument escapes to a closure context which might be destroyed when called.
   private(set) var foundConsumingPartialApply = false
 
+  init(_ context: FunctionPassContext) {
+    self.calleeAnalysis = context.calleeAnalysis
+  }
+
   mutating func hasUnknownUses(argument: FunctionArgument) -> Bool {
     if argument.type.isAddress {
       return walkDownUses(ofAddress: argument, path: UnusedWalkingPath()) == .abortWalk
@@ -444,14 +449,23 @@ private struct ArgumentEscapingWalker : ValueDefUseWalker, AddressDefUseWalker {
       return .continueWalk
 
     case let apply as ApplySite:
-      if apply.isCallee(operand: value) {
-        // `CollectedEffects.handleApply` only handles argument operands of an apply, but not the callee operand.
-        return .abortWalk
-      }
       if let pa = apply as? PartialApplyInst, !pa.isOnStack {
         foundConsumingPartialApply = true
       }
-      return .continueWalk
+      // `CollectedEffects.handleApply` only handles argument operands of an apply, but not the callee operand.
+      if let calleeArgIdx = apply.calleeArgumentIndex(of: value),
+         let callees = calleeAnalysis.getCallees(callee: apply.callee)
+      {
+        // If an argument escapes in a called function, we don't know anything about the argument's side effects.
+        // For example, it could escape to the return value and effects might occur in the caller.
+        for callee in callees {
+          if callee.effects.escapeEffects.canEscape(argumentIndex: calleeArgIdx, path: SmallProjectionPath.init(.anyValueFields)) {
+            return .abortWalk
+          }
+        }
+        return .continueWalk
+      }
+      return .abortWalk
     default:
       return .abortWalk
     }

test/SILOptimizer/side_effects.sil

@@ -87,10 +87,11 @@ sil @$s4test1ZCfD : $@convention(method) (@owned Z) -> () {
 // CHECK-LABEL: sil @load_store_to_args
 // CHECK-NEXT:  [%0: read v**]
 // CHECK-NEXT:  [%1: write v**]
-// CHECK-NEXT:  [%2: write c0.v**]
+// CHECK-NEXT:  [%2: noescape **, write c0.v**]
 // CHECK-NEXT:  [global: ]
 // CHECK-NEXT:  {{^[^[]}}
 sil @load_store_to_args : $@convention(thin) (@inout Int32, @inout Int32, @guaranteed X) -> () {
+[%2: noescape **]
 bb0(%0 : $*Int32, %1 : $*Int32, %2 : $X):
   %l1 = load %0 : $*Int32
   store %l1 to %1  : $*Int32
@@ -971,11 +972,13 @@ bb0(%0 : @guaranteed $S):
 }
 
 // CHECK-LABEL: sil [ossa] @storeToArgs
-// CHECK-NEXT:  [%0: write c0.v**]
-// CHECK-NEXT:  [%1: write c0.v**]
+// CHECK-NEXT:  [%0: noescape **, write c0.v**]
+// CHECK-NEXT:  [%1: noescape **, write c0.v**]
 // CHECK-NEXT:  [global: ]
 // CHECK-NEXT:  {{^[^[]}}
 sil [ossa] @storeToArgs : $@convention(thin) (@guaranteed List, @guaranteed List) -> () {
+[%0: noescape **]
+[%1: noescape **]
 bb0(%1 : @guaranteed $List, %2 : @guaranteed $List):
   cond_br undef, bb1, bb2
 
@@ -1195,3 +1198,39 @@ bb0(%0 : $*T):
   %r = tuple()
   return %r : $()
 }
+
+sil @store_owned_to_out : $@convention(thin) (@owned SP) -> @out SP {
+[%0: noescape **, write v**]
+[%1: escape v** -> %0.v**]
+[global: ]
+}  
+
+// CHECK-LABEL: sil @test_escaping_arg1
+// CHECK-NEXT:  [%0: write v**.c*.v**, copy v**.c*.v**, destroy v**.c*.v**]
+// CHECK-NEXT:  [global: write,copy,destroy]
+// CHECK-NEXT:  {{^[^[]}}
+sil @test_escaping_arg1 : $@convention(thin) (@owned SP) -> () {
+bb0(%0 : $SP):
+  %1 = alloc_stack $SP
+  %2 = function_ref @store_owned_to_out : $@convention(thin) (@owned SP) -> @out SP
+  %3 = apply %2(%1, %0) : $@convention(thin) (@owned SP) -> @out SP
+  %4 = struct_element_addr %1 : $*SP, #SP.value
+  %5 = load %4 : $*X
+  strong_release %5 : $X
+  dealloc_stack %1 : $*SP
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil @test_escaping_arg2
+// CHECK-NEXT:  [%0: read v**.c*.v**, write v**.c*.v**, copy v**.c*.v**, destroy v**.c*.v**]
+// CHECK-NEXT:  [global: read,write,copy,destroy,allocate,deinit_barrier]
+// CHECK-NEXT:  {{^[^[]}}
+sil @test_escaping_arg2 : $@convention(thin) (@owned SP) -> () {
+bb0(%0 : $SP):
+  %1 = function_ref @forward_to_return : $@convention(thin) (@owned SP) -> @owned SP
+  %2 = apply %1(%0) : $@convention(thin) (@owned SP) -> @owned SP
+  release_value %2 : $SP
+  %r = tuple ()
+  return %r : $()
+}