Merge pull request #64013 from nate-chandler/mandatory-generic-specializer/partial-apply-stack-discipline

[MandatoryGenericSpecializer] Fix stack nesting when lowering OSSA during inlining.

Author: nate-chandler

lib/SILOptimizer/Transforms/GenericSpecializer.cpp

@@ -353,6 +353,8 @@ optimizeInst(SILInstruction *inst, SILOptFunctionBuilder &funcBuilder,
     // If the de-virtualized callee is a transparent function, inline it.
     SILInliner::inlineFullApply(fas, SILInliner::InlineKind::MandatoryInline,
                                 funcBuilder, deleter);
+    if (callee->hasOwnership() && !inst->getFunction()->hasOwnership())
+      invalidatedStackNesting = true;
     return true;
   }
   if (auto *bi = dyn_cast<BuiltinInst>(inst)) {

test/SILOptimizer/mandatory_generic_specialization.sil

@@ -0,0 +1,37 @@
+// RUN: %target-sil-opt -enable-sil-verify-all %s -mandatory-generic-specializer | %FileCheck %s
+
+import Builtin
+
+sil @paable : $@convention(thin) (Builtin.Int64) -> ()
+
+sil [ossa] [transparent] @partial_apply_on_stack_nesting_violator : $@convention(thin) <T> () -> () {
+    %paable = function_ref @paable : $@convention(thin) (Builtin.Int64) -> ()
+    %one = integer_literal $Builtin.Int64, 1
+    %first = partial_apply [callee_guaranteed] [on_stack] %paable(%one) : $@convention(thin) (Builtin.Int64) -> ()
+    %two = integer_literal $Builtin.Int64, 2
+    %second = partial_apply [callee_guaranteed] [on_stack] %paable(%two) : $@convention(thin) (Builtin.Int64) -> ()
+    // Note that the destroy_values do not occur in an order which coincides
+    // with stack disciplined dealloc_stacks.
+    destroy_value %first : $@noescape @callee_guaranteed () -> ()
+    destroy_value %second : $@noescape @callee_guaranteed () -> ()
+    %retval = tuple ()
+    return %retval : $()
+}
+
+// Verify that when inlining partial_apply_on_stack_nesting_violator, the stack
+// nesting of the on_stack closures is fixed.
+// CHECK-LABEL: sil [no_locks] @test_inline_stack_violating_ossa_func : {{.*}} {
+// CHECK:         [[PAABLE:%[^,]+]] = function_ref @paable
+// CHECK:         [[FIRST:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[PAABLE]]
+// CHECK:         [[SECOND:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[PAABLE]]
+// CHECK:         dealloc_stack [[SECOND]]
+// CHECK:         dealloc_stack [[FIRST]]
+// CHECK-LABEL: } // end sil function 'test_inline_stack_violating_ossa_func'
+sil [no_locks] @test_inline_stack_violating_ossa_func : $@convention(thin) () -> () {
+    %callee = function_ref @partial_apply_on_stack_nesting_violator : $@convention(thin) <T> () -> ()
+    apply %callee<Builtin.Int64>() : $@convention(thin) <T> () -> ()
+    %retval = tuple ()
+    return %retval : $()
+}
+
+