SILCombiner: Fix partial_apply optimization of @callee_guaranteed

closures

@callee_guaranteed closure contexts must not be released when replacing
a closure application.

SR-5441
rdar://33255593

Author: aschwaighofer

include/swift/SIL/SILInstruction.h

@@ -1914,6 +1914,9 @@ class PartialApplyInst
   CanSILFunctionType getFunctionType() const {
     return getType().castTo<SILFunctionType>();
   }
+  bool hasCalleeGuaranteedContext() const {
+    return getType().castTo<SILFunctionType>()->isCalleeGuaranteed();
+  }
 };
 
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -334,14 +334,18 @@ bool PartialApplyCombiner::processSingleApply(FullApplySite AI) {
     for (auto Arg : ToBeReleasedArgs) {
       Builder.emitDestroyValueOperation(PAI->getLoc(), Arg);
     }
-    Builder.createStrongRelease(AI.getLoc(), PAI, Builder.getDefaultAtomicity());
+    if (!PAI->hasCalleeGuaranteedContext())
+      Builder.createStrongRelease(AI.getLoc(), PAI,
+                                  Builder.getDefaultAtomicity());
     Builder.setInsertionPoint(AI.getInstruction());
   } else {
     // Release the non-consumed parameters.
     for (auto Arg : ToBeReleasedArgs) {
       Builder.emitDestroyValueOperation(PAI->getLoc(), Arg);
     }
-    Builder.createStrongRelease(AI.getLoc(), PAI, Builder.getDefaultAtomicity());
+    if (!PAI->hasCalleeGuaranteedContext())
+      Builder.createStrongRelease(AI.getLoc(), PAI,
+                                  Builder.getDefaultAtomicity());
   }
 
   if (auto apply = dyn_cast<ApplyInst>(AI))

test/SILOptimizer/sil_combine_apply.sil

@@ -437,3 +437,29 @@ bb99:
   %empty = tuple ()
   return %empty : $()
 }
+
+class C {}
+
+sil @guaranteed_closure : $@convention(thin) (@guaranteed C) -> ()
+
+// CHECK-LABEL: sil @test_guaranteed_closure
+// CHECK: bb0([[ARG:%.*]] : $C):
+// CHECK: strong_retain [[ARG]]
+// CHECK: [[F:%.*]] = function_ref @guaranteed_closure
+// CHECK: [[C:%.*]] = partial_apply [callee_guaranteed] [[F]]([[ARG]])
+// CHECK: strong_retain [[ARG]]
+// CHECK: apply [[F]]
+// Don't release the closure context -- it is @callee_guaranteed.
+// CHECK-NOT: strong_release [[C]] : $@callee_guaranteed () -> ()
+// CHECK: strong_release [[ARG]]
+// CHECK-NOT: strong_release [[C]] : $@callee_guaranteed () -> ()
+// CHECK: return [[C]]
+
+sil @test_guaranteed_closure : $@convention(thin) (@guaranteed C) -> @owned @callee_guaranteed () -> () {
+bb0(%0: $C):
+  strong_retain %0 : $C
+  %closure_fun = function_ref @guaranteed_closure : $@convention(thin) (@guaranteed C) -> ()
+  %closure = partial_apply [callee_guaranteed] %closure_fun(%0) : $@convention(thin) (@guaranteed C) -> ()
+  apply %closure() : $@callee_guaranteed () -> ()
+  return %closure : $@callee_guaranteed () -> ()
+}