Merge pull request #65780 from meg-gupta/paonstackfix

Fix insertion of dealloc_stack for partial_apply [on_stack] while going to non-OSSA

Author: swift-ci

lib/SILOptimizer/Mandatory/OwnershipModelEliminator.cpp

@@ -485,19 +485,11 @@ bool OwnershipModelEliminatorVisitor::visitDestroyValueInst(
     DestroyValueInst *dvi) {
   // Nonescaping closures are represented ultimately as trivial pointers to
   // their context, but we use ownership to do borrow checking of their captures
-  // in OSSA. Now that we're eliminating ownership, fold away destroys, unless
-  // we're destroying the original partial_apply, in which case this is where
-  // we dealloc_stack the context.
+  // in OSSA. Now that we're eliminating ownership, fold away destroys.
   auto operand = dvi->getOperand();
   auto operandTy = operand->getType();
   if (auto operandFnTy = operandTy.getAs<SILFunctionType>()){
     if (operandFnTy->isTrivialNoEscape()) {
-      if (auto origPA = dvi->getNonescapingClosureAllocation()) {
-        withBuilder<void>(dvi, [&](SILBuilder &b, SILLocation loc) {
-          b.createDeallocStack(loc, origPA);
-        });
-      }
-      
       eraseInstruction(dvi);
       return true;
     }
@@ -627,9 +619,41 @@ static bool stripOwnership(SILFunction &func) {
   if (func.isExternalDeclaration())
     return false;
 
+  llvm::DenseMap<PartialApplyInst *, SmallVector<SILInstruction *>>
+      lifetimeEnds;
+
+  // Nonescaping closures are represented ultimately as trivial pointers to
+  // their context, but we use ownership to do borrow checking of their captures
+  // in OSSA. Now that we're eliminating ownership, we need to dealloc_stack the
+  // context at its lifetime ends.
+  // partial_apply's lifetime ends has to be gathered before we begin to leave
+  // OSSA, but no dealloc_stack can be emitted until after we leave OSSA.
+  for (auto &block : func) {
+    for (auto &ii : block) {
+      auto *pai = dyn_cast<PartialApplyInst>(&ii);
+      if (!pai || !pai->isOnStack()) {
+        continue;
+      }
+      pai->visitOnStackLifetimeEnds([&](Operand *op) {
+        lifetimeEnds[pai].push_back(op->getUser());
+        return true;
+      });
+    }
+  }
+
   // Set F to have unqualified ownership.
   func.setOwnershipEliminated();
 
+  // Now that we are in non-ossa, create dealloc_stack at partial_apply's
+  // lifetime ends
+  for (auto &it : lifetimeEnds) {
+    auto *pai = it.first;
+    for (auto *lifetimeEnd : it.second) {
+      SILBuilderWithScope(lifetimeEnd->getNextInstruction())
+          .createDeallocStack(lifetimeEnd->getLoc(), pai);
+    }
+  }
+
   bool madeChange = false;
   SmallVector<SILInstruction *, 32> createdInsts;
   OwnershipModelEliminatorVisitor visitor(func);

test/SILOptimizer/ownership_model_eliminator.sil

@@ -393,4 +393,21 @@ bb0(%0 : $*C):
   dealloc_stack %1 : $*C
   %9999 = tuple()
   return %9999 : $()
-}
+}
+
+sil @closure1: $@convention(thin) (@guaranteed C, @inout_aliasable C) -> ()
+sil @closure2 : $@convention(thin) (@guaranteed @noescape @callee_guaranteed () -> ()) -> ()
+
+// Ensure no assertion about missing dealloc_stack for partial_apply [on_stack] after OME
+sil [ossa] @test_partial_apply_on_stack: $@convention(thin) (@guaranteed C, @inout C) -> () {
+bb0(%0 : @guaranteed $C, %1 : $*C):
+  %f1 = function_ref @closure1: $@convention(thin) (@guaranteed C, @inout_aliasable C) -> ()
+  %pa1 = partial_apply [callee_guaranteed] [on_stack] %f1(%0, %1) : $@convention(thin) (@guaranteed C, @inout_aliasable C) -> ()
+  %md = mark_dependence %pa1 : $@noescape @callee_guaranteed () -> () on %1 : $*C
+  %f2 = function_ref @closure2 : $@convention(thin) (@guaranteed @noescape @callee_guaranteed () -> ()) -> ()
+  %pa2 = partial_apply [callee_guaranteed] %f2(%md) : $@convention(thin) (@guaranteed @noescape @callee_guaranteed () -> ()) -> ()
+  destroy_value %pa2 : $@callee_guaranteed () -> ()
+  %t = tuple ()
+  return %t : $()
+}
+