Merge pull request #68381 from nate-chandler/cherrypick/release/5.9/gh68328

5.9: [MoveChecker] Complete lifetimes before checking.

Author: nate-chandler

include/swift/SIL/OwnershipUseVisitor.h

@@ -283,11 +283,23 @@ bool OwnershipUseVisitor<Impl>::visitInnerBorrowScopeEnd(Operand *borrowEnd) {
   case OperandOwnership::EndBorrow:
     return handleUsePoint(borrowEnd, UseLifetimeConstraint::NonLifetimeEnding);
 
-  case OperandOwnership::Reborrow:
+  case OperandOwnership::Reborrow: {
     if (!asImpl().handleInnerReborrow(borrowEnd))
       return false;
 
     return handleUsePoint(borrowEnd, UseLifetimeConstraint::NonLifetimeEnding);
+  }
+  case OperandOwnership::DestroyingConsume: {
+    // partial_apply [on_stack] can introduce borrowing operand and can have
+    // destroy_value consumes.
+    auto *pai = dyn_cast<PartialApplyInst>(borrowEnd->get());
+    // TODO: When we have ForwardingInstruction abstraction, walk the use-def
+    // chain to ensure we have a partial_apply [on_stack] def.
+    assert(pai && pai->isOnStack() ||
+           OwnershipForwardingMixin::get(
+               cast<SingleValueInstruction>(borrowEnd->get())));
+    return handleUsePoint(borrowEnd, UseLifetimeConstraint::NonLifetimeEnding);
+  }
 
   default:
     llvm_unreachable("expected borrow scope end");
@@ -388,8 +400,9 @@ bool OwnershipUseVisitor<Impl>::visitGuaranteedUse(Operand *use) {
   case OperandOwnership::ForwardingUnowned:
   case OperandOwnership::UnownedInstantaneousUse:
   case OperandOwnership::BitwiseEscape:
-  case OperandOwnership::EndBorrow:
     return handleUsePoint(use, UseLifetimeConstraint::NonLifetimeEnding);
+  case OperandOwnership::EndBorrow:
+    return handleUsePoint(use, UseLifetimeConstraint::LifetimeEnding);
 
   case OperandOwnership::Reborrow:
     if (!asImpl().handleOuterReborrow(use))

include/swift/SIL/OwnershipUtils.h

@@ -477,6 +477,9 @@ struct BorrowingOperand {
   /// valid BorrowedValue instance.
   BorrowedValue getBorrowIntroducingUserResult();
 
+  /// Return the borrowing operand's value.
+  SILValue getScopeIntroducingUserResult();
+
   /// Compute the implicit uses that this borrowing operand "injects" into the
   /// set of its operands uses.
   ///

include/swift/SIL/SILInstruction.h

@@ -1261,6 +1261,11 @@ class OwnershipForwardingMixin {
   ///
   /// \p inst is an OwnershipForwardingMixin
   static bool isAddressOnly(SILInstruction *inst);
+
+  // Call \p visitor on all forwarded results of the current forwarding
+  // operation.
+  static bool visitForwardedValues(SILInstruction *inst,
+                                   function_ref<bool(SILValue)> visitor);
 };
 
 /// A single value inst that forwards a static ownership from its first operand.

lib/SIL/IR/SILInstructions.cpp

@@ -3276,6 +3276,31 @@ bool OwnershipForwardingMixin::isAddressOnly(SILInstruction *inst) {
   return inst->getOperand(0)->getType().isAddressOnly(*inst->getFunction());
 }
 
+bool OwnershipForwardingMixin::visitForwardedValues(
+    SILInstruction *inst, function_ref<bool(SILValue)> visitor) {
+  if (auto *svi = dyn_cast<SingleValueInstruction>(inst)) {
+    return visitor(svi);
+  }
+  if (auto *mvri = dyn_cast<MultipleValueInstruction>(inst)) {
+    return llvm::all_of(mvri->getResults(), [&](SILValue value) {
+      if (value->getOwnershipKind() == OwnershipKind::None)
+        return true;
+      return visitor(value);
+    });
+  }
+  auto *ti = cast<TermInst>(inst);
+  assert(ti->mayHaveTerminatorResult());
+  return llvm::all_of(ti->getSuccessorBlocks(), [&](SILBasicBlock *succBlock) {
+    // If we do not have any arguments, then continue.
+    if (succBlock->args_empty())
+      return true;
+
+    auto args = succBlock->getSILPhiArguments();
+    assert(args.size() == 1 && "Transforming terminator with multiple args?!");
+    return visitor(args[0]);
+  });
+}
+
 // This may be called in an invalid SIL state. SILCombine creates new
 // terminators in non-terminator position and defers deleting the original
 // terminator until after all modification.

lib/SIL/Utils/OSSALifetimeCompletion.cpp

@@ -129,7 +129,7 @@ static bool endLifetimeAtUnreachableBlocks(SILValue value,
       changed = true;
     }
     for (auto *successor : block->getSuccessorBlocks()) {
-      deadEndBlocks.push(successor);
+      deadEndBlocks.pushIfNotVisited(successor);
     }
   }
   return changed;

lib/SIL/Utils/OwnershipLiveness.cpp

@@ -140,8 +140,10 @@ struct InteriorLivenessVisitor :
   /// Handles begin_borrow, load_borrow, store_borrow, begin_apply.
   bool handleInnerBorrow(BorrowingOperand borrowingOperand) {
     if (handleInnerScopeCallback) {
-      handleInnerScopeCallback(
-        borrowingOperand.getBorrowIntroducingUserResult().value);
+      auto value = borrowingOperand.getScopeIntroducingUserResult();
+      if (value) {
+        handleInnerScopeCallback(value);
+      }
     }
     return true;
   }

lib/SIL/Utils/OwnershipUtils.cpp

@@ -751,6 +751,30 @@ BorrowedValue BorrowingOperand::getBorrowIntroducingUserResult() {
   llvm_unreachable("covered switch");
 }
 
+SILValue BorrowingOperand::getScopeIntroducingUserResult() {
+  switch (kind) {
+  case BorrowingOperandKind::Invalid:
+  case BorrowingOperandKind::Yield:
+  case BorrowingOperandKind::Apply:
+  case BorrowingOperandKind::TryApply:
+    return SILValue();
+
+  case BorrowingOperandKind::BeginAsyncLet:
+  case BorrowingOperandKind::PartialApplyStack:
+  case BorrowingOperandKind::BeginBorrow:
+    return cast<SingleValueInstruction>(op->getUser());
+
+  case BorrowingOperandKind::BeginApply:
+    return cast<BeginApplyInst>(op->getUser())->getTokenResult();
+
+  case BorrowingOperandKind::Branch: {
+    PhiOperand phiOp(op);
+    return phiOp.getValue();
+  }
+  }
+  llvm_unreachable("covered switch");
+}
+
 void BorrowingOperand::getImplicitUses(
     SmallVectorImpl<Operand *> &foundUses) const {
   // FIXME: this visitScopeEndingUses should never return false once dead

lib/SILOptimizer/Mandatory/MoveOnlyChecker.cpp

@@ -30,6 +30,7 @@
 #include "swift/SIL/FieldSensitivePrunedLiveness.h"
 #include "swift/SIL/InstructionUtils.h"
 #include "swift/SIL/MemAccessUtils.h"
+#include "swift/SIL/OSSALifetimeCompletion.h"
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/PrunedLiveness.h"
 #include "swift/SIL/SILArgument.h"
@@ -86,6 +87,7 @@ struct MoveOnlyChecker {
   }
 
   void checkObjects();
+  void completeObjectLifetimes(ArrayRef<MarkMustCheckInst *>);
   void checkAddresses();
 };
 
@@ -109,10 +111,75 @@ void MoveOnlyChecker::checkObjects() {
     return;
   }
 
+  completeObjectLifetimes(moveIntroducersToProcess.getArrayRef());
+
   MoveOnlyObjectChecker checker{diagnosticEmitter, domTree, poa, allocator};
   madeChange |= checker.check(moveIntroducersToProcess);
 }
 
+void MoveOnlyChecker::completeObjectLifetimes(
+    ArrayRef<MarkMustCheckInst *> insts) {
+  // TODO: Delete once OSSALifetimeCompletion is run as part of SILGenCleanup.
+  OSSALifetimeCompletion completion(fn, domTree);
+
+  // Collect all values derived from each mark_unresolved_non_copyable_value
+  // instruction via ownership instructions and phis.
+  ValueWorklist transitiveValues(fn);
+  for (auto *inst : insts) {
+    transitiveValues.push(inst);
+  }
+  while (auto value = transitiveValues.pop()) {
+    for (auto *use : value->getUses()) {
+      auto *user = use->getUser();
+      switch (user->getKind()) {
+      case SILInstructionKind::BeginBorrowInst:
+      case SILInstructionKind::CopyValueInst:
+      case SILInstructionKind::MoveValueInst:
+        transitiveValues.pushIfNotVisited(cast<SingleValueInstruction>(user));
+        break;
+      case SILInstructionKind::BranchInst: {
+        PhiOperand po(use);
+        transitiveValues.pushIfNotVisited(po.getValue());
+        break;
+      }
+      default: {
+        auto *forward = OwnershipForwardingMixin::get(user);
+        if (!forward)
+          continue;
+        OwnershipForwardingMixin::visitForwardedValues(
+            user, [&transitiveValues](auto forwarded) {
+              transitiveValues.pushIfNotVisited(forwarded);
+              return true;
+            });
+        break;
+      }
+      }
+    }
+  }
+  // Complete the lifetime of each collected value.  This is a subset of the
+  // work that SILGenCleanup will do.
+  for (auto *block : poa->get(fn)->getPostOrder()) {
+    for (SILInstruction &inst : reverse(*block)) {
+      for (auto result : inst.getResults()) {
+        if (!transitiveValues.isVisited(result))
+          continue;
+        if (completion.completeOSSALifetime(result) ==
+            LifetimeCompletion::WasCompleted) {
+          madeChange = true;
+        }
+      }
+    }
+    for (SILArgument *arg : block->getArguments()) {
+      if (!transitiveValues.isVisited(arg))
+        continue;
+      if (completion.completeOSSALifetime(arg) ==
+          LifetimeCompletion::WasCompleted) {
+        madeChange = true;
+      }
+    }
+  }
+}
+
 void MoveOnlyChecker::checkAddresses() {
   unsigned diagCount = diagnosticEmitter.getDiagnosticCount();
   SmallSetVector<MarkMustCheckInst *, 32> moveIntroducersToProcess;