Fix SILCombiner for partial_apply [stack]

Author: aschwaighofer

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -74,9 +74,22 @@ static bool foldInverseReabstractionThunks(PartialApplyInst *PAI,
 SILInstruction *SILCombiner::visitPartialApplyInst(PartialApplyInst *PAI) {
   // partial_apply without any substitutions or arguments is just a
   // thin_to_thick_function.
-  if (!PAI->hasSubstitutions() && (PAI->getNumArguments() == 0))
-    return Builder.createThinToThickFunction(PAI->getLoc(), PAI->getCallee(),
-                                             PAI->getType());
+  if (!PAI->hasSubstitutions() && (PAI->getNumArguments() == 0)) {
+    if (!PAI->isOnStack())
+      return Builder.createThinToThickFunction(PAI->getLoc(), PAI->getCallee(),
+                                               PAI->getType());
+
+    // Remove dealloc_stack of partial_apply [stack].
+    for (auto *Use : PAI->getUses())
+      if (auto *dealloc = dyn_cast<DeallocStackInst>(Use->getUser()))
+        eraseInstFromFunction(*dealloc);
+    auto *thinToThick = Builder.createThinToThickFunction(
+        PAI->getLoc(), PAI->getCallee(), PAI->getType());
+    replaceInstUsesWith(*PAI, thinToThick);
+    eraseInstFromFunction(*PAI);
+    return nullptr;
+  }
+
 
   // partial_apply %reabstraction_thunk_typeAtoB(
   //    partial_apply %reabstraction_thunk_typeBtoA %closure_typeB))
@@ -138,6 +151,12 @@ class PartialApplyCombiner {
 
 /// Returns true on success.
 bool PartialApplyCombiner::allocateTemporaries() {
+  // A partial_apply [stack]'s argument are not owned by the partial_apply and
+  // therefore their lifetime must outlive any uses.
+  if (PAI->isOnStack()) {
+    return true;
+  }
+
   // Copy the original arguments of the partial_apply into newly created
   // temporaries and use these temporaries instead of the original arguments
   // afterwards.
@@ -397,6 +416,15 @@ SILInstruction *PartialApplyCombiner::combine() {
       Uses.append(CFI->getUses().begin(), CFI->getUses().end());
       continue;
     }
+    // Look through mark_dependence users of partial_apply [stack].
+    if (auto *MD = dyn_cast<MarkDependenceInst>(User)) {
+      if (MD->getValue() == Use->get() &&
+          MD->getValue()->getType().is<SILFunctionType>() &&
+          MD->getValue()->getType().castTo<SILFunctionType>()->isNoEscape()) {
+        Uses.append(MD->getUses().begin(), MD->getUses().end());
+      }
+      continue;
+    }
     // If this use of a partial_apply is not
     // an apply which uses it as a callee, bail.
     auto AI = FullApplySite::isa(User);

test/SILOptimizer/sil_combine_apply.sil

@@ -571,3 +571,68 @@ bb0:
   %999 = tuple ()
   return %999 : $()
 }
+
+// CHECK-LABEL: sil shared @applied_on_stack : $@convention(thin) () -> () {
+// CHECK: bb0:
+// CHECK:   [[F:%.*]] = function_ref @indirectClosure : $@convention(thin) (@in Builtin.Int64) -> @out Builtin.Int64
+// CHECK:   [[STK:%.*]] = alloc_stack $Builtin.Int64
+// CHECK:   [[STK2:%.*]] = alloc_stack $Builtin.Int64
+// CHECK:   [[I:%.*]] = integer_literal $Builtin.Int64, 3
+// CHECK:   store [[I]] to [[STK2]] : $*Builtin.Int64
+// CHECK:   apply [[F]]([[STK]], [[STK2]]) : $@convention(thin) (@in Builtin.Int64) -> @out Builtin.Int64
+// CHECK:   dealloc_stack [[STK2]] : $*Builtin.Int64
+// CHECK:   dealloc_stack [[STK]] : $*Builtin.Int64
+// CHECK:   return %8 : $()
+
+sil shared @applied_on_stack : $@convention(thin) () -> () {
+bb0:
+  %fn = function_ref @indirectClosure : $@convention(thin) (@in Builtin.Int64) -> @out Builtin.Int64
+  %pa = partial_apply [callee_guaranteed] [on_stack] %fn() : $@convention(thin) (@in Builtin.Int64) -> @out Builtin.Int64
+  %out = alloc_stack $Builtin.Int64
+  %in = alloc_stack $Builtin.Int64
+  %c3 = integer_literal $Builtin.Int64, 3
+  store %c3 to %in : $*Builtin.Int64
+  %call = apply %pa(%out, %in) : $@noescape @callee_guaranteed (@in Builtin.Int64) -> @out Builtin.Int64
+  dealloc_stack %in : $*Builtin.Int64
+  dealloc_stack %out : $*Builtin.Int64
+  dealloc_stack %pa : $@noescape @callee_guaranteed (@in Builtin.Int64) -> @out Builtin.Int64
+  %999 = tuple ()
+  return %999 : $()
+}
+
+sil @guaranteed_closure_throws2 : $@convention(thin) (Builtin.Int32, @guaranteed C) -> @error Error
+
+// CHECK-LABEL: sil @test_guaranteed_closure_try_apply_on_stack
+// CHECK: bb0
+// CHECK: strong_retain
+// CHECK: strong_retain
+// CHECK: try_apply
+// CHECK: bb1
+// CHECK: strong_release
+// CHECK: strong_release
+// CHECK: br bb3
+// CHECK: bb2
+// CHECK: strong_release
+// CHECK: strong_release
+// CHECK: br bb3
+sil @test_guaranteed_closure_try_apply_on_stack : $@convention(thin) (@guaranteed C, Builtin.Int32) ->  () {
+bb0(%0: $C, %1: $Builtin.Int32):
+  strong_retain %0 : $C
+  %closure_fun = function_ref @guaranteed_closure_throws2 : $@convention(thin) (Builtin.Int32, @guaranteed C) -> @error Error
+  %closure = partial_apply [callee_guaranteed] [on_stack] %closure_fun(%0) : $@convention(thin) (Builtin.Int32, @guaranteed C) -> @error Error
+  try_apply %closure(%1) : $@noescape @callee_guaranteed (Builtin.Int32) -> @error Error, normal bb7, error bb11
+
+bb7(%callret : $()):
+	dealloc_stack %closure : $@noescape @callee_guaranteed (Builtin.Int32) -> @error Error
+	strong_release %0: $C
+  br bb99
+
+bb11(%128 : $Error):
+	dealloc_stack %closure : $@noescape @callee_guaranteed (Builtin.Int32) -> @error Error
+	strong_release %0: $C
+  br bb99
+
+bb99:
+  %t = tuple()
+  return %t : $()
+}