Merge pull request #22038 from gottesmm/pr-48a80b46d0806aa50603d26425b66bdb807aecc2

diff --git a/include/swift/SIL/OwnershipUtils.h b/include/swift/SIL/OwnershipUtils.h
index 7a60169..9550260 100644
--- a/include/swift/SIL/OwnershipUtils.h
+++ b/include/swift/SIL/OwnershipUtils.h
@@ -35,8 +35,10 @@
     ReturnFalse = 1,
     PrintMessage = 2,
     Assert = 4,
+    ReturnFalseOnLeak = 8,
     PrintMessageAndReturnFalse = PrintMessage | ReturnFalse,
     PrintMessageAndAssert = PrintMessage | Assert,
+    ReturnFalseOnLeakAssertOtherwise = ReturnFalseOnLeak | Assert,
   } Value;
 
   ErrorBehaviorKind() : Value(Invalid) {}
@@ -47,6 +49,11 @@
     return Value & Assert;
   }
 
+  bool shouldReturnFalseOnLeak() const {
+    assert(Value != Invalid);
+    return Value & ReturnFalseOnLeak;
+  }
+
   bool shouldPrintMessage() const {
     assert(Value != Invalid);
     return Value & PrintMessage;
@@ -95,12 +102,21 @@
 /// non-consuming uses, or from the producer instruction.
 /// 2. The consuming use set jointly post dominates producers and all non
 /// consuming uses.
-bool valueHasLinearLifetime(SILValue value,
-                            ArrayRef<BranchPropagatedUser> consumingUses,
-                            ArrayRef<BranchPropagatedUser> nonConsumingUses,
-                            SmallPtrSetImpl<SILBasicBlock *> &visitedBlocks,
-                            DeadEndBlocks &deadEndBlocks,
-                            ownership::ErrorBehaviorKind errorBehavior);
+///
+/// \p value The value whose lifetime we are checking.
+/// \p consumingUses the array of users that destroy or consume a value.
+/// \p nonConsumingUses regular uses
+/// \p deadEndBlocks a cache for the dead end block computation
+/// \p errorBehavior If we detect an error, should we return false or hard
+/// error.
+/// \p leakingBlocks If non-null a list of blocks where the value was detected
+/// to leak. Can be used to insert missing destroys.
+bool valueHasLinearLifetime(
+    SILValue value, ArrayRef<BranchPropagatedUser> consumingUses,
+    ArrayRef<BranchPropagatedUser> nonConsumingUses,
+    SmallPtrSetImpl<SILBasicBlock *> &visitedBlocks,
+    DeadEndBlocks &deadEndBlocks, ownership::ErrorBehaviorKind errorBehavior,
+    SmallVectorImpl<SILBasicBlock *> *leakingBlocks = nullptr);
 
 /// Returns true if v is an address or trivial.
 bool isValueAddressOrTrivial(SILValue v, SILModule &m);
diff --git a/lib/SIL/LinearLifetimeChecker.cpp b/lib/SIL/LinearLifetimeChecker.cpp
index 858dd2b..ed11ee1 100644
--- a/lib/SIL/LinearLifetimeChecker.cpp
+++ b/lib/SIL/LinearLifetimeChecker.cpp
@@ -53,6 +53,11 @@
   /// The blocks that we have already visited.
   SmallPtrSetImpl<SILBasicBlock *> &visitedBlocks;
 
+  /// If non-null a list that we should place any detected leaking blocks for
+  /// our caller. The intention is that this can be used in a failing case to
+  /// put in missing destroys.
+  SmallVectorImpl<SILBasicBlock *> *leakingBlocks;
+
   /// The set of blocks with consuming uses.
   SmallPtrSet<SILBasicBlock *, 8> blocksWithConsumingUses;
 
@@ -66,12 +71,13 @@
 
   /// A list of successor blocks that we must visit by the time the algorithm
   /// terminates.
-  SmallPtrSet<SILBasicBlock *, 8> successorBlocksThatMustBeVisited;
+  SmallSetVector<SILBasicBlock *, 8> successorBlocksThatMustBeVisited;
 
   State(SILValue value, SmallPtrSetImpl<SILBasicBlock *> &visitedBlocks,
-        ErrorBehaviorKind errorBehavior)
+        ErrorBehaviorKind errorBehavior,
+        SmallVectorImpl<SILBasicBlock *> *leakingBlocks)
       : value(value), errorBehavior(errorBehavior),
-        visitedBlocks(visitedBlocks) {}
+        visitedBlocks(visitedBlocks), leakingBlocks(leakingBlocks) {}
 
   void initializeAllNonConsumingUses(
       ArrayRef<BranchPropagatedUser> nonConsumingUsers);
@@ -331,7 +337,7 @@
     // First remove BB from the SuccessorBlocksThatMustBeVisited list. This
     // ensures that when the algorithm terminates, we know that BB was not the
     // beginning of a non-covered path to the exit.
-    successorBlocksThatMustBeVisited.erase(block);
+    successorBlocksThatMustBeVisited.remove(block);
 
     // Then remove BB from BlocksWithNonLifetimeEndingUses so we know that
     // this block was properly joint post-dominated by our lifetime ending
@@ -395,23 +401,41 @@
 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;
+
   if (!successorBlocksThatMustBeVisited.empty()) {
-    return handleError([&] {
-      llvm::errs()
-          << "Function: '" << value->getFunction()->getName() << "'\n"
-          << "Error! Found a leak due to a consuming post-dominance failure!\n"
-          << "    Value: " << *value << "    Post Dominating Failure Blocks:\n";
-      for (auto *succBlock : successorBlocksThatMustBeVisited) {
-        llvm::errs() << "        bb" << succBlock->getDebugID();
-      }
-      llvm::errs() << '\n';
-    });
+    // If we are asked to store any leaking blocks, put them in the leaking
+    // blocks array.
+    if (leakingBlocks) {
+      copy(successorBlocksThatMustBeVisited,
+           std::back_inserter(*leakingBlocks));
+    }
+
+    // If we are supposed to error on leaks, do so now.
+    if (!errorBehavior.shouldReturnFalseOnLeak()) {
+      return handleError([&] {
+        llvm::errs() << "Function: '" << value->getFunction()->getName()
+                     << "'\n"
+                     << "Error! Found a leak due to a consuming post-dominance "
+                        "failure!\n"
+                     << "    Value: " << *value
+                     << "    Post Dominating Failure Blocks:\n";
+        for (auto *succBlock : successorBlocksThatMustBeVisited) {
+          llvm::errs() << "        bb" << succBlock->getDebugID();
+        }
+        llvm::errs() << '\n';
+      });
+    }
+
+    // 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 true;
+    return doesntHaveAnyLeaks;
   }
 
   // If we do have remaining blocks, then these non lifetime ending uses must be
@@ -436,7 +460,7 @@
 
   // If all of our remaining blocks were dead uses, then return true. We are
   // good.
-  return true;
+  return doesntHaveAnyLeaks;
 }
 
 //===----------------------------------------------------------------------===//
@@ -447,10 +471,11 @@
     SILValue value, ArrayRef<BranchPropagatedUser> consumingUses,
     ArrayRef<BranchPropagatedUser> nonConsumingUses,
     SmallPtrSetImpl<SILBasicBlock *> &visitedBlocks, DeadEndBlocks &deBlocks,
-    ErrorBehaviorKind errorBehavior) {
+    ErrorBehaviorKind errorBehavior,
+    SmallVectorImpl<SILBasicBlock *> *leakingBlocks) {
   assert(!consumingUses.empty() && "Must have at least one consuming user?!");
 
-  State state(value, visitedBlocks, errorBehavior);
+  State state(value, visitedBlocks, errorBehavior, leakingBlocks);
 
   // First add our non-consuming uses and their blocks to the
   // blocksWithNonConsumingUses map. While we do this, if we have multiple uses