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

TLDR: This fixes an ownership verifier assert caused by not placing end_borrows
along paths where an enum is provable to have a trivial case. It only happens if
all non-trivial cases in a switch_enum are "dead end blocks" where the program
will end and we leak objects.

The Problem
-----------

The actual bug here only occurs in cases where we have a switch_enum on an enum
with mixed trivial, non-trivial cases and all of the non-trivial payloaded cases
are "dead end blocks". As an example, lets look at a simple switch_enum over an
optional where the .some case is a dead end block and we leak the Klass object
into program termination:

```
%0 = load [copy] %mem : $Klass
switch_enum %0 : $Optional<Klass>, case #Optional.some: bbDeadEnd, case #Optional.none: bbContinue

bbDeadEnd(%0a : @owned $Klass): // %0 is leaked into program end!
  unreachable

bbContinue:
  ... // program continue.
```

In this case, if we were only looking at final destroying uses, we would pass a
def without any uses to the ValueLifetimeChecker causing us to not have a
frontier at all causing us to not insert any end_borrows, yielding:

```
%0 = load_borrow %mem : $Klass
switch_enum %0 : $Optional<Klass>, case #Optional.some: bbDeadEnd, case #Optional.none: bbContinue

bbDeadEnd(%0a : @guaranteed $Klass): // %0 is leaked into program end and
                                     // doesnt need an end_borrow!
  unreachable

bbContinue:
  ... // program continue... we need an end_borrow here though!
```

This then trips the ownership verifier since switch_enum is a transforming
terminator that acts like a forwarding instruction implying we need an
end_borrow on the base value along all non-dead end paths through the program.

Importantly this is not actually a leak of a value or unsafe behavior since the
only time that we enter into unsafe territory is along paths where the enum was
actually trivial. So the load_borrow is actually just loaded the trivial enum
value.

The Fix
-------

In order to work around this, I realized that the right solution is to also
include the forwarding consuming uses (in this case the switch_enum use) when
determining the lifetime and that this solves the problem.

That being said, after I made that change, I noticed that I needed to remove my
previous manner of computing the insertion point to use for arguments when
finding the lifetime using ValueLifetimeAnalysis. Previously since I was using
only the destroying uses I knew that the destroy_value could not be the first
instruction in the block of my argument since I handled that case individually
before using the ValueLifetimeAnalysis. That invariant is no longer true as can
be seen in the case above if %0 was from a SILArgument itself instead of a load
[copy] and we were converting that argument to be a guaranteed argument.

To fix this, I taught ValueLifetimeAnalysis how to handle defs from
Arguments. The key thing is I noticed while reading the code that the analysis
only generally cared about the instruction's parent block. Beyond that, the def
being from an instruction was only needed to determine if a user is earlier in
the same block as the def instruction. Those concerns not apply to SILArgument
which dominate all instructions in the same block, so in this patch, we just
skip those conditional checks when we have a SILArgument. The rest of the code
that uses the parent block is the same for both SILArgument/SILInstructions.

rdar://65244617

Author: gottesmm

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) {