Fix use-after-free in SILCombine (#34145)

SILCombine maintains a worklist of instructions and deleting of instructions is valid only via callbacks that remove them from the worklist as well. It calls swift::tryDeleteDeadClosure which in turn calls SILBuilder apis like emitStrongRelease/emitReleaseValue/emitDestroyValue which can delete instructions via SILInstruction::eraseFromParent leaving behind a stale entry in SILCombine's worklist causing a crash.

This PR adds swift::emitDestroyOperation which correctly calls the appropriate InstModCallbacks on added/removed instructions. This comes from swift::releasePartialApplyCapturedArg which was handling creation of destroys with callbacks correctly.

Author: meg-gupta

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -357,6 +357,11 @@ void getConsumedPartialApplyArgs(PartialApplyInst *pai,
                                  SmallVectorImpl<Operand *> &argOperands,
                                  bool includeTrivialAddrArgs);
 
+/// Emit destroy operation for \p operand, and call appropriate functions from
+/// \p callbacks for newly created instructions and deleted instructions.
+void emitDestroyOperation(SILBuilder &builder, SILLocation loc,
+                          SILValue operand, InstModCallbacks callbacks);
+
 /// Collect all (transitive) users of \p inst which just copy or destroy \p
 /// inst.
 ///
@@ -366,6 +371,7 @@ void getConsumedPartialApplyArgs(PartialApplyInst *pai,
 /// destroys, i.e. if \p inst can be considered as "dead".
 bool collectDestroys(SingleValueInstruction *inst,
                      SmallVectorImpl<SILInstruction *> &destroys);
+
 /// If Closure is a partial_apply or thin_to_thick_function with only local
 /// ref count users and a set of post-dominating releases:
 ///

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -1187,28 +1187,16 @@ static bool shouldDestroyPartialApplyCapturedArg(SILValue arg,
   return true;
 }
 
-// *HEY YOU, YES YOU, PLEASE READ*. Even though a textual partial apply is
-// printed with the convention of the closed over function upon it, all
-// non-inout arguments to a partial_apply are passed at +1. This includes
-// arguments that will eventually be passed as guaranteed or in_guaranteed to
-// the closed over function. This is because the partial apply is building up a
-// boxed aggregate to send off to the closed over function. Of course when you
-// call the function, the proper conventions will be used.
-void swift::releasePartialApplyCapturedArg(SILBuilder &builder, SILLocation loc,
-                                           SILValue arg,
-                                           SILParameterInfo paramInfo,
-                                           InstModCallbacks callbacks) {
-  if (!shouldDestroyPartialApplyCapturedArg(arg, paramInfo,
-                                            builder.getFunction()))
-    return;
-
-  // Otherwise, we need to destroy the argument. If we have an address, we
-  // insert a destroy_addr and return. Any live range issues must have been
-  // dealt with by our caller.
-  if (arg->getType().isAddress()) {
-    // Then emit the destroy_addr for this arg
-    SILInstruction *newInst = builder.emitDestroyAddrAndFold(loc, arg);
-    callbacks.createdNewInst(newInst);
+void swift::emitDestroyOperation(SILBuilder &builder, SILLocation loc,
+                                 SILValue operand, InstModCallbacks callbacks) {
+  // If we have an address, we insert a destroy_addr and return. Any live range
+  // issues must have been dealt with by our caller.
+  if (operand->getType().isAddress()) {
+    // Then emit the destroy_addr for this operand. This function does not
+    // delete any instructions
+    SILInstruction *newInst = builder.emitDestroyAddrAndFold(loc, operand);
+    if (newInst != nullptr)
+      callbacks.createdNewInst(newInst);
     return;
   }
 
@@ -1217,12 +1205,12 @@ void swift::releasePartialApplyCapturedArg(SILBuilder &builder, SILLocation loc,
 
   // If we have qualified ownership, we should just emit a destroy value.
   if (builder.getFunction().hasOwnership()) {
-    callbacks.createdNewInst(builder.createDestroyValue(loc, arg));
+    callbacks.createdNewInst(builder.createDestroyValue(loc, operand));
     return;
   }
 
-  if (arg->getType().hasReferenceSemantics()) {
-    auto u = builder.emitStrongRelease(loc, arg);
+  if (operand->getType().hasReferenceSemantics()) {
+    auto u = builder.emitStrongRelease(loc, operand);
     if (u.isNull())
       return;
 
@@ -1235,7 +1223,7 @@ void swift::releasePartialApplyCapturedArg(SILBuilder &builder, SILLocation loc,
     return;
   }
 
-  auto u = builder.emitReleaseValue(loc, arg);
+  auto u = builder.emitReleaseValue(loc, operand);
   if (u.isNull())
     return;
 
@@ -1247,6 +1235,24 @@ void swift::releasePartialApplyCapturedArg(SILBuilder &builder, SILLocation loc,
   callbacks.createdNewInst(u.get<ReleaseValueInst *>());
 }
 
+// *HEY YOU, YES YOU, PLEASE READ*. Even though a textual partial apply is
+// printed with the convention of the closed over function upon it, all
+// non-inout arguments to a partial_apply are passed at +1. This includes
+// arguments that will eventually be passed as guaranteed or in_guaranteed to
+// the closed over function. This is because the partial apply is building up a
+// boxed aggregate to send off to the closed over function. Of course when you
+// call the function, the proper conventions will be used.
+void swift::releasePartialApplyCapturedArg(SILBuilder &builder, SILLocation loc,
+                                           SILValue arg,
+                                           SILParameterInfo paramInfo,
+                                           InstModCallbacks callbacks) {
+  if (!shouldDestroyPartialApplyCapturedArg(arg, paramInfo,
+                                            builder.getFunction()))
+    return;
+
+  emitDestroyOperation(builder, loc, arg, callbacks);
+}
+
 void swift::deallocPartialApplyCapturedArg(SILBuilder &builder, SILLocation loc,
                                            SILValue arg,
                                            SILParameterInfo paramInfo) {
@@ -1415,11 +1421,7 @@ bool swift::tryDeleteDeadClosure(SingleValueInstruction *closure,
       for (Operand *argOp : argsToHandle) {
         SILValue arg = argOp->get();
         SILBuilderWithScope builder(pai, builderCtxt);
-        if (arg->getType().isObject()) {
-          builder.emitDestroyValueOperation(pai->getLoc(), arg);
-        } else {
-          builder.emitDestroyAddr(pai->getLoc(), arg);
-        }
+        emitDestroyOperation(builder, pai->getLoc(), arg, callbacks);
       }
     }
   }