Swift SIL: fix ApplySite.callerArgIndex

It needs to return nil in case the callee argument is not really applied by the ApplySite.
This fixes a crash in EscapeInfo when handling "escaping-to" effects:
If the callee operand of an apply instruction is a `partial_apply` of a `function_ref`, the callee-analysis looks through the `partial_apply` and we get the referenced function as a callee.
In this case there are less apply arguments than callee arguments, because the `partial_apply` already "applies" some of the callee arguments.

rdar://96830171

Author: eeckstein

SwiftCompilerSources/Sources/SIL/ApplySite.swift

@@ -21,8 +21,46 @@ public protocol ApplySite : Instruction {
   var operands: OperandArray { get }
   var numArguments: Int { get }
   var substitutionMap: SubstitutionMap { get }
+  
+  /// Converts an argument index of the apply to the corresponding argument index of the callee.
+  ///
+  /// For a FullApplySite this is always a 1-to-1 mapping.
+  /// For a `partial_apply` the callee index can be higher than the caller's argument index
+  /// because the arguments to `partial_apply` are a suffix of the callee.
+  ///
+  /// Example:
+  /// ```
+  /// func callee(a, b, c, d, e) { }
+  ///
+  /// %pa = partial_apply @callee(c, d, e)
+  /// // caller indices:          0, 1, 2
+  /// // callee indices:          2, 3, 4
+  ///
+  ///  %a = apply         %pa    (a, b)
+  /// // caller indices:          0, 1
+  /// // callee indices:          0, 1
+  /// ```
   func calleeArgIndex(callerArgIndex: Int) -> Int
+
+  /// Converts an argument index of a callee to the corresponding argument index of the apply.
+  ///
+  /// If the apply does not actually apply that argument, it returns nil.
+  /// Otherwise, for a FullApplySite this is always a 1-to-1 mapping.
+  /// For a `partial_apply` the caller index can be lower than the callee's argument index
+  /// because the arguments to `partial_apply` are a suffix of the callee.
+  ///
+  /// Example:
+  /// ```
+  ///                func callee(a, b, c, d, e) { }
+  /// // callee indices:         0, 1, 2, 3, 4
+  /// // caller indices in %pa:  -, -, 0, 1, 2     ("-" == nil)
+  /// // caller indices in %a:   0, 1, -, -, -
+  ///
+  /// %pa = partial_apply @callee(c, d, e)
+  ///  %a = apply         %pa    (a, b)
+  /// ```
   func callerArgIndex(calleeArgIndex: Int) -> Int?
+
   func getArgumentConvention(calleeArgIndex: Int) -> ArgumentConvention
 }
 
@@ -66,6 +104,15 @@ public protocol FullApplySite : ApplySite {
 }
 
 extension FullApplySite {
-  public func calleeArgIndex(callerArgIndex: Int) -> Int { callerArgIndex }
-  public func callerArgIndex(calleeArgIndex: Int) -> Int? { calleeArgIndex }
+  public func calleeArgIndex(callerArgIndex: Int) -> Int {
+    assert(callerArgIndex >= 0 && callerArgIndex < numArguments)
+    return callerArgIndex
+  }
+
+  public func callerArgIndex(calleeArgIndex: Int) -> Int? {
+    if calleeArgIndex < numArguments {
+      return calleeArgIndex
+    }
+    return nil
+  }
 }

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -514,7 +514,10 @@ final public class PartialApplyInst : SingleValueInstruction, ApplySite {
   public func callerArgIndex(calleeArgIndex: Int) -> Int? {
     let firstIdx = PartialApply_getCalleeArgIndexOfFirstAppliedArg(bridged)
     if calleeArgIndex >= firstIdx {
-      return calleeArgIndex - firstIdx
+      let callerIdx = calleeArgIndex - firstIdx
+      if callerIdx < numArguments {
+        return callerIdx
+      }
     }
     return nil
   }

test/SILOptimizer/escape_info.sil

@@ -968,6 +968,8 @@ bb0(%0 : @owned $Z):
 sil [ossa] [escapes !%1] @closure1 : $@convention(thin) (@guaranteed Y, @guaranteed Y) -> ()
 sil [ossa] [escapes !%0] @closure2 : $@convention(thin) (@guaranteed Y, @guaranteed Y) -> ()
 sil [escapes !%0.**] @closure3 : $@convention(thin) (@guaranteed Z) -> ()
+sil [ossa] [escapes %0 => %1, !%1] @closure4 : $@convention(thin) (@guaranteed X, @guaranteed Y) -> ()
+sil [ossa] [escapes %1 => %0, !%1] @closure5 : $@convention(thin) (@guaranteed Y, @guaranteed X) -> ()
 
 // CHECK-LABEL: Escape information for callClosure1:
 // CHECK:  -    :   %0 = alloc_ref $Y
@@ -1019,6 +1021,40 @@ bb0:
   return %6 : $()
 }
 
+// CHECK-LABEL: Escape information for callClosureWithEscapesToEffect1:
+// CHECK:  -    :   %0 = alloc_ref $Y
+// CHECK: global:   %1 = alloc_ref $X
+// CHECK: End function callClosureWithEscapesToEffect1
+sil [ossa] @callClosureWithEscapesToEffect1 : $@convention(thin) () -> () {
+bb0:
+  %0 = alloc_ref $Y
+  %1 = alloc_ref $X
+  %2 = function_ref @closure4 : $@convention(thin) (@guaranteed X, @guaranteed Y) -> ()
+  %3 = partial_apply [callee_guaranteed] %2(%0) : $@convention(thin) (@guaranteed X, @guaranteed Y) -> ()
+  %8 = apply %3(%1) : $@callee_guaranteed (@guaranteed X) -> ()
+  destroy_value %3 : $@callee_guaranteed (@guaranteed X) -> ()
+  destroy_value %1 : $X
+  %11 = tuple ()
+  return %11 : $()
+}
+
+// CHECK-LABEL: Escape information for callClosureWithEscapesToEffect2:
+// CHECK: global:   %0 = alloc_ref $Y
+// CHECK:  -    :   %1 = alloc_ref $X
+// CHECK: End function callClosureWithEscapesToEffect2
+sil [ossa] @callClosureWithEscapesToEffect2 : $@convention(thin) () -> () {
+bb0:
+  %0 = alloc_ref $Y
+  %1 = alloc_ref $X
+  %2 = function_ref @closure5 : $@convention(thin) (@guaranteed Y, @guaranteed X) -> ()
+  %3 = partial_apply [callee_guaranteed] %2(%1) : $@convention(thin) (@guaranteed Y, @guaranteed X) -> ()
+  %8 = apply %3(%0) : $@callee_guaranteed (@guaranteed Y) -> ()
+  destroy_value %3 : $@callee_guaranteed (@guaranteed Y) -> ()
+  destroy_value %0 : $Y
+  %11 = tuple ()
+  return %11 : $()
+}
+
 // CHECK-LABEL: Escape information for escape_via_destroy_of_closure:
 // CHECK:  -    :   %0 = alloc_ref $Z
 // CHECK: global:   %1 = alloc_ref $Y