[sil] Do not eliminate (apply (partial_apply)) if the partial_apply captures an @in parameter and handle @in_guaranteed appropriately.

Otherwise the code as written miscompiles code like:

```
@inline(never)
func consumeSelf<T>(_ t : __owned T) {
  print("Consuming self!")
  print(t)
}

class Klass {}
struct S<T> {
  let t: T? = (Klass() as! T)

  @inline(__always)
  __consuming func foo(_ t: T) {
    consumeSelf(self)
  }
}

public func test<T>(_ t: __owned T) {
  let k = S<T>()
  let f = k.foo

  for _ in 0..<1024 {
    f(t)
  }
}

test(Klass())
```

As one can tell, without annotations it is hard to create an example like the
above, but there is no reason why we couldn't emit more code like this from the
frontend.

If the parameter is guaranteed though, the current impl is fine for
@in_guaranteed since in the loop, we do not need to retain or release the
value.

rdar://58885352

Author: gottesmm

lib/SILOptimizer/Utils/PartialApplyCombiner.cpp

@@ -90,29 +90,38 @@ bool PartialApplyCombiner::allocateTemporaries() {
   for (unsigned i : indices(argList)) {
     SILValue arg = argList[i];
     SILParameterInfo param = paramList[i];
+    // If we do not have an indirect parameter, we do not need to do any further
+    // work here.
+    if (!pai->getSubstCalleeConv().isSILIndirect(param))
+      continue;
+
+    // If our indirect parameter is mutating, we can just skip it as well.
     if (param.isIndirectMutating())
       continue;
 
-    // Create a temporary and copy the argument into it, if:
-    // - the argument stems from an alloc_stack
-    // - the argument is consumed by the callee and is indirect
-    //   (e.g. it is an @in argument)
-    if (isa<AllocStackInst>(arg) ||
-        (param.isConsumed() &&
-         pai->getSubstCalleeConv().isSILIndirect(param))) {
-      // If the argument has a dependent type, then we can not create a
-      // temporary for it at the beginning of the function, so we must bail.
-      //
-      // TODO: This is because we are inserting alloc_stack at the beginning/end
-      // of functions where the dependent type may not exist yet.
-      if (arg->getType().hasOpenedExistential())
-        return false;
-
-      // If the temporary is non-trivial, we need to destroy it later.
-      if (!arg->getType().isTrivial(*pai->getFunction()))
-        needsDestroys = true;
-      argsToHandle.push_back(std::make_pair(arg, i));
+    // If we are consuming an indirect parameter. Bail! We do not support that
+    // today!
+    if (param.isConsumed()) {
+      return false;
+    }
+
+    // Otherwise, we must have a guaranteed parameter. If we don't, bail.
+    if (!param.isGuaranteed()) {
+      return false;
     }
+
+    // If the argument has a dependent type, then we can not create a
+    // temporary for it at the beginning of the function, so we must bail.
+    //
+    // TODO: This is because we are inserting alloc_stack at the beginning/end
+    // of functions where the dependent type may not exist yet.
+    if (arg->getType().hasOpenedExistential())
+      return false;
+
+    // If the temporary is non-trivial, we need to destroy it later.
+    if (!arg->getType().isTrivial(*pai->getFunction()))
+      needsDestroys = true;
+    argsToHandle.push_back(std::make_pair(arg, i));
   }
 
   if (needsDestroys) {

test/Interpreter/apply_partial_apply_consuming_elim.swift

@@ -0,0 +1,30 @@
+// RUN: %target-run-simple-swift(-Osize)
+
+// REQUIRES: executable_test
+
+@inline(never)
+func consumeSelf<T>(_ t : __owned T) {
+  print("Consuming self!")
+  print(t)
+}
+
+class Klass {}
+struct S<T> {
+  let t: T? = (Klass() as! T)
+
+  @inline(__always)
+  __consuming func foo(_ t: T) {
+    consumeSelf(self)
+  }
+}
+
+public func test<T>(_ t: __owned T) {
+  let k = S<T>()
+  let f = k.foo
+
+  for _ in 0..<1024 {
+    f(t)
+  }
+}
+
+test(Klass())

test/SILOptimizer/sil_combine.sil

@@ -2905,231 +2905,6 @@ bb0:
   return %2 : $Builtin.Int1
 }
 
-class CC1 {
-   deinit
-  init()
-}
-
-class CC2 {
-   deinit
-  init()
-}
-
-class CC3 {
-   deinit
-  init()
-}
-
-class CC4 {
-   deinit
-  init()
-}
-
-
-// Function that takes different kinds of arguments: @in, @guaranteed, @owned,
-sil @closure_with_in_guaranteed_owned_in_args : $@convention(method) (@in CC2, @guaranteed CC1, @owned CC3, @in CC4) -> Optional<Int>
-
-// Test the peephole performing apply{partial_apply(x, y, z)}(a) -> apply(a, x, y, z)
-// We need to check the following:
-// - all arguments of a partial_apply, which are either results of a stack_alloc or consumed indirect arguments
-//  should be copied into temporaries. This should happen just before that partial_apply instruction.
-// - The temporaries are allocated at the beginning of the function and deallocated at the end.
-// - Before each apply of the partial_apply, we retain values of any arguments which are of non-address type.
-//   This is required because they could be consumed (i.e. released by the callee).
-// - After each apply of the partial_apply, we release values of any arguments which are non-consumed by the callee (e.g. @guaranteed ones)
-
-// CHECK-LABEL: sil @test_apply_of_partial_apply
-
-// CHECK: bb0{{.*}}:
-// A new temporary should have been created for each alloc_stack argument passed to partial_apply
-// CHECK: [[TMP:%[0-9]+]] = alloc_stack $CC4
-// CHECK: [[CLOSURE:%[0-9]+]] = function_ref @closure_with_in_guaranteed_owned_in_args
-// Copy the original value of the argument into a temporary
-// CHECK: copy_addr {{.*}} to [initialization] [[TMP]] : $*CC4
-// CHECK-NOT: partial_apply
-
-// CHECK: bb1:
-// CHECK-NOT: partial_apply
-// Check that the peephole inserted a retain of the closure's guaranteed argument
-// CHECK: strong_retain %{{[0-9]+}} : $CC1
-// Check that the peephole inserted a retain of the closure's owned argument
-// CHECK: strong_retain %{{[0-9]+}} : $CC3
-// CHECK: apply [[CLOSURE]]
-// Check that the peephole inserted a release the closure's guaranteed argument
-// CHECK: strong_release %{{[0-9]+}} : $CC1
-// CHECK-NOT: partial_apply
-// Retain the guaranteed argument
-// CHECK: strong_retain %{{[0-9]+}} : $CC1
-// Retain the owned argument
-// CHECK: strong_retain %{{[0-9]+}} : $CC3
-// CHECK: apply [[CLOSURE]]({{.*}}, [[TMP]])
-// Check that the peephole inserted a release the guaranteed argument
-// CHECK: strong_release %{{[0-9]+}} : $CC1
-// Release the @owned CC4 argument of the function
-// CHECK: load {{%[0-9]+}} : $*CC4
-// CHECK: strong_release {{%[0-9]+}} : $CC4
-// CHECK: br bb3
-
-// CHECK: bb2:
-// CHECK-NOT: partial_apply
-// Check that the peephole inserted a retain of the closure's guaranteed argument
-// CHECK: strong_retain %{{[0-9]+}} : $CC1
-// Check that the peephole inserted a retain of the closure's owned argument
-// CHECK: strong_retain %{{[0-9]+}} : $CC3
-// CHECK: apply [[CLOSURE]]({{.*}}, [[TMP]])
-// Check that the peephole inserted a release the closure's guaranteed argument
-// CHECK: strong_release %{{[0-9]+}} : $CC1
-// Release the @owned CC4 argument of the function
-// CHECK: load {{%[0-9]+}} : $*CC4
-// CHECK: strong_release {{%[0-9]+}} : $CC4
-// CHECK: br bb3
-
-// bb3:
-// dealloc_stack [[TMP]] : $CC4
-
-sil @test_apply_of_partial_apply : $@convention(thin) (@in Optional<CC2>, @guaranteed CC1, @guaranteed CC3, @guaranteed CC4, @guaranteed CC2) -> Optional<Int> {
-bb0(%0 : $*Optional<CC2>, %1 : $CC1, %2 : $CC3, %3 : $CC4, %4 : $CC2):
-
-  %5 = function_ref @closure_with_in_guaranteed_owned_in_args : $@convention(method) (@in CC2, @guaranteed CC1, @owned CC3, @in CC4) -> Optional<Int>
-  %6 = alloc_stack $CC3
-  store %2 to %6 : $*CC3
-  %8 = load %6 : $*CC3
-  %9 = alloc_stack $CC4
-  store %3 to %9 : $*CC4
-  %11 = load %9 : $*CC4
-  strong_retain %1 : $CC1
-  %12 = partial_apply %5(%1, %8, %9) : $@convention(method) (@in CC2, @guaranteed CC1, @owned CC3, @in CC4) -> Optional<Int>
-  dealloc_stack %9 : $*CC4
-  dealloc_stack %6 : $*CC3
-  %15 = convert_function %12 : $@callee_owned (@in CC2) -> Optional<Int> to $@callee_owned (@in CC2) -> (Optional<Int>, @error Error)
-  %16 = alloc_stack $Optional<Int>
-  %17 = alloc_stack $Optional<CC2>
-  copy_addr %0 to [initialization] %17 : $*Optional<CC2>
-  switch_enum_addr %17 : $*Optional<CC2>, case #Optional.some!enumelt.1: bb1, case #Optional.none!enumelt: bb2
-
-bb1:
-  %21 = unchecked_take_enum_data_addr %17 : $*Optional<CC2>, #Optional.some!enumelt.1
-
-  %22 = alloc_stack $CC2
-  copy_addr [take] %21 to [initialization] %22 : $*CC2
-  %24 = alloc_stack $CC2
-  copy_addr %22 to [initialization] %24 : $*CC2
-  destroy_addr %22 : $*CC2
-  strong_retain %15 : $@callee_owned (@in CC2) -> (Optional<Int>, @error Error)
-  %28 = apply %12(%24) : $@callee_owned (@in CC2) -> Optional<Int>
-  store %28 to %16 : $*Optional<Int>
-  dealloc_stack %24 : $*CC2
-  dealloc_stack %22 : $*CC2
-
-  %102 = alloc_stack $CC2
-  copy_addr [take] %21 to [initialization] %102 : $*CC2
-  %104 = alloc_stack $CC2
-  copy_addr %102 to [initialization] %104 : $*CC2
-  destroy_addr %102 : $*CC2
-  strong_retain %15 : $@callee_owned (@in CC2) -> (Optional<Int>, @error Error)
-  %108 = apply %12(%104) : $@callee_owned (@in CC2) -> Optional<Int>
-  store %108 to %16 : $*Optional<Int>
-  dealloc_stack %104 : $*CC2
-  dealloc_stack %102 : $*CC2
-
-  dealloc_stack %17 : $*Optional<CC2>
-
-  br bb3
-
-bb2:
-  %39 = alloc_stack $CC2
-  store %4 to %39 : $*CC2
-  strong_retain %15 : $@callee_owned (@in CC2) -> (Optional<Int>, @error Error)
-  %42 = apply %12(%39) : $@callee_owned (@in CC2) -> Optional<Int>
-  store %42 to %16 : $*Optional<Int>
-  dealloc_stack %39 : $*CC2
-  dealloc_stack %17 : $*Optional<CC2>
-  br bb3
-
-bb3:
-  destroy_addr %0 : $*Optional<CC2>
-  strong_release %15 : $@callee_owned (@in CC2) -> (Optional<Int>, @error Error)
-  %36 = load %16 : $*Optional<Int>
-  dealloc_stack %16 : $*Optional<Int>
-  return %36 : $Optional<Int>
-}
-
-// Test if we insert the right stack/dealloc-stack when converting apply{partial_apply}.
-
-// CHECK-LABEL: sil @test_stack_insertion_for_partial_apply_apply
-// CHECK:     bb0({{.*}}):
-// CHECK-NEXT:  [[T:%[0-9]+]] = alloc_stack $Int
-// CHECK-NEXT:  alloc_stack $Bool
-// CHECK-NEXT:  [[S2:%[0-9]+]] = alloc_stack $Int
-// CHECK:       copy_addr [[S2]] to [initialization] [[T]]
-// CHECK:       dealloc_stack [[S2]] : $*Int
-// CHECK:       apply
-// CHECK:     bb1:
-// CHECK-NOT:   dealloc_stack
-// CHECK:     bb2:
-// CHECK:       dealloc_stack {{.*}} : $*Bool
-// CHECK:       dealloc_stack [[T]] : $*Int
-// CHECK:       return
-sil @test_stack_insertion_for_partial_apply_apply : $@convention(thin) (Int, Double) -> () {
-bb0(%0 : $Int, %1 : $Double):
-  %s1 = alloc_stack $Bool
-  %s2 = alloc_stack $Int
-  store %0 to %s2 : $*Int
-  %f1 = function_ref @callee : $@convention(thin) (Double, @in_guaranteed Int) -> ()
-  %pa = partial_apply %f1(%s2) : $@convention(thin) (Double, @in_guaranteed Int) -> ()
-  dealloc_stack %s2 : $*Int
-  %a1 = apply %pa(%1) : $@callee_owned (Double) -> ()
-  cond_br undef, bb1, bb2
-
-bb1:
-  %f2 = function_ref @noreturn_func : $@convention(thin) () -> Never
-  %a2 = apply %f2() : $@convention(thin) () -> Never
-  unreachable
-
-bb2:
-  dealloc_stack %s1 : $*Bool
-  %r = tuple ()
-  return %r : $()
-}
-
-// CHECK-LABEL: sil @test_generic_partial_apply_apply
-// CHECK: bb0([[ARG0:%.*]] : $*T, [[ARG1:%.*]] : $*T):
-// CHECK-NEXT: [[TMP:%.*]] = alloc_stack $T
-// CHECK: [[FN:%.*]] = function_ref @generic_callee
-// CHECK-NEXT: copy_addr [[ARG1]] to [initialization] [[TMP]] : $*T
-// CHECK-NEXT: destroy_addr [[ARG1]]
-// CHECK-NEXT: apply [[FN]]<T, T>([[ARG0]], [[TMP]])
-// CHECK-NEXT: destroy_addr [[TMP]]
-// CHECK-NEXT: tuple
-// CHECK-NEXT: dealloc_stack [[TMP]]
-// CHECK-NEXT: return
-sil @test_generic_partial_apply_apply : $@convention(thin) <T> (@in T, @in T) -> () {
-bb0(%0 : $*T, %1 : $*T):
-  %f1 = function_ref @generic_callee : $@convention(thin) <T, U> (@in T, @in U) -> ()
-  %pa = partial_apply %f1<T, T>(%1) : $@convention(thin) <T, U> (@in T, @in U) -> ()
-  %a1 = apply %pa(%0) : $@callee_owned (@in T) -> ()
-  %r = tuple ()
-  return %r : $()
-}
-
-// CHECK-LABEL: sil @test_existential_partial_apply_apply
-// CHECK: bb0(%0 : $*FakeProtocol):
-// CHECK-NEXT: [[OPEN:%.*]] = open_existential_addr immutable_access
-// CHECK-NEXT: [[FN:%.*]] = witness_method
-// CHECK-NEXT: apply [[FN]]<@opened("5DD6F3D0-808A-11E6-93A0-34363BD08DA0") FakeProtocol>([[OPEN]])
-// CHECK-NEXT: tuple
-// CHECK-NEXT: return
-sil @test_existential_partial_apply_apply : $@convention(thin) (@in FakeProtocol) -> () {
-bb0(%0: $*FakeProtocol):
-  %o = open_existential_addr immutable_access %0 : $*FakeProtocol to $*@opened("5DD6F3D0-808A-11E6-93A0-34363BD08DA0") FakeProtocol
-  %f1 = witness_method $@opened("5DD6F3D0-808A-11E6-93A0-34363BD08DA0") FakeProtocol, #FakeProtocol.requirement!1, %o : $*@opened("5DD6F3D0-808A-11E6-93A0-34363BD08DA0") FakeProtocol : $@convention(witness_method: FakeProtocol) <T where T : FakeProtocol> (@in_guaranteed T) -> ()
-  %pa = partial_apply %f1<@opened("5DD6F3D0-808A-11E6-93A0-34363BD08DA0") FakeProtocol>() : $@convention(witness_method: FakeProtocol) <T where T : FakeProtocol> (@in_guaranteed T) -> ()
-  %a1 = apply %pa(%o) : $@callee_owned (@in_guaranteed @opened("5DD6F3D0-808A-11E6-93A0-34363BD08DA0") FakeProtocol) -> ()
-
-  %r = tuple ()
-  return %r : $()
-}
-
 sil @callee : $@convention(thin) (Double, @in_guaranteed Int) -> ()
 sil @generic_callee : $@convention(thin) <T, U> (@in T, @in U) -> ()
 sil @noreturn_func : $@convention(thin) () -> Never