Merge pull request #29645 from gottesmm/pr-4435ac48330e0dd08b30bf6fe8f29a93f25ff835

gottesmm · web-flow · commit 37655db07af3 · 2020-02-04T17:34:35.000-08:00
[sil] Do not eliminate (apply (partial_apply)) if the partial_apply captures an @in parameter and handle @in_guaranteed appropriately.
diff --git a/lib/SILOptimizer/Utils/PartialApplyCombiner.cpp b/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) {
diff --git a/test/Interpreter/apply_partial_apply_consuming_elim.swift b/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())
diff --git a/test/SILOptimizer/sil_combine.sil b/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
diff --git a/test/SILOptimizer/sil_combine_apply.sil b/test/SILOptimizer/sil_combine_apply.sil
