Merge pull request swiftlang#63798 from jckarter/owned-nonescaping-closure-asserts

[SIL] Add assertions to check assumptions about owned noescape closures.

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

@@ -1681,11 +1681,11 @@ static bool visitRecursivelyLifetimeEndingUses(
 
     // There shouldn't be any dead-end consumptions of a nonescaping
     // partial_apply that don't forward it along, aside from destroy_value.
-    assert(use->getUser()->hasResults());
-    for (auto result : use->getUser()->getResults()) {
-      if (!visitRecursivelyLifetimeEndingUses(result, noUsers, func)) {
-        return false;
-      }
+    assert(use->getUser()->hasResults()
+           && use->getUser()->getNumResults() == 1);
+    if (!visitRecursivelyLifetimeEndingUses(use->getUser()->getResult(0),
+                                            noUsers, func)) {
+      return false;
     }
   }
   return true;
@@ -1705,6 +1705,53 @@ 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)
+               || isa<ThinToThickFunctionInst>(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.
+      assert((isa<CopyValueInst>(operand)
+              || isa<SILArgument>(operand)
+              || isa<DifferentiableFunctionInst>(operand)
+              || isa<DifferentiableFunctionExtractInst>(operand)
+              || isa<LoadInst>(operand)
+              || (operand->dump(), false))
+             && "unexpected forwarding instruction for noescape closure");
+      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;
     }