Merge pull request #83624 from nate-chandler/rdar157291161

[SemanticARCOpts] Fix miscompile at dead-ends.

Author: nate-chandler

include/swift/SIL/OwnershipUtils.h

@@ -630,6 +630,21 @@ struct BorrowedValue {
   /// necessarilly dominated by the borrow scope.
   void computeTransitiveLiveness(MultiDefPrunedLiveness &liveness) const;
 
+  /// Whether \p insts are completely within this borrow introducer's local
+  /// scope.
+  ///
+  /// Precondition: \p insts are dominated by the local borrow introducer.
+  ///
+  /// This ignores reborrows. The assumption is that, since \p insts are
+  /// dominated by this local scope, checking the extended borrow scope should
+  /// not be necessary to determine they are within the scope.
+  ///
+  /// \p deadEndBlocks is optional during transition. It will be completely
+  /// removed in an upcoming commit.
+  template <typename Instructions>
+  bool areWithinExtendedScope(Instructions insts,
+                              DeadEndBlocks *deadEndBlocks) const;
+
   /// Returns true if \p uses are completely within this borrow introducer's
   /// local scope.
   ///
@@ -1426,6 +1441,43 @@ bool isRedundantMoveValue(MoveValueInst *mvi);
 /// `forEndBorrowValue`, which is the operand value of an `end_borrow`.
 void updateReborrowFlags(SILValue forEndBorrowValue);
 
+/// A location at which a value is used.  Abstracts over explicit uses
+/// (operands) and implicit uses (instructions).
+struct UsePoint {
+  using Value = llvm::PointerUnion<SILInstruction *, Operand *>;
+  Value value;
+
+  UsePoint(Operand *op) : value(op) {}
+  UsePoint(SILInstruction *inst) : value(inst) {}
+  UsePoint(Value value) : value(value) {}
+
+  SILInstruction *getInstruction() const {
+    if (auto *op = dyn_cast<Operand *>(value)) {
+      return op->getUser();
+    }
+    return cast<SILInstruction *>(value);
+  }
+
+  Operand *getOperandOrNull() const { return dyn_cast<Operand *>(value); }
+
+  Operand *getOperand() const { return cast<Operand *>(value); }
+};
+
+struct UsePointToInstruction {
+  SILInstruction *operator()(const UsePoint point) const {
+    return point.getInstruction();
+  }
+};
+
+using UsePointInstructionRange =
+    TransformRange<ArrayRef<UsePoint>, UsePointToInstruction>;
+
+struct PointToOperand {
+  Operand *operator()(const UsePoint point) const { return point.getOperand(); }
+};
+
+using PointOperandRange = TransformRange<ArrayRef<UsePoint>, PointToOperand>;
+
 } // namespace swift
 
 #endif

include/swift/SIL/PrunedLiveness.h

@@ -729,6 +729,20 @@ class PrunedLiveRange : public PrunedLiveness {
   bool isWithinBoundary(SILInstruction *inst,
                         DeadEndBlocks *deadEndBlocks) const;
 
+  /// Whether all \p insts are between this def and the liveness boundary;
+  /// \p deadEndBlocks is optional.
+  template <typename Instructions>
+  bool areWithinBoundary(Instructions insts,
+                         DeadEndBlocks *deadEndBlocks) const {
+    assert(asImpl().isInitialized());
+
+    for (auto *inst : insts) {
+      if (!isWithinBoundary(inst, deadEndBlocks))
+        return false;
+    }
+    return true;
+  };
+
   /// Returns true when all \p uses are between this def and the liveness
   /// boundary \p deadEndBlocks is optional.
   bool areUsesWithinBoundary(ArrayRef<Operand *> uses,

include/swift/SIL/SILInstruction.h

@@ -670,6 +670,10 @@ class SILInstruction : public llvm::ilist_node<SILInstruction> {
   using TransformedOperandValueRange =
     OptionalTransformRange<ArrayRef<Operand>, OperandToTransformedValue>;
 
+  /// Functor for Operand::getUser()
+  struct OperandToUser;
+  using OperandUserRange = TransformRange<ArrayRef<Operand *>, OperandToUser>;
+
   static OperandValueRange getOperandValues(ArrayRef<Operand*> operands);
 
   OperandRefValueRange getOperandValues() const;
@@ -1021,6 +1025,12 @@ struct SILInstruction::OperandRefToValue {
   }
 };
 
+struct SILInstruction::OperandToUser {
+  SILInstruction *operator()(const Operand *use) const {
+    return const_cast<Operand *>(use)->getUser();
+  }
+};
+
 struct SILInstruction::FilterOperandToRealOperand {
   const SILInstruction &i;
 

lib/SIL/Utils/OwnershipUtils.cpp

@@ -975,8 +975,9 @@ computeTransitiveLiveness(MultiDefPrunedLiveness &liveness) const {
   });
 }
 
-bool BorrowedValue::areUsesWithinExtendedScope(
-    ArrayRef<Operand *> uses, DeadEndBlocks *deadEndBlocks) const {
+template <typename Instructions>
+bool BorrowedValue::areWithinExtendedScope(Instructions insts,
+                                           DeadEndBlocks *deadEndBlocks) const {
   // First make sure that we actually have a local scope. If we have a non-local
   // scope, then we have something (like a SILFunctionArgument) where a larger
   // semantic construct (in the case of SILFunctionArgument, the function
@@ -988,7 +989,20 @@ bool BorrowedValue::areUsesWithinExtendedScope(
   // Compute the local scope's liveness.
   MultiDefPrunedLiveness liveness(value->getFunction());
   computeTransitiveLiveness(liveness);
-  return liveness.areUsesWithinBoundary(uses, deadEndBlocks);
+  return liveness.areWithinBoundary(insts, deadEndBlocks);
+}
+
+template bool BorrowedValue::areWithinExtendedScope<UsePointInstructionRange>(
+    UsePointInstructionRange insts, DeadEndBlocks *deadEndBlocks) const;
+
+template bool
+BorrowedValue::areWithinExtendedScope<SILInstruction::OperandUserRange>(
+    SILInstruction::OperandUserRange insts, DeadEndBlocks *deadEndBlocks) const;
+
+bool BorrowedValue::areUsesWithinExtendedScope(
+    ArrayRef<Operand *> uses, DeadEndBlocks *deadEndBlocks) const {
+  SILInstruction::OperandUserRange users(uses, SILInstruction::OperandToUser());
+  return areWithinExtendedScope(users, deadEndBlocks);
 }
 
 // The visitor \p func is only called on final scope-ending uses, not reborrows.

lib/SIL/Utils/PrunedLiveness.cpp

@@ -796,14 +796,8 @@ static FunctionTest SSAPrunedLiveness__areUsesWithinBoundary(
 template <typename LivenessWithDefs>
 bool PrunedLiveRange<LivenessWithDefs>::areUsesWithinBoundary(
     ArrayRef<Operand *> uses, DeadEndBlocks *deadEndBlocks) const {
-  assert(asImpl().isInitialized());
-
-  for (auto *use : uses) {
-    auto *user = use->getUser();
-    if (!isWithinBoundary(user, deadEndBlocks))
-      return false;
-  }
-  return true;
+  SILInstruction::OperandUserRange users(uses, SILInstruction::OperandToUser());
+  return areWithinBoundary(users, deadEndBlocks);
 }
 
 template <typename LivenessWithDefs>

lib/SILOptimizer/SemanticARC/CopyValueOpts.cpp

@@ -27,6 +27,7 @@
 #include "swift/SIL/MemAccessUtils.h"
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/Projection.h"
+#include "swift/SIL/Test.h"
 
 using namespace swift;
 using namespace swift::semanticarc;
@@ -170,7 +171,7 @@ bool SemanticARCOptVisitor::performGuaranteedCopyValueOptimization(
         return borrowScope.isLocalScope();
       });
 
-  auto destroys = lr.getDestroyingUses();
+  auto destroys = lr.getDestroyingInsts();
   if (destroys.empty() && haveAnyLocalScopes) {
     return false;
   }
@@ -196,8 +197,8 @@ bool SemanticARCOptVisitor::performGuaranteedCopyValueOptimization(
   // block.
   {
     if (llvm::any_of(borrowScopeIntroducers, [&](BorrowedValue borrowScope) {
-          return !borrowScope.areUsesWithinExtendedScope(
-              lr.getAllConsumingUses(), nullptr);
+          return !borrowScope.areWithinExtendedScope(lr.getAllConsumingInsts(),
+                                                     nullptr);
         })) {
       LLVM_DEBUG(llvm::dbgs() << "copy_value is extending borrow introducer "
                                  "lifetime, bailing out\n");
@@ -238,8 +239,8 @@ bool SemanticARCOptVisitor::performGuaranteedCopyValueOptimization(
       }
 
       if (llvm::any_of(borrowScopeIntroducers, [&](BorrowedValue borrowScope) {
-            return !borrowScope.areUsesWithinExtendedScope(
-                phiArgLR.getAllConsumingUses(), nullptr);
+            return !borrowScope.areWithinExtendedScope(
+                phiArgLR.getAllConsumingInsts(), nullptr);
           })) {
         return false;
       }
@@ -828,6 +829,20 @@ bool SemanticARCOptVisitor::tryPerformOwnedCopyValueOptimization(
   return true;
 }
 
+namespace swift::test {
+static FunctionTest SemanticARCOptsCopyValueOptsGuaranteedValueOptTest(
+    "semantic_arc_opts__copy_value_opts__guaranteed_value_opt",
+    [](auto &function, auto &arguments, auto &test) {
+      SemanticARCOptVisitor visitor(function, test.getPassManager(),
+                                    *test.getDeadEndBlocks(),
+                                    /*onlyMandatoryOpts=*/false);
+
+      visitor.performGuaranteedCopyValueOptimization(
+          cast<CopyValueInst>(arguments.takeInstruction()));
+      function.print(llvm::errs());
+    });
+} // end namespace swift::test
+
 //===----------------------------------------------------------------------===//
 //                            Top Level Entrypoint
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/SemanticARC/OwnedToGuaranteedPhiOpt.cpp

@@ -212,7 +212,11 @@ bool swift::semanticarc::tryConvertOwnedPhisToGuaranteedPhis(Context &ctx) {
     SmallVector<std::pair<SILValue, unsigned>, 8> incomingValueUpdates;
     for (auto introducer : ownedValueIntroducers) {
       SILValue v = introducer.value;
-      OwnershipLiveRange lr(v);
+      SmallVector<std::pair<Operand *, SILValue>, 4> extraConsumes;
+      for (auto op : incomingValueOperandList) {
+        extraConsumes.push_back({op.getOperand(), op.getOriginal()});
+      }
+      OwnershipLiveRange lr(v, extraConsumes);
 
       // For now, we only handle copy_value for simplicity.
       //

lib/SILOptimizer/SemanticARC/OwnershipLiveRange.cpp

@@ -15,20 +15,25 @@
 
 #include "swift/Basic/Assertions.h"
 #include "swift/SIL/BasicBlockUtils.h"
+#include "swift/SIL/OSSALifetimeCompletion.h"
 #include "swift/SIL/OwnershipUtils.h"
 
 using namespace swift;
 using namespace swift::semanticarc;
 
-OwnershipLiveRange::OwnershipLiveRange(SILValue value)
+OwnershipLiveRange::OwnershipLiveRange(
+    SILValue value, ArrayRef<std::pair<Operand *, SILValue>> extraUses)
     : introducer(OwnedValueIntroducer::get(value)), destroyingUses(),
       ownershipForwardingUses(), unknownConsumingUses() {
   assert(introducer);
   assert(introducer.value->getOwnershipKind() == OwnershipKind::Owned);
 
-  SmallVector<Operand *, 32> tmpDestroyingUses;
-  SmallVector<Operand *, 32> tmpForwardingConsumingUses;
-  SmallVector<Operand *, 32> tmpUnknownConsumingUses;
+  SmallVector<UsePoint, 32> tmpDestroyingUses;
+  SmallVector<UsePoint, 32> tmpForwardingConsumingUses;
+  SmallVector<UsePoint, 32> tmpUnknownConsumingUses;
+
+  SmallVector<SILValue, 4> defs;
+  defs.push_back(value);
 
   // We know that our silvalue produces an @owned value. Look through all of our
   // uses and classify them as either consuming or not.
@@ -111,6 +116,7 @@ OwnershipLiveRange::OwnershipLiveRange(SILValue value)
           continue;
         }
         llvm::copy(v->getUses(), std::back_inserter(worklist));
+        defs.push_back(v);
       }
       continue;
     }
@@ -138,8 +144,32 @@ OwnershipLiveRange::OwnershipLiveRange(SILValue value)
         // Otherwise add all users of this BBArg to the worklist to visit
         // recursively.
         llvm::copy(succArg->getUses(), std::back_inserter(worklist));
+        defs.push_back(succArg);
+      }
+    }
+  }
+
+  for (auto def : defs) {
+    if (def->use_begin() == def->use_end()) {
+      continue;
+    }
+    SmallVector<SILBasicBlock *, 32> blocks;
+    SSAPrunedLiveness liveness(def->getFunction(), &blocks);
+    liveness.initializeDef(def);
+    liveness.computeSimple();
+    for (auto use : extraUses) {
+      if (use.second != def) {
+        continue;
       }
+      liveness.updateForUse(use.first->getUser(), /*lifetimeEnding=*/true);
     }
+    OSSALifetimeCompletion::visitAvailabilityBoundary(
+        def, liveness, [&tmpDestroyingUses](auto *inst, auto end) {
+          if (end != OSSALifetimeCompletion::LifetimeEnd::Boundary) {
+            return;
+          }
+          tmpDestroyingUses.push_back(inst);
+        });
   }
 
   // The order in which we append these to consumingUses matters since we assume
@@ -248,7 +278,8 @@ void OwnershipLiveRange::insertEndBorrowsAtDestroys(
 
 void OwnershipLiveRange::convertOwnedGeneralForwardingUsesToGuaranteed() && {
   while (!ownershipForwardingUses.empty()) {
-    auto *use = ownershipForwardingUses.back();
+    auto point = ownershipForwardingUses.back();
+    auto *use = point.getOperand();
     ownershipForwardingUses = ownershipForwardingUses.drop_back();
     ForwardingOperand operand(use);
     operand.replaceOwnershipKind(OwnershipKind::Owned,
@@ -260,9 +291,12 @@ void OwnershipLiveRange::convertToGuaranteedAndRAUW(
     SILValue newGuaranteedValue, InstModCallbacks callbacks) && {
   auto *value = cast<SingleValueInstruction>(introducer.value);
   while (!destroyingUses.empty()) {
-    auto *d = destroyingUses.back();
+    auto point = destroyingUses.back();
+    auto *destroy = point.getInstruction();
     destroyingUses = destroyingUses.drop_back();
-    callbacks.deleteInst(d->getUser());
+    if (isa<TermInst>(destroy))
+      continue;
+    callbacks.deleteInst(destroy);
   }
 
   callbacks.eraseAndRAUWSingleValueInst(value, newGuaranteedValue);
@@ -318,9 +352,12 @@ void OwnershipLiveRange::convertJoinedLiveRangePhiToGuaranteed(
 
   // Then eliminate all of the destroys...
   while (!destroyingUses.empty()) {
-    auto *d = destroyingUses.back();
+    auto point = destroyingUses.back();
+    auto *destroy = point.getInstruction();
     destroyingUses = destroyingUses.drop_back();
-    callbacks.deleteInst(d->getUser());
+    if (isa<TermInst>(destroy))
+      continue;
+    callbacks.deleteInst(destroy);
   }
 
   // and change all of our guaranteed forwarding insts to have guaranteed
@@ -362,13 +399,3 @@ OwnershipLiveRange::hasUnknownConsumingUse(bool assumingAtFixPoint) const {
   // to our introducer.
   return HasConsumingUse_t::YesButAllPhiArgs;
 }
-
-OwnershipLiveRange::DestroyingInstsRange
-OwnershipLiveRange::getDestroyingInsts() const {
-  return DestroyingInstsRange(getDestroyingUses(), OperandToUser());
-}
-
-OwnershipLiveRange::ConsumingInstsRange
-OwnershipLiveRange::getAllConsumingInsts() const {
-  return ConsumingInstsRange(consumingUses, OperandToUser());
-}