[ownership] Fix the linear lifetime checker to emit leak fixups for certain loops it didn't handle correctly before.

Specifically:

bb0:
  br bb1

bb1:
  cond_br ..., bb2, bb3

bb2:
  br bb1

bb3:
  return

What would happen in this case is that we wouldn't revisit bb1 after we found
the double consume to grab the leak (and would early exit as well). So we would
not insert a destroy for the out of loop value causing a leak from the
perspective of the ownership checker. This was due to us early exiting and also
due to us not revisiting bb1 after we went around the backedge from bb2 ->
bb1. Now what we do instead is if we catch a double consume in a block we have
already visited, we check if we haven't visited any of that block's
successors. If we haven't, we add that to the list of successors we need to
visit.

Author: gottesmm

lib/SIL/LinearLifetimeChecker.cpp

@@ -98,19 +98,19 @@ struct State {
                                      SILBasicBlock *userBlock);
 
   /// Once we have marked all of our producing blocks.
-  bool checkPredsForDoubleConsume(BranchPropagatedUser consumingUser,
+  void checkPredsForDoubleConsume(BranchPropagatedUser consumingUser,
                                   SILBasicBlock *userBlock);
-  bool checkPredsForDoubleConsume(SILBasicBlock *userBlock);
+  void checkPredsForDoubleConsume(SILBasicBlock *userBlock);
 
   /// Once we have setup all of our consuming/non-consuming blocks and have
   /// validated that all intra-block dataflow is safe, perform the inter-block
   /// dataflow.
-  bool performDataflow(DeadEndBlocks &deBlocks);
+  void performDataflow(DeadEndBlocks &deBlocks);
 
   /// After we have performed the dataflow, check the end state of our dataflow
   /// for validity. If this is a linear typed value, return true. Return false
   /// otherwise.
-  bool checkDataflowEndState(DeadEndBlocks &deBlocks);
+  void checkDataflowEndState(DeadEndBlocks &deBlocks);
 };
 
 } // end anonymous namespace
@@ -166,8 +166,7 @@ void State::initializeAllNonConsumingUses(
 
 void State::initializeAllConsumingUses(
     ArrayRef<BranchPropagatedUser> consumingUses,
-    SmallVectorImpl<std::pair<BranchPropagatedUser, SILBasicBlock *>>
-        &predsToAddToWorklist) {
+    SmallVectorImpl<BrPropUserAndBlockPair> &predsToAddToWorklist) {
   for (BranchPropagatedUser user : consumingUses) {
     SILBasicBlock *userBlock = user.getParent();
 
@@ -271,41 +270,58 @@ void State::checkForSameBlockUseAfterFree(BranchPropagatedUser consumingUser,
   return;
 }
 
-bool State::checkPredsForDoubleConsume(BranchPropagatedUser consumingUser,
+void State::checkPredsForDoubleConsume(BranchPropagatedUser consumingUser,
                                        SILBasicBlock *userBlock) {
   if (!blocksWithConsumingUses.count(userBlock))
-    return false;
+    return;
+
+  // Check if this is a block that we have already visited. This means that we
+  // had a back edge of some sort. Double check that we haven't missed any
+  // successors.
+  if (visitedBlocks.count(userBlock)) {
+    for (auto *succ : userBlock->getSuccessorBlocks()) {
+      if (!visitedBlocks.count(succ)) {
+        successorBlocksThatMustBeVisited.insert(succ);
+      }
+    }
+  }
 
   error.handleOverConsume([&] {
     llvm::errs() << "Function: '" << value->getFunction()->getName() << "'\n"
                  << "Found over consume?!\n"
                  << "Value: " << *value << "User: " << *consumingUser
                  << "Block: bb" << userBlock->getDebugID() << "\n\n";
   });
-
-  // If we reached this point, then we did not assert, but we did flag an
-  // error. Return true so we continue the walk.
-  return true;
 }
 
-bool State::checkPredsForDoubleConsume(SILBasicBlock *userBlock) {
+void State::checkPredsForDoubleConsume(SILBasicBlock *userBlock) {
   if (!blocksWithConsumingUses.count(userBlock))
-    return false;
+    return;
+
+  // Check if this is a block that we have already visited. This means that we
+  // had a back edge of some sort. Double check that we haven't missed any
+  // successors.
+  if (visitedBlocks.count(userBlock)) {
+    for (auto *succ : userBlock->getSuccessorBlocks()) {
+      if (!visitedBlocks.count(succ)) {
+        successorBlocksThatMustBeVisited.insert(succ);
+      }
+    }
+  }
 
   error.handleOverConsume([&] {
     llvm::errs() << "Function: '" << value->getFunction()->getName() << "'\n"
                  << "Found over consume?!\n"
                  << "Value: " << *value << "Block: bb"
                  << userBlock->getDebugID() << "\n\n";
   });
-  return true;
 }
 
 //===----------------------------------------------------------------------===//
 //                                  Dataflow
 //===----------------------------------------------------------------------===//
 
-bool State::performDataflow(DeadEndBlocks &deBlocks) {
+void State::performDataflow(DeadEndBlocks &deBlocks) {
   LLVM_DEBUG(llvm::dbgs() << "    Beginning to check dataflow constraints\n");
   // Until the worklist is empty...
   while (!worklist.empty()) {
@@ -366,11 +382,7 @@ bool State::performDataflow(DeadEndBlocks &deBlocks) {
     // 2. We add the predecessor to the worklist if we have not visited it yet.
     for (auto *predBlock : block->getPredecessorBlocks()) {
       // Check if we have an over consume.
-      if (checkPredsForDoubleConsume(predBlock)) {
-        // If we were to assert, it is handled inside check preds for double
-        // consume. So just return false so that we bail.
-        return false;
-      }
+      checkPredsForDoubleConsume(predBlock);
 
       if (visitedBlocks.count(predBlock)) {
         continue;
@@ -380,15 +392,9 @@ bool State::performDataflow(DeadEndBlocks &deBlocks) {
       worklist.push_back(predBlock);
     }
   }
-
-  return true;
 }
 
-bool State::checkDataflowEndState(DeadEndBlocks &deBlocks) {
-  // Make sure that we visited all successor blocks that we needed to visit to
-  // make sure we didn't leak.
-  bool doesntHaveAnyLeaks = true;
-
+void State::checkDataflowEndState(DeadEndBlocks &deBlocks) {
   if (!successorBlocksThatMustBeVisited.empty()) {
     // If we are asked to store any leaking blocks, put them in the leaking
     // blocks array.
@@ -412,13 +418,12 @@ bool State::checkDataflowEndState(DeadEndBlocks &deBlocks) {
 
     // Otherwise... see if we have any other failures. This signals the user
     // wants us to tell it where to insert compensating destroys.
-    doesntHaveAnyLeaks = false;
   }
 
   // Make sure that we do not have any lifetime ending uses left to visit that
   // are not transitively unreachable blocks.... so return early.
   if (blocksWithNonConsumingUses.empty()) {
-    return doesntHaveAnyLeaks;
+    return;
   }
 
   // If we do have remaining blocks, then these non lifetime ending uses must be
@@ -439,12 +444,7 @@ bool State::checkDataflowEndState(DeadEndBlocks &deBlocks) {
       }
       llvm::errs() << "\n";
     });
-    return false;
   }
-
-  // If all of our remaining blocks were dead uses, then return true. We are
-  // good.
-  return doesntHaveAnyLeaks;
 }
 
 //===----------------------------------------------------------------------===//
@@ -506,19 +506,17 @@ LinearLifetimeError swift::valueHasLinearLifetime(
 
     // Make sure that the predecessor is not in our blocksWithConsumingUses
     // list.
-    if (state.checkPredsForDoubleConsume(user, predBlock)) {
-      return state.error;
-    }
+    state.checkPredsForDoubleConsume(user, predBlock);
 
     if (!state.visitedBlocks.insert(predBlock).second)
       continue;
+
     state.worklist.push_back(predBlock);
   }
 
   // Now that our algorithm is completely prepared, run the
   // dataflow... If we find a failure, return false.
-  if (!state.performDataflow(deBlocks))
-    return state.error;
+  state.performDataflow(deBlocks);
 
   // ...and then check that the end state shows that we have a valid linear
   // typed value.

lib/SILOptimizer/Mandatory/PredictableMemOpt.cpp

@@ -425,7 +425,7 @@ class AvailableValueAggregator {
   /// If as a result of us copying values, we may have unconsumed destroys, find
   /// the appropriate location and place the values there. Only used when
   /// ownership is enabled.
-  void addMissingDestroysForCopiedValues(LoadInst *li);
+  LoadInst *addMissingDestroysForCopiedValues(LoadInst *li, SILValue newVal);
 
   void print(llvm::raw_ostream &os) const;
   void dump() const LLVM_ATTRIBUTE_USED;
@@ -746,51 +746,69 @@ SILValue AvailableValueAggregator::handlePrimitiveValue(SILType loadTy,
   return eltVal;
 }
 
-void AvailableValueAggregator::addMissingDestroysForCopiedValues(LoadInst *li) {
+LoadInst *
+AvailableValueAggregator::addMissingDestroysForCopiedValues(LoadInst *li,
+                                                            SILValue newVal) {
   // If ownership is not enabled... bail. We do not need to do this since we do
   // not need to insert an extra copy unless we have ownership since without
   // ownership stores do not consume.
   if (!B.hasOwnership())
-    return;
+    return li;
 
-  SmallVector<BranchPropagatedUser, 1> consumingUses;
   SmallPtrSet<SILBasicBlock *, 8> visitedBlocks;
   SmallVector<SILBasicBlock *, 8> leakingBlocks;
+  bool foundLoop = false;
+  auto loc = RegularLocation::getAutoGeneratedLocation();
   while (!insertedInsts.empty()) {
     auto *cvi = dyn_cast<CopyValueInst>(insertedInsts.pop_back_val());
     if (!cvi)
       continue;
 
-    // Clear our worklist.
-    consumingUses.clear();
+    // Clear our state.
     visitedBlocks.clear();
     leakingBlocks.clear();
-
     // The linear lifetime checker doesn't care if the passed in load is
     // actually a user of our copy_value. What we care about is that the load is
     // guaranteed to be in the block where we have reformed the tuple in a
     // consuming manner. This means if we add it as the consuming use of the
     // copy, we can find the leaking places if any exist.
-    consumingUses.push_back(li);
-
+    //
     // Then perform the linear lifetime check. If we succeed, continue. We have
     // no further work to do.
-    auto errorKind =
-        ownership::ErrorBehaviorKind::ReturnFalseOnLeakAssertOtherwise;
-    auto error =
-        valueHasLinearLifetime(cvi, consumingUses, {}, visitedBlocks,
-                               deadEndBlocks, errorKind, &leakingBlocks);
+    auto errorKind = ownership::ErrorBehaviorKind::ReturnFalse;
+    auto error = valueHasLinearLifetime(
+        cvi, {li}, {}, visitedBlocks, deadEndBlocks, errorKind, &leakingBlocks);
     if (!error.getFoundError())
       continue;
 
+    // Ok, we found some leaking blocks. Since we are using the linear lifetime
+    // checker with memory, we do not have any guarantees that the store is out
+    // side of a loop and a load is in a loop. In such a case, we want to
+    // replace the load with a copy_value.
+    foundLoop |= error.getFoundOverConsume();
+
     // Ok, we found some leaking blocks. Insert destroys at the
     // beginning of these blocks for our copy_value.
-    auto loc = RegularLocation::getAutoGeneratedLocation();
     for (auto *bb : leakingBlocks) {
       SILBuilderWithScope b(bb->begin());
       b.emitDestroyValueOperation(loc, cvi);
     }
   }
+
+  // If we didn't find a loop, we are done, just return li to get RAUWed.
+  if (!foundLoop)
+    return li;
+
+  // If we found a loop, then we know that our leaking blocks are the exiting
+  // blocks of the loop. Thus we need to change the load inst to a copy_value
+  // instead of deleting it.
+  newVal = SILBuilderWithScope(li).emitCopyValueOperation(loc, newVal);
+  li->replaceAllUsesWith(newVal);
+  SILValue addr = li->getOperand();
+  li->eraseFromParent();
+  if (auto *addrI = addr->getDefiningInstruction())
+    recursivelyDeleteTriviallyDeadInstructions(addrI);
+  return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
@@ -1374,27 +1392,31 @@ bool AllocOptimize::promoteLoad(SILInstruction *Inst) {
   auto *Load = cast<LoadInst>(Inst);
   AvailableValueAggregator Agg(Load, AvailableValues, Uses, deadEndBlocks,
                                false /*isTake*/);
-  SILValue NewVal = Agg.aggregateValues(LoadTy, Load->getOperand(), FirstElt);
+  SILValue newVal = Agg.aggregateValues(LoadTy, Load->getOperand(), FirstElt);
+
+  LLVM_DEBUG(llvm::dbgs() << "  *** Promoting load: " << *Load << "\n");
+  LLVM_DEBUG(llvm::dbgs() << "      To value: " << *newVal << "\n");
 
   // If we inserted any copies, we created the copies at our stores. We know
   // that in our load block, we will reform the aggregate as appropriate at the
   // load implying that the value /must/ be fully consumed. Thus any leaking
   // blocks that we may have can be found by performing a linear lifetime check
   // over all copies that we found using the load as the "consuming uses" (just
   // for the purposes of identifying the consuming block).
-  Agg.addMissingDestroysForCopiedValues(Load);
+  auto *oldLoad = Agg.addMissingDestroysForCopiedValues(Load, newVal);
 
   ++NumLoadPromoted;
-  
-  // Simply replace the load.
-  LLVM_DEBUG(llvm::dbgs() << "  *** Promoting load: " << *Load << "\n");
-  LLVM_DEBUG(llvm::dbgs() << "      To value: " << *NewVal << "\n");
-  
-  Load->replaceAllUsesWith(NewVal);
-  SILValue Addr = Load->getOperand();
-  Load->eraseFromParent();
-  if (auto *AddrI = Addr->getDefiningInstruction())
-    recursivelyDeleteTriviallyDeadInstructions(AddrI);
+
+  // If we are returned the load, eliminate it. Otherwise, it was already
+  // handled for us... so return true.
+  if (!oldLoad)
+    return true;
+
+  oldLoad->replaceAllUsesWith(newVal);
+  SILValue addr = oldLoad->getOperand();
+  oldLoad->eraseFromParent();
+  if (auto *addrI = addr->getDefiningInstruction())
+    recursivelyDeleteTriviallyDeadInstructions(addrI);
   return true;
 }