Fix SILCombine of partial_apply to correctly extend store_borrow when needed

Author: meg-gupta

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -247,7 +247,7 @@ void emitDestroyOperation(SILBuilder &builder, SILLocation loc,
 /// Returns true, if there are no other users beside those collected in \p
 /// destroys, i.e. if \p inst can be considered as "dead".
 bool collectDestroys(SingleValueInstruction *inst,
-                     SmallVectorImpl<SILInstruction *> &destroys);
+                     SmallVectorImpl<Operand *> &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:

include/swift/SILOptimizer/Utils/ValueLifetime.h

@@ -138,6 +138,20 @@ class ValueLifetimeAnalysis {
     propagateLiveness();
   }
 
+  ValueLifetimeAnalysis(SILArgument *def, ArrayRef<Operand *> useRange)
+      : defValue(def), inLiveBlocks(def->getFunction()), userSet() {
+    for (auto *use : useRange)
+      userSet.insert(use->getUser());
+    propagateLiveness();
+  }
+
+  ValueLifetimeAnalysis(SILInstruction *def, ArrayRef<Operand *> useRange)
+      : defValue(def), inLiveBlocks(def->getFunction()), userSet() {
+    for (auto *use : useRange)
+      userSet.insert(use->getUser());
+    propagateLiveness();
+  }
+
   /// Compute the LifetimeBoundary--the last users and boundary edges. This
   /// always succeeds.
   ///

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -25,6 +25,7 @@
 #include "swift/SIL/SILDebugInfoExpression.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SIL/SILUndef.h"
+#include "swift/SIL/ScopedAddressUtils.h"
 #include "swift/SIL/TypeLowering.h"
 #include "swift/SILOptimizer/Analysis/ARCAnalysis.h"
 #include "swift/SILOptimizer/Analysis/Analysis.h"
@@ -971,16 +972,16 @@ void swift::getConsumedPartialApplyArgs(PartialApplyInst *pai,
 }
 
 bool swift::collectDestroys(SingleValueInstruction *inst,
-                            SmallVectorImpl<SILInstruction *> &destroys) {
+                            SmallVectorImpl<Operand *> &destroys) {
   bool isDead = true;
   for (Operand *use : inst->getUses()) {
     SILInstruction *user = use->getUser();
     if (useHasTransitiveOwnership(user)) {
       if (!collectDestroys(cast<SingleValueInstruction>(user), destroys))
         isDead = false;
-      destroys.push_back(user);
+      destroys.push_back(use);
     } else if (useDoesNotKeepValueAlive(user)) {
-      destroys.push_back(user);
+      destroys.push_back(use);
     } else {
       isDead = false;
     }
@@ -996,7 +997,7 @@ bool swift::collectDestroys(SingleValueInstruction *inst,
 ///       weirdness of the old retain/release ARC model. Most likely this will
 ///       not be needed anymore with OSSA.
 static bool keepArgsOfPartialApplyAlive(PartialApplyInst *pai,
-                                        ArrayRef<SILInstruction *> paiUsers,
+                                        ArrayRef<Operand *> paiUses,
                                         SILBuilderContext &builderCtxt,
                                         InstModCallbacks callbacks) {
   SmallVector<Operand *, 8> argsToHandle;
@@ -1008,7 +1009,7 @@ static bool keepArgsOfPartialApplyAlive(PartialApplyInst *pai,
   // Compute the set of endpoints, which will be used to insert destroys of
   // temporaries. This may fail if the frontier is located on a critical edge
   // which we may not split.
-  ValueLifetimeAnalysis vla(pai, paiUsers);
+  ValueLifetimeAnalysis vla(pai, paiUses);
 
   ValueLifetimeAnalysis::Frontier partialApplyFrontier;
   if (!vla.computeFrontier(partialApplyFrontier,
@@ -1068,7 +1069,7 @@ bool swift::tryDeleteDeadClosure(SingleValueInstruction *closure,
 
   // Collect all destroys of the closure (transitively including destorys of
   // copies) and check if those are the only uses of the closure.
-  SmallVector<SILInstruction *, 16> closureDestroys;
+  SmallVector<Operand *, 16> closureDestroys;
   if (!collectDestroys(closure, closureDestroys))
     return false;
 
@@ -1102,7 +1103,8 @@ bool swift::tryDeleteDeadClosure(SingleValueInstruction *closure,
 
   // Delete all copy and destroy instructions in order so that leaf uses are
   // deleted first.
-  for (SILInstruction *user : closureDestroys) {
+  for (auto *use : closureDestroys) {
+    auto *user = use->getUser();
     assert(
         (useDoesNotKeepValueAlive(user) || useHasTransitiveOwnership(user)) &&
         "We expect only ARC operations without "

lib/SILOptimizer/Utils/PartialApplyCombiner.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/SIL/SILValue.h"
+#include "swift/SIL/ScopedAddressUtils.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
 #include "swift/SILOptimizer/Utils/ValueLifetime.h"
 
@@ -59,34 +60,44 @@ class PartialApplyCombiner {
 /// apply instructions.
 bool PartialApplyCombiner::copyArgsToTemporaries(
     ArrayRef<FullApplySite> applies) {
+  SmallVector<Operand *, 8> argsToHandle;
+
+  // Find args that need extension for a non-stack partial_apply
   // A partial_apply [stack]'s argument are not owned by the partial_apply and
   // therefore their lifetime must outlive any uses.
-  if (pai->isOnStack())
-    return true;
-
-  SmallVector<Operand *, 8> argsToHandle;
-  getConsumedPartialApplyArgs(pai, argsToHandle,
-                              /*includeTrivialAddrArgs*/ true);
-  if (argsToHandle.empty())
-    return true;
+  if (!pai->isOnStack()) {
+    getConsumedPartialApplyArgs(pai, argsToHandle,
+                                /*includeTrivialAddrArgs*/ true);
+  }
 
   // Compute the set of endpoints, which will be used to insert destroys of
   // temporaries.
-  SmallVector<SILInstruction *, 16> paiUsers;
+  SmallVector<Operand *, 16> paiUses;
 
   // Of course we must inlude all apply instructions which we want to optimize.
   for (FullApplySite ai : applies) {
-    paiUsers.push_back(ai.getInstruction());
+    paiUses.push_back(ai.getCalleeOperand());
+  }
+
+  SmallVector<StoreBorrowInst *, 8> storeBorrowsToHandle;
+  for (auto arg : pai->getArguments()) {
+    if (auto *sbi = dyn_cast<StoreBorrowInst>(arg)) {
+      storeBorrowsToHandle.push_back(sbi);
+    }
   }
 
+  if (argsToHandle.empty() && storeBorrowsToHandle.empty()) {
+    return true;
+  }
   // Also include all destroys in the liferange for the arguments.
   // This is needed for later processing in tryDeleteDeadClosure: in case the
-  // pai gets dead after this optimization, tryDeleteDeadClosure relies on that
-  // we already copied the pai arguments to extend their lifetimes until the pai
-  // is finally destroyed.
-  collectDestroys(pai, paiUsers);
+  // pai gets dead after this optimization, tryDeleteDeadClosure relies on
+  // that we already copied the pai arguments to extend their lifetimes until
+  // the pai is finally destroyed.
+  collectDestroys(pai, paiUses);
 
-  ValueLifetimeAnalysis vla(pai, paiUsers);
+  ValueLifetimeAnalysis vla(pai,
+                            llvm::makeArrayRef(paiUses.begin(), paiUses.end()));
   ValueLifetimeAnalysis::Frontier partialApplyFrontier;
 
   // Computing the frontier may fail if the frontier is located on a critical
@@ -114,6 +125,23 @@ bool PartialApplyCombiner::copyArgsToTemporaries(
     // allocated temporary) at the end of the partial_apply's lifetime.
     endLifetimeAtFrontier(tmp, partialApplyFrontier, builderCtxt, callbacks);
   }
+
+  DeadEndBlocks deBlocks(pai->getFunction());
+  for (auto *storeBorrow : storeBorrowsToHandle) {
+    if (extendStoreBorrow(storeBorrow, paiUses, &deBlocks, callbacks)) {
+      continue;
+    }
+    SILBuilderWithScope builder(pai, builderCtxt);
+    // Copy address-arguments into a stack-allocated temporary.
+    auto *asi = builder.createAllocStack(pai->getLoc(), storeBorrow->getType());
+    builder.createCopyAddr(pai->getLoc(), storeBorrow, asi, IsTake_t::IsNotTake,
+                           IsInitialization_t::IsInitialization);
+    argToTmpCopy.insert(std::make_pair(storeBorrow, asi));
+
+    // Destroy the argument value (either as SSA value or in the stack-
+    // allocated temporary) at the end of the partial_apply's lifetime.
+    endLifetimeAtFrontier(asi, partialApplyFrontier, builderCtxt, callbacks);
+  }
   return true;
 }