Factor out helper to find the original PartialApplyInst deallocated by a DestroyValueInst in OSSA, if any

Author: jckarter

include/swift/SIL/SILCloner.h

@@ -1962,23 +1962,10 @@ void SILCloner<ImplClass>::visitDestroyValueInst(DestroyValueInst *Inst) {
       if (fnTy->isTrivialNoEscape()) {
         // Destroying the partial_apply [stack] becomes the stack deallocation
         // of the context.
-        // Look through mark_dependence and other wrapper instructions.
-        SILValue deallocOperand = Inst->getOperand();
-        while (true) {
-          if (auto mdi = dyn_cast<MarkDependenceInst>(deallocOperand)) {
-            deallocOperand = mdi->getValue();
-          } else if (isa<ConvertEscapeToNoEscapeInst>(deallocOperand)) {
-            break;
-          } else if (auto conv = dyn_cast<ConversionInst>(deallocOperand)) {
-            deallocOperand = conv->getConverted();
-          } else {
-            break;
-          }
-        }
-        if (isa<PartialApplyInst>(deallocOperand)) {
+        if (auto origPA = Inst->getNonescapingClosureAllocation()) {
           recordClonedInstruction(Inst,
             getBuilder().createDeallocStack(getOpLocation(Inst->getLoc()),
-                                            getOpValue(deallocOperand)));
+                                            getOpValue(origPA)));
         }
 
         return;

include/swift/SIL/SILInstruction.h

@@ -8193,6 +8193,10 @@ class DestroyValueInst
   void setPoisonRefs(bool poisonRefs = true) {
     sharedUInt8().DestroyValueInst.poisonRefs = poisonRefs;
   }
+  
+  /// If the value being destroyed is a stack allocation of a nonescaping
+  /// closure, then return the PartialApplyInst that allocated the closure.
+  PartialApplyInst *getNonescapingClosureAllocation() const;
 };
 
 class MoveValueInst

lib/SIL/IR/SILInstruction.cpp

@@ -1705,6 +1705,47 @@ PartialApplyInst::visitOnStackLifetimeEnds(
   return !noUsers;
 }
 
+PartialApplyInst *
+DestroyValueInst::getNonescapingClosureAllocation() const {
+  SILValue operand = getOperand();
+  auto operandFnTy = operand->getType().getAs<SILFunctionType>();
+  // The query doesn't make sense if we aren't operating on a noescape closure
+  // to begin with.
+  if (!operandFnTy || !operandFnTy->isTrivialNoEscape()) {
+    return nullptr;
+  }
+
+  // Look through marker and conversion instructions that would forward
+  // ownership of the original partial application.
+  while (true) {
+    if (auto mdi = dyn_cast<MarkDependenceInst>(operand)) {
+      operand = mdi->getValue();
+      continue;
+    } else if (isa<ConvertEscapeToNoEscapeInst>(operand)) {
+      // Stop at a conversion from escaping closure, since there's no stack
+      // allocation in that case.
+      return nullptr;
+    } else if (auto conv = dyn_cast<ConversionInst>(operand)) {
+      operand = conv->getConverted();
+      continue;
+    } else if (auto pa = dyn_cast<PartialApplyInst>(operand)) {
+      // If we found the `[on_stack]` partial apply, we're done.
+      if (pa->isOnStack()) {
+        return pa;
+      }
+      // Any other kind of partial apply fails to pass muster.
+      return nullptr;
+    } else {
+      // The original partial_apply instruction should only be forwarded
+      // through one of the above instructions. Anything else should lead us
+      // to a copy or borrow of the closure from somewhere else.
+      // TODO: Assert that this is one of those cases, such as a SILArgument,
+      // copy_value, etc.
+      return nullptr;
+    }
+  }
+}
+
 #ifndef NDEBUG
 
 //---

lib/SILOptimizer/Mandatory/OwnershipModelEliminator.cpp

@@ -486,30 +486,14 @@ bool OwnershipModelEliminatorVisitor::visitDestroyValueInst(
   // we dealloc_stack the context.
   auto operand = dvi->getOperand();
   auto operandTy = operand->getType();
-  if (auto operandFnTy = operandTy.getAs<SILFunctionType>()) {
+  if (auto operandFnTy = operandTy.getAs<SILFunctionType>()){
     if (operandFnTy->isTrivialNoEscape()) {
-      // Look through mark_dependence and other wrapper instructions.
-      SILValue deallocOperand = operand;
-      while (true) {
-        if (auto mdi = dyn_cast<MarkDependenceInst>(deallocOperand)) {
-          deallocOperand = mdi->getValue();
-        } else if (isa<ConvertEscapeToNoEscapeInst>(deallocOperand)) {
-          // If there's a surviving convert_escape_to_noescape that wasn't
-          // turned into a stack closure, then stop here and just delete the
-          // destroy_value, since the original escaping closure's lifetime
-          // will persist.
-          break;
-        } else if (auto conv = dyn_cast<ConversionInst>(deallocOperand)) {
-          deallocOperand = conv->getConverted();
-        } else {
-          break;
-        }
-      }
-      if (isa<PartialApplyInst>(deallocOperand)) {
+      if (auto origPA = dvi->getNonescapingClosureAllocation()) {
         withBuilder<void>(dvi, [&](SILBuilder &b, SILLocation loc) {
-          b.createDeallocStack(loc, deallocOperand);
+          b.createDeallocStack(loc, origPA);
         });
       }
+      
       eraseInstruction(dvi);
       return true;
     }