[semantic-arc] When computing Lifetimes include all consuming uses, not just the final destroying use.

include/swift/SILOptimizer/Utils/ValueLifetime.h

@@ -28,20 +28,38 @@ namespace swift {
 /// of the analysis and can be a super set of the uses of the SILValue
 /// e.g. it can be the set of transitive uses of the SILValue.
 class ValueLifetimeAnalysis {
+  /// The instruction or argument that define the value.
+  PointerUnion<SILInstruction *, SILArgument *> defValue;
+
+  /// The set of blocks where the value is live.
+  llvm::SmallSetVector<SILBasicBlock *, 16> liveBlocks;
+
+  /// The set of instructions where the value is used, or the users-list
+  /// provided with the constructor.
+  llvm::SmallPtrSet<SILInstruction *, 16> userSet;
+
+  /// Indicates whether the basic block containing def has users of def that
+  /// precede def. This field is initialized by propagateLiveness.
+  bool hasUsersBeforeDef;
+
+  /// Critical edges that couldn't be split to compute the frontier. This could
+  /// be non-empty when the analysis is invoked with DontModifyCFG mode.
+  llvm::SmallVector<std::pair<TermInst *, unsigned>, 16> criticalEdges;
+
 public:
 
   /// The lifetime frontier for the value. It is the list of instructions
   /// following the last uses of the value. All the frontier instructions
   /// end the value's lifetime.
   using Frontier = SmallVector<SILInstruction *, 4>;
 
-  /// Constructor for the value \p Def with a specific range of users.
+  /// Constructor for the value \p def with a specific range of users.
   ///
   /// We templatize over the RangeTy so that we can initialize
   /// ValueLifetimeAnalysis with misc iterators including transform
   /// iterators.
   template <typename RangeTy>
-  ValueLifetimeAnalysis(SILInstruction *def, const RangeTy &userRange)
+  ValueLifetimeAnalysis(decltype(defValue) def, const RangeTy &userRange)
       : defValue(def), userSet(userRange.begin(), userRange.end()) {
     propagateLiveness();
   }
@@ -102,7 +120,7 @@ class ValueLifetimeAnalysis {
   /// Returns true if the value is alive at the begin of block \p bb.
   bool isAliveAtBeginOfBlock(SILBasicBlock *bb) {
     return liveBlocks.count(bb) &&
-           (bb != defValue->getParent() || hasUsersBeforeDef);
+           (hasUsersBeforeDef || bb != getDefValueParentBlock());
   }
 
   /// Checks if there is a dealloc_ref inside the value's live range.
@@ -112,24 +130,19 @@ class ValueLifetimeAnalysis {
   void dump() const;
 
 private:
+  SILFunction *getFunction() const {
+    if (auto *inst = defValue.dyn_cast<SILInstruction *>()) {
+      return inst->getFunction();
+    }
+    return defValue.get<SILArgument *>()->getFunction();
+  }
 
-  /// The value.
-  SILInstruction *defValue;
-
-  /// The set of blocks where the value is live.
-  llvm::SmallSetVector<SILBasicBlock *, 16> liveBlocks;
-
-  /// The set of instructions where the value is used, or the users-list
-  /// provided with the constructor.
-  llvm::SmallPtrSet<SILInstruction *, 16> userSet;
-
-  /// Indicates whether the basic block containing def has users of def that
-  /// precede def. This field is initialized by propagateLiveness.
-  bool hasUsersBeforeDef;
-
-  /// Critical edges that couldn't be split to compute the frontier. This could
-  /// be non-empty when the analysis is invoked with DontModifyCFG mode.
-  llvm::SmallVector<std::pair<TermInst *, unsigned>, 16> criticalEdges;
+  SILBasicBlock *getDefValueParentBlock() const {
+    if (auto *inst = defValue.dyn_cast<SILInstruction *>()) {
+      return inst->getParent();
+    }
+    return defValue.get<SILArgument *>()->getParent();
+  }
 
   /// Propagates the liveness information up the control flow graph.
   void propagateLiveness();

lib/SILOptimizer/Transforms/SemanticARCOpts.cpp

@@ -245,6 +245,10 @@ class OwnershipLiveRange {
       TransformRange<ArrayRef<Operand *>, OperandToUser>;
   DestroyingInstsRange getDestroyingInsts() const;
 
+  using ConsumingInstsRange =
+      TransformRange<ArrayRef<Operand *>, OperandToUser>;
+  ConsumingInstsRange getAllConsumingInsts() const;
+
   /// If this LiveRange has a single destroying use, return that use. Otherwise,
   /// return nullptr.
   Operand *getSingleDestroyingUse() const {
@@ -335,6 +339,11 @@ OwnershipLiveRange::getDestroyingInsts() const {
   return DestroyingInstsRange(getDestroyingUses(), OperandToUser());
 }
 
+OwnershipLiveRange::ConsumingInstsRange
+OwnershipLiveRange::getAllConsumingInsts() const {
+  return ConsumingInstsRange(consumingUses, OperandToUser());
+}
+
 OwnershipLiveRange::OwnershipLiveRange(SILValue value)
     : introducer(*OwnedValueIntroducer::get(value)), destroyingUses(),
       ownershipForwardingUses(), unknownConsumingUses() {
@@ -464,36 +473,17 @@ void OwnershipLiveRange::insertEndBorrowsAtDestroys(
   //
   // TODO: Hoist this out?
   SILInstruction *inst = introducer.value->getDefiningInstruction();
+  Optional<ValueLifetimeAnalysis> analysis;
   if (!inst) {
-    // If our introducer was not for an inst, it should be from an arg. In such
-    // a case, we handle one of two cases:
-    //
-    // 1. If we have one destroy and that destroy is the initial instruction in
-    // the arguments block, we just insert the end_borrow here before the
-    // destroy_value and bail. If the destroy is not the initial instruction in
-    // the arg block, we delegate to the ValueLifetimeAnalysis code.
-    //
-    // 2. If we have multiple destroys, by the properties of owned values having
-    // a linear lifetime, we know that the destroys can not both be first in the
-    // args block since the only way that we could have two such destroys in the
-    // arg's block is if we destructured the arg. In such a case, the
-    // destructure instruction would have to be between the argument and any
-    // destroy meaning the destroys could not be first. In such a case, we
-    // delegate to the ValueLifetimeAnalysis code.
-    auto *arg = cast<SILArgument>(introducer.value);
-    auto *beginInst = &*arg->getParent()->begin();
-    if (auto *singleDestroyingUse = getSingleDestroyingUse()) {
-      if (singleDestroyingUse->getUser() == beginInst) {
-        auto loc = RegularLocation::getAutoGeneratedLocation();
-        SILBuilderWithScope builder(beginInst);
-        builder.createEndBorrow(loc, newGuaranteedValue);
-        return;
-      }
-    }
-    inst = beginInst;
+    analysis.emplace(cast<SILArgument>(introducer.value),
+                     getAllConsumingInsts());
+  } else {
+    analysis.emplace(inst, getAllConsumingInsts());
   }
-  ValueLifetimeAnalysis analysis(inst, getDestroyingInsts());
-  bool foundCriticalEdges = !analysis.computeFrontier(
+
+  // Use all consuming uses in our value lifetime analysis to ensure correctness
+  // in the face of unreachable code.
+  bool foundCriticalEdges = !analysis->computeFrontier(
       scratch, ValueLifetimeAnalysis::DontModifyCFG, &deadEndBlocks);
   (void)foundCriticalEdges;
   assert(!foundCriticalEdges);

lib/SILOptimizer/Utils/ValueLifetime.cpp

@@ -11,17 +11,23 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/SILOptimizer/Utils/ValueLifetime.h"
+#include "swift/Basic/STLExtras.h"
 #include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SILOptimizer/Utils/CFGOptUtils.h"
 
 using namespace swift;
 
 void ValueLifetimeAnalysis::propagateLiveness() {
+  bool defIsInstruction = defValue.is<SILInstruction *>();
   assert(liveBlocks.empty() && "frontier computed twice");
-  assert(!userSet.count(defValue) && "definition cannot be its own use");
+  assert(
+      (!defIsInstruction || !userSet.count(defValue.get<SILInstruction *>())) &&
+      "definition cannot be its own use");
 
-  auto defBB = defValue->getParentBlock();
-  llvm::SmallVector<SILBasicBlock *, 64> worklist;
+  // Compute the def block only if we have a SILInstruction. If we have a
+  // SILArgument, this will be nullptr.
+  auto *defBB = getDefValueParentBlock();
+  SmallVector<SILBasicBlock *, 64> worklist;
   int numUsersBeforeDef = 0;
 
   // Find the initial set of blocks where the value is live, because
@@ -31,20 +37,28 @@ void ValueLifetimeAnalysis::propagateLiveness() {
     if (liveBlocks.insert(userBlock))
       worklist.push_back(userBlock);
 
-    // A user in the defBB could potentially be located before the defValue.
-    if (userBlock == defBB)
+    // A user in the defBB could potentially be located before the defValue. If
+    // we had a SILArgument, defBB will be nullptr, so we should always have
+    // numUsersBeforeDef is 0. We assert this at the end of the loop.
+    if (defIsInstruction && userBlock == defBB)
       ++numUsersBeforeDef;
   }
-  // Don't count any users in the defBB which are actually located _after_
-  // the defValue.
-  auto instIter = defValue->getIterator();
-  while (numUsersBeforeDef > 0 && ++instIter != defBB->end()) {
-    if (userSet.count(&*instIter))
-      --numUsersBeforeDef;
+  assert((defValue.is<SILInstruction *>() || (numUsersBeforeDef == 0)) &&
+         "Non SILInstruction defValue with users before the def?!");
+
+  // Don't count any users in the defBB which are actually located _after_ the
+  // defValue.
+  if (defIsInstruction) {
+    auto instIter = defValue.get<SILInstruction *>()->getIterator();
+    while (numUsersBeforeDef > 0 && ++instIter != defBB->end()) {
+      if (userSet.count(&*instIter))
+        --numUsersBeforeDef;
+    }
   }
 
   // Initialize the hasUsersBeforeDef field.
   hasUsersBeforeDef = numUsersBeforeDef > 0;
+  assert(defIsInstruction || !hasUsersBeforeDef);
 
   // Now propagate liveness backwards until we hit the block that defines the
   // value.
@@ -55,30 +69,31 @@ void ValueLifetimeAnalysis::propagateLiveness() {
     if (bb == defBB && !hasUsersBeforeDef)
       continue;
 
-    for (SILBasicBlock *Pred : bb->getPredecessorBlocks()) {
+    for (auto *predBB : bb->getPredecessorBlocks()) {
       // If it's already in the set, then we've already queued and/or
       // processed the predecessors.
-      if (liveBlocks.insert(Pred))
-        worklist.push_back(Pred);
+      if (liveBlocks.insert(predBB))
+        worklist.push_back(predBB);
     }
   }
 }
 
 SILInstruction *ValueLifetimeAnalysis::findLastUserInBlock(SILBasicBlock *bb) {
   // Walk backwards in bb looking for last use of the value.
-  for (auto ii = bb->rbegin(); ii != bb->rend(); ++ii) {
-    assert(defValue != &*ii && "Found def before finding use!");
+  for (auto &inst : llvm::reverse(*bb)) {
+    assert(defValue.dyn_cast<SILInstruction *>() != &inst &&
+           "Found def before finding use!");
 
-    if (userSet.count(&*ii))
-      return &*ii;
+    if (userSet.count(&inst))
+      return &inst;
   }
   llvm_unreachable("Expected to find use of value in block!");
 }
 
 bool ValueLifetimeAnalysis::computeFrontier(Frontier &frontier, Mode mode,
                                             DeadEndBlocks *deBlocks) {
-  assert(!isAliveAtBeginOfBlock(defValue->getFunction()->getEntryBlock())
-         && "Can't compute frontier for def which does not dominate all uses");
+  assert(!isAliveAtBeginOfBlock(getFunction()->getEntryBlock()) &&
+         "Can't compute frontier for def which does not dominate all uses");
 
   bool noCriticalEdges = true;
 
@@ -101,10 +116,16 @@ bool ValueLifetimeAnalysis::computeFrontier(Frontier &frontier, Mode mode,
     for (const SILSuccessor &succ : bb->getSuccessors()) {
       if (isAliveAtBeginOfBlock(succ)) {
         liveInSucc = true;
-        if (succ == defValue->getParent()) {
+        if (succ == getDefValueParentBlock()) {
           // Here, the basic block bb uses the value but also redefines the
           // value inside bb. The new value could be used by the successors
           // of succ and therefore could be live at the end of succ as well.
+          //
+          // This should never happen if we have a SILArgument since the
+          // SILArgument can not have any uses before it in a block.
+          assert(defValue.is<SILInstruction *>() &&
+                 "SILArguments dominate all instructions in their defining "
+                 "blocks");
           usedAndRedefinedInSucc = true;
         }
       } else if (!deBlocks || !deBlocks->isDeadEnd(succ)) {
@@ -115,7 +136,10 @@ bool ValueLifetimeAnalysis::computeFrontier(Frontier &frontier, Mode mode,
       // Here, the basic block bb uses the value and later redefines the value.
       // Therefore, this value's lifetime ends after its last use preceding the
       // re-definition of the value.
-      auto ii = defValue->getReverseIterator();
+      //
+      // We know that we can not have a SILArgument here since the SILArgument
+      // dominates all instructions in the same block.
+      auto ii = defValue.get<SILInstruction *>()->getReverseIterator();
       for (; ii != bb->rend(); ++ii) {
         if (userSet.count(&*ii)) {
           frontier.push_back(&*std::next(ii));
@@ -245,15 +269,19 @@ bool ValueLifetimeAnalysis::isWithinLifetime(SILInstruction *inst) {
 // Searches \p bb backwards from the instruction before \p frontierInst
 // to the beginning of the list and returns true if we find a dealloc_ref
 // /before/ we find \p defValue (the instruction that defines our target value).
-static bool blockContainsDeallocRef(SILBasicBlock *bb, SILInstruction *defValue,
-                                    SILInstruction *frontierInst) {
+static bool
+blockContainsDeallocRef(SILBasicBlock *bb,
+                        PointerUnion<SILInstruction *, SILArgument *> defValue,
+                        SILInstruction *frontierInst) {
   SILBasicBlock::reverse_iterator End = bb->rend();
   SILBasicBlock::reverse_iterator iter = frontierInst->getReverseIterator();
   for (++iter; iter != End; ++iter) {
     SILInstruction *inst = &*iter;
     if (isa<DeallocRefInst>(inst))
       return true;
-    if (inst == defValue)
+    // We know that inst is not a nullptr, so if we have a SILArgument, this
+    // will always fail as we want.
+    if (inst == defValue.dyn_cast<SILInstruction *>())
       return false;
   }
   return false;
@@ -281,9 +309,14 @@ bool ValueLifetimeAnalysis::containsDeallocRef(const Frontier &frontier) {
 }
 
 void ValueLifetimeAnalysis::dump() const {
-  llvm::errs() << "lifetime of def: " << *defValue;
-  for (SILInstruction *Use : userSet) {
-    llvm::errs() << "  use: " << *Use;
+  llvm::errs() << "lifetime of def: ";
+  if (auto *ii = defValue.dyn_cast<SILInstruction *>()) {
+    llvm::errs() << *ii;
+  } else {
+    llvm::errs() << *defValue.get<SILArgument *>();
+  }
+  for (SILInstruction *use : userSet) {
+    llvm::errs() << "  use: " << *use;
   }
   llvm::errs() << "  live blocks:";
   for (SILBasicBlock *bb : liveBlocks) {