TempRValueElimination: support consuming indirect apply operands

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/TempRValueElimination.swift

@@ -177,6 +177,10 @@ private func hoistDestroy(of allocStack: AllocStackInst, after lastUse: Instruct
   case let li as LoadInst where li.loadOwnership == .take:
     assert(li.address == allocStack, "load must be not take a projected address")
     return
+  case let apply as ApplyInst where apply.consumes(address: allocStack):
+    if allocStack.uses.contains(where: { $0.instruction == apply && apply.convention(of: $0)!.isGuaranteed }) {
+      return
+    }
   default:
     break
   }
@@ -266,6 +270,14 @@ private extension AllocStackInst {
   }
 }
 
+private extension ApplySite {
+  func consumes(address: Value) -> Bool {
+    argumentOperands.contains { argOp in
+      argOp.value == address && convention(of: argOp)!.isConsumed
+    }
+  }
+}
+
 /// Tries to move an `end_access` down to extend the access scope over all uses of the `alloc_stack`.
 /// For example:
 /// ```
@@ -537,7 +549,12 @@ private struct UseCollector : AddressDefUseWalker {
   }
 
   private mutating func visitApply(address: Operand, apply: ApplySite) -> WalkResult {
-    if !apply.convention(of: address)!.isGuaranteed {
+    let argConvention = apply.convention(of: address)!
+    guard argConvention.isGuaranteed ||
+          // Only accept consuming-in arguments if it consumes the whole `alloc_stack`. A consume from
+          // a projection would destroy only a part of the `alloc_stack` and we don't handle this.
+          (argConvention == .indirectIn && (copy.isTakeOfSource && address.value == copy.destinationAddress))
+    else {
       return .abortWalk
     }
     uses.insert(apply)

test/Interop/Cxx/reference/rvalue-reference.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift(-I %S/Inputs/ -Xfrontend -enable-experimental-cxx-interop -Onone) | %FileCheck -check-prefix=CHECK-DASH-ONONE %s
+// RUN: %target-run-simple-swift(-I %S/Inputs/ -Xllvm -sil-disable-pass=mandatory-temp-rvalue-elimination -Xfrontend -enable-experimental-cxx-interop -Onone) | %FileCheck -check-prefix=CHECK-DASH-ONONE %s
 // RUN: %target-run-simple-swift(-I %S/Inputs/ -Xfrontend -enable-experimental-cxx-interop -O) | %FileCheck -check-prefix=CHECK-DASH-O %s
 //
 // REQUIRES: executable_test

test/SILOptimizer/temp_rvalue_opt_ossa.sil

@@ -65,6 +65,8 @@ sil [ossa] @guaranteed_user_with_result : $@convention(thin) (@guaranteed Klass)
 sil [ossa] @inguaranteed_user_without_result_NTS : $@convention(thin) (@in_guaranteed NonTrivialStruct) -> ()
 sil [ossa] @inguaranteed_user_without_result_MOS : $@convention(thin) (@in_guaranteed MOS) -> ()
 
+sil [ossa] @consuming_indirect_arg : $@convention(thin) (@in Klass) -> ()
+
 sil [ossa] @inguaranteed_user_without_result : $@convention(thin) (@in_guaranteed Klass) -> () {
 bb0(%0 : $*Klass):
   %9999 = tuple()
@@ -1643,3 +1645,33 @@ bb0(%0 : @owned $Klass):
   return %4
 }
 
+// CHECK-LABEL: sil [ossa] @take_and_consuming_arg :
+// CHECK-NOT:     alloc_stack
+// CHECK:         apply {{%[0-9]+}}(%0)
+// CHECK-LABEL: } // end sil function 'take_and_consuming_arg'
+sil [ossa] @take_and_consuming_arg : $@convention(thin) (@in Klass) -> () {
+bb0(%0 : $*Klass):
+  %1 = alloc_stack $Klass
+  copy_addr [take] %0 to [init] %1
+  %3 = function_ref @consuming_indirect_arg : $@convention(thin) (@in Klass) -> ()
+  %4 = apply %3(%1) : $@convention(thin) (@in Klass) -> ()
+  dealloc_stack %1 : $*Klass
+  %9 = tuple ()
+  return %9 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @copy_and_consuming_arg :
+// CHECK:         alloc_stack
+// CHECK-NEXT:    copy_addr
+// CHECK-LABEL: } // end sil function 'copy_and_consuming_arg'
+sil [ossa] @copy_and_consuming_arg : $@convention(thin) (@inout Klass) -> () {
+bb0(%0 : $*Klass):
+  %1 = alloc_stack $Klass
+  copy_addr %0 to [init] %1
+  %3 = function_ref @consuming_indirect_arg : $@convention(thin) (@in Klass) -> ()
+  %4 = apply %3(%1) : $@convention(thin) (@in Klass) -> ()
+  dealloc_stack %1 : $*Klass
+  %9 = tuple ()
+  return %9 : $()
+}
+