Fix another use-after-free in SILCombine (#34168)

* Fix another use-after-free in SILCombine

swift::endLifetimeAtFrontier also needs to use
swift::emitDestroyOperation and delete instructions via callbacks that
can correctly remove it from the worklist that SILCombine maintains

* Add test for use-after-free in SILCombine

Author: meg-gupta

include/swift/SILOptimizer/Utils/ValueLifetime.h

@@ -17,8 +17,9 @@
 #ifndef SWIFT_SILOPTIMIZER_UTILS_CFG_H
 #define SWIFT_SILOPTIMIZER_UTILS_CFG_H
 
-#include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILBuilder.h"
+#include "swift/SIL/SILInstruction.h"
+#include "swift/SILOptimizer/Utils/InstOptUtils.h"
 
 namespace swift {
 
@@ -159,7 +160,8 @@ class ValueLifetimeAnalysis {
 /// destroy_value at each instruction of the \p frontier.
 void endLifetimeAtFrontier(SILValue valueOrStackLoc,
                            const ValueLifetimeAnalysis::Frontier &frontier,
-                           SILBuilderContext &builderCtxt);
+                           SILBuilderContext &builderCtxt,
+                           InstModCallbacks callbacks);
 
 } // end namespace swift
 

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -1323,7 +1323,8 @@ bool swift::collectDestroys(SingleValueInstruction *inst,
 ///       not be needed anymore with OSSA.
 static bool keepArgsOfPartialApplyAlive(PartialApplyInst *pai,
                                         ArrayRef<SILInstruction *> paiUsers,
-                                        SILBuilderContext &builderCtxt) {
+                                        SILBuilderContext &builderCtxt,
+                                        InstModCallbacks callbacks) {
   SmallVector<Operand *, 8> argsToHandle;
   getConsumedPartialApplyArgs(pai, argsToHandle,
                               /*includeTrivialAddrArgs*/ false);
@@ -1357,7 +1358,7 @@ static bool keepArgsOfPartialApplyAlive(PartialApplyInst *pai,
 
     // Delay the destroy of the value (either as SSA value or in the stack-
     // allocated temporary) at the end of the partial_apply's lifetime.
-    endLifetimeAtFrontier(tmp, partialApplyFrontier, builderCtxt);
+    endLifetimeAtFrontier(tmp, partialApplyFrontier, builderCtxt, callbacks);
   }
   return true;
 }
@@ -1406,7 +1407,8 @@ bool swift::tryDeleteDeadClosure(SingleValueInstruction *closure,
            "partial_apply [stack] should have been handled before");
     SILBuilderContext builderCtxt(pai->getModule());
     if (needKeepArgsAlive) {
-      if (!keepArgsOfPartialApplyAlive(pai, closureDestroys, builderCtxt))
+      if (!keepArgsOfPartialApplyAlive(pai, closureDestroys, builderCtxt,
+                                       callbacks))
         return false;
     } else {
       // A preceeding partial_apply -> apply conversion (done in

lib/SILOptimizer/Utils/PartialApplyCombiner.cpp

@@ -115,7 +115,7 @@ bool PartialApplyCombiner::copyArgsToTemporaries(
 
     // Destroy the argument value (either as SSA value or in the stack-
     // allocated temporary) at the end of the partial_apply's lifetime.
-    endLifetimeAtFrontier(tmp, partialApplyFrontier, builderCtxt);
+    endLifetimeAtFrontier(tmp, partialApplyFrontier, builderCtxt, callbacks);
   }
   return true;
 }

lib/SILOptimizer/Utils/ValueLifetime.cpp

@@ -327,15 +327,12 @@ void ValueLifetimeAnalysis::dump() const {
 
 void swift::endLifetimeAtFrontier(
     SILValue valueOrStackLoc, const ValueLifetimeAnalysis::Frontier &frontier,
-    SILBuilderContext &builderCtxt) {
+    SILBuilderContext &builderCtxt, InstModCallbacks callbacks) {
   for (SILInstruction *endPoint : frontier) {
     SILBuilderWithScope builder(endPoint, builderCtxt);
     SILLocation loc = RegularLocation(endPoint->getLoc().getSourceLoc());
-    if (valueOrStackLoc->getType().isObject()) {
-      builder.emitDestroyValueOperation(loc, valueOrStackLoc);
-    } else {
-      assert(isa<AllocStackInst>(valueOrStackLoc));
-      builder.createDestroyAddr(loc, valueOrStackLoc);
+    emitDestroyOperation(builder, loc, valueOrStackLoc, callbacks);
+    if (isa<AllocStackInst>(valueOrStackLoc)) {
       builder.createDeallocStack(loc, valueOrStackLoc);
     }
   }

test/SILOptimizer/sil_combine_apply.sil

@@ -991,3 +991,29 @@ sil @test_partial_apply_method : $@convention(thin) () -> @owned @callee_owned (
 // CHECK: [[FN:%.*]] = function_ref @method
 // CHECK-NEXT: partial_apply [[FN]]()
 // CHECK-NEXT: return
+
+sil [noinline] @$foo : $@convention(thin) (@owned { var Int64 }) -> () 
+
+// CHECK-LABEL: sil private [noinline] @$testdeadpartialapply :
+// CHECK-NOT: partial_apply
+// CHECK-LABEL: } // end sil function '$testdeadpartialapply'
+sil private [noinline] @$testdeadpartialapply : $@convention(thin) (@owned { var Int64 }) -> () {
+bb0(%0 : ${ var Int64 }):
+  %3 = function_ref @$foo : $@convention(thin) (@owned { var Int64 }) -> ()
+  %5 = partial_apply [callee_guaranteed] %3(%0) : $@convention(thin) (@owned { var Int64 }) -> ()
+  cond_br undef, bb1, bb2
+bb1:
+  strong_retain %0 : ${ var Int64 }
+  strong_retain %5 : $@callee_guaranteed () -> ()
+  unreachable
+bb2:
+  strong_retain %0 : ${ var Int64 }
+  strong_retain %5 : $@callee_guaranteed () -> ()
+  br bb3
+bb3:
+  strong_release %5 : $@callee_guaranteed () -> ()
+  strong_release %5 : $@callee_guaranteed () -> ()
+  strong_release %0 : ${ var Int64 }
+  %rv = tuple()
+  return %rv : $()
+}