Merge pull request swiftlang#67947 from nate-chandler/opaque-values/20230815/1/partial-apply-indexing

[AddressLowering] Fix partial_apply argument indexing.

Author: nate-chandler

include/swift/SIL/ApplySite.h

@@ -295,9 +295,20 @@ class ApplySite {
   /// Return the apply operand for the given applied argument index.
   Operand &getArgumentRef(unsigned i) const { return getArgumentOperands()[i]; }
 
+  // The apply operand at the given index into the callee's function's
+  // arguments.
+  Operand &getArgumentRefAtCalleeArgIndex(unsigned i) const {
+    return getArgumentRef(i - getCalleeArgIndexOfFirstAppliedArg());
+  }
+
   /// Return the ith applied argument.
   SILValue getArgument(unsigned i) const { return getArguments()[i]; }
 
+  // The argument at the given index into the callee's function's arguments.
+  SILValue getArgumentAtCalleeArgIndex(unsigned i) const {
+    return getArgument(i - getCalleeArgIndexOfFirstAppliedArg());
+  }
+
   /// Set the ith applied argument.
   void setArgument(unsigned i, SILValue V) const {
     getArgumentOperands()[i].set(V);

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -2190,14 +2190,17 @@ bool CallArgRewriter::rewriteArguments() {
   bool changed = false;
 
   auto origConv = apply.getSubstCalleeConv();
-  assert(apply.getNumArguments() == origConv.getNumParameters() &&
-         "results should not yet be rewritten");
+  assert((apply.getNumArguments() == origConv.getNumParameters() &&
+          apply.asFullApplySite()) ||
+         (apply.getNumArguments() <= origConv.getNumParameters() &&
+          !apply.asFullApplySite()) &&
+             "results should not yet be rewritten");
 
   for (unsigned argIdx = apply.getCalleeArgIndexOfFirstAppliedArg(),
                 endArgIdx = argIdx + apply.getNumArguments();
        argIdx < endArgIdx; ++argIdx) {
 
-    Operand &operand = apply.getArgumentRef(argIdx);
+    Operand &operand = apply.getArgumentRefAtCalleeArgIndex(argIdx);
     // Ignore arguments that have already been rewritten with an address.
     if (operand.get()->getType().isAddress())
       continue;

test/SILOptimizer/address_lowering.sil

@@ -2619,6 +2619,22 @@ sil [ossa] @test_partial_apply_4_loadable_stack : $() -> () {
   return %retval : $()
 }
 
+// CHECK-LABEL: sil [ossa] @test_partial_apply_5_indexing : {{.*}} {
+// CHECK:       {{bb[0-9]+}}([[C:%[^,]+]] :
+// CHECK:         [[CALLEE:%[^,]+]] = function_ref @test_partial_apply_5_indexing_callee
+// CHECK:         [[CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [[CALLEE]]([[C]])
+// CHECK:         destroy_value [[CLOSURE]]
+// CHECK-LABEL: } // end sil function 'test_partial_apply_5_indexing'
+sil @test_partial_apply_5_indexing_callee : $@convention(thin) (@inout Int, @guaranteed Klass) -> (@out Int)
+sil [ossa] @test_partial_apply_5_indexing : $@convention(thin) (@owned Klass) -> () {
+bb0(%c : @owned $Klass):
+  %callee = function_ref @test_partial_apply_5_indexing_callee : $@convention(thin) (@inout Int, @guaranteed Klass) -> (@out Int)
+  %closure = partial_apply [callee_guaranteed] %callee(%c) : $@convention(thin) (@inout Int, @guaranteed Klass) -> (@out Int)
+  destroy_value %closure : $@callee_guaranteed (@inout Int) -> (@out Int)
+  %retval = tuple ()
+  return %retval : $()
+}
+
 // CHECK-LABEL: sil hidden [ossa] @test_store_1 : {{.*}} {
 // CHECK:         [[MAYBE_ADDR:%[^,]+]] = alloc_stack $Optional<Self>
 // CHECK:         [[LOAD_ADDR:%[^,]+]] = alloc_stack $Self