Merge pull request #83907 from nate-chandler/rdar158149082

[AllocBoxToStack] Don't destroy in dead-ends.

Author: nate-chandler

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/AllocBoxToStack.swift

@@ -362,9 +362,34 @@ private func createAllocStack(for allocBox: AllocBoxInst, flags: Flags, _ contex
   for destroy in getFinalDestroys(of: allocBox, context) {
     let loc = allocBox.location.asCleanup.withScope(of: destroy.location)
     Builder.insert(after: destroy, location: loc, context) { builder in
-      if  !unboxedType.isTrivial(in: allocBox.parentFunction), !(destroy is DeallocBoxInst) {
+      if !(destroy is DeallocBoxInst),
+         context.deadEndBlocks.isDeadEnd(destroy.parentBlock),
+         !isInLoop(block: destroy.parentBlock, context) {
+        // "Last" releases in dead-end regions may not actually destroy the box
+        // and consequently may not actually release the stored value.  That's
+        // because values (including boxes) may be leaked along paths into
+        // dead-end regions.  Thus it is invalid to lower such final releases of
+        // the box to destroy_addr's/dealloc_box's of the stack-promoted storage.
+        //
+        // There is one exception: if the alloc_box is in a dead-end loop.  In
+        // that case SIL invariants require that the final releases actually
+        // destroy the box; otherwise, a box would leak once per loop.  To check
+        // for this, it is sufficient check that the LastRelease is in a dead-end
+        // loop: if the alloc_box is not in that loop, then the entire loop is in
+        // the live range, so no release within the loop would be a "final
+        // release".
+        //
+        // None of this applies to dealloc_box instructions which always destroy
+        // the box.
+        return
+      }
+      if !unboxedType.isTrivial(in: allocBox.parentFunction), !(destroy is DeallocBoxInst) {
         builder.createDestroyAddr(address: stackLocation)
       }
+      if let dbi = destroy as? DeallocBoxInst, dbi.isDeadEnd {
+        // Don't bother to create dealloc_stack instructions in dead-ends.
+        return
+      }
       builder.createDeallocStack(asi)
     }
   }

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/StackPromotion.swift

@@ -330,17 +330,3 @@ private extension BasicBlockRange {
     exits.contains { !deadEndBlocks.isDeadEnd($0) && !$0.hasSinglePredecessor }
   }
 }
-
-private func isInLoop(block startBlock: BasicBlock, _ context: FunctionPassContext) -> Bool {
-  var worklist = BasicBlockWorklist(context)
-  defer { worklist.deinitialize() }
-
-  worklist.pushIfNotVisited(contentsOf: startBlock.successors)
-  while let block = worklist.pop() {
-    if block == startBlock {
-      return true
-    }
-    worklist.pushIfNotVisited(contentsOf: block.successors)
-  }
-  return false
-}

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -1012,3 +1012,17 @@ func eraseIfDead(functions: [Function], _ context: ModulePassContext) {
     toDelete = remaining
   }
 }
+
+func isInLoop(block startBlock: BasicBlock, _ context: FunctionPassContext) -> Bool {
+  var worklist = BasicBlockWorklist(context)
+  defer { worklist.deinitialize() }
+
+  worklist.pushIfNotVisited(contentsOf: startBlock.successors)
+  while let block = worklist.pop() {
+    if block == startBlock {
+      return true
+    }
+    worklist.pushIfNotVisited(contentsOf: block.successors)
+  }
+  return false
+}

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -699,7 +699,9 @@ final public class DeallocRefInst : Instruction, UnaryInstruction, Deallocation
 
 final public class DeallocPartialRefInst : Instruction, Deallocation {}
 
-final public class DeallocBoxInst : Instruction, UnaryInstruction, Deallocation {}
+final public class DeallocBoxInst : Instruction, UnaryInstruction, Deallocation {
+  public var isDeadEnd: Bool { bridged.DeallocBoxInst_isDeadEnd() }
+}
 
 final public class DeallocExistentialBoxInst : Instruction, UnaryInstruction, Deallocation {}
 

include/swift/SIL/SILBridging.h

@@ -828,6 +828,7 @@ struct BridgedInstruction {
   BRIDGED_INLINE bool CopyAddrInst_isInitializationOfDest() const;
   BRIDGED_INLINE void CopyAddrInst_setIsTakeOfSrc(bool isTakeOfSrc) const;
   BRIDGED_INLINE void CopyAddrInst_setIsInitializationOfDest(bool isInitializationOfDest) const;
+  BRIDGED_INLINE bool DeallocBoxInst_isDeadEnd() const;
   BRIDGED_INLINE bool ExplicitCopyAddrInst_isTakeOfSrc() const;
   BRIDGED_INLINE bool ExplicitCopyAddrInst_isInitializationOfDest() const;
   BRIDGED_INLINE SwiftInt MarkUninitializedInst_getKind() const;

include/swift/SIL/SILBridgingImpl.h

@@ -1511,6 +1511,10 @@ void BridgedInstruction::CopyAddrInst_setIsInitializationOfDest(bool isInitializ
       isInitializationOfDest ? swift::IsInitialization : swift::IsNotInitialization);
 }
 
+bool BridgedInstruction::DeallocBoxInst_isDeadEnd() const {
+  return getAs<swift::DeallocBoxInst>()->isDeadEnd();
+}
+
 bool BridgedInstruction::ExplicitCopyAddrInst_isTakeOfSrc() const {
   return getAs<swift::ExplicitCopyAddrInst>()->isTakeOfSrc();
 }

include/swift/SIL/SILFunction.h

@@ -1697,7 +1697,8 @@ class SILFunction
   }
 
   /// Verifies the lifetime of memory locations in the function.
-  void verifyMemoryLifetime(CalleeCache *calleeCache);
+  void verifyMemoryLifetime(CalleeCache *calleeCache,
+                            DeadEndBlocks *deadEndBlocks);
 
   /// Verifies ownership of the function.
   /// Since we don't have complete lifetimes everywhere, computes DeadEndBlocks

lib/SIL/Verifier/MemoryLifetimeVerifier.cpp

@@ -12,13 +12,14 @@
 
 #define DEBUG_TYPE "sil-memory-lifetime-verifier"
 #include "swift/Basic/Assertions.h"
-#include "swift/SIL/MemoryLocations.h"
+#include "swift/SIL/ApplySite.h"
+#include "swift/SIL/BasicBlockDatastructures.h"
+#include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/BitDataflow.h"
 #include "swift/SIL/CalleeCache.h"
+#include "swift/SIL/MemoryLocations.h"
 #include "swift/SIL/SILBasicBlock.h"
 #include "swift/SIL/SILFunction.h"
-#include "swift/SIL/ApplySite.h"
-#include "swift/SIL/BasicBlockDatastructures.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace swift;
@@ -43,6 +44,7 @@ class MemoryLifetimeVerifier {
 
   SILFunction *function;
   CalleeCache *calleeCache;
+  DeadEndBlocks *deadEndBlocks;
   MemoryLocations locations;
 
   /// alloc_stack memory locations which are used for store_borrow.
@@ -140,11 +142,12 @@ class MemoryLifetimeVerifier {
   }
 
 public:
-  MemoryLifetimeVerifier(SILFunction *function, CalleeCache *calleeCache) :
-    function(function),
-    calleeCache(calleeCache),
-    locations(/*handleNonTrivialProjections*/ true,
-              /*handleTrivialLocations*/ true) {}
+  MemoryLifetimeVerifier(SILFunction *function, CalleeCache *calleeCache,
+                         DeadEndBlocks *deadEndBlocks)
+      : function(function), calleeCache(calleeCache),
+        deadEndBlocks(deadEndBlocks),
+        locations(/*handleNonTrivialProjections*/ true,
+                  /*handleTrivialLocations*/ true) {}
 
   /// The main entry point to verify the lifetime of all memory locations in
   /// the function.
@@ -884,7 +887,12 @@ void MemoryLifetimeVerifier::checkBlock(SILBasicBlock *block, Bits &bits) {
       }
       case SILInstructionKind::DeallocStackInst: {
         SILValue opVal = cast<DeallocStackInst>(&I)->getOperand();
-        requireBitsClear(bits & nonTrivialLocations, opVal, &I);
+        if (!deadEndBlocks->isDeadEnd(I.getParent())) {
+          // TODO: rdar://159311784: Maybe at some point the invariant will be
+          //                         enforced that values stored into addresses
+          //                         don't leak in dead-ends.
+          requireBitsClear(bits & nonTrivialLocations, opVal, &I);
+        }
         // Needed to clear any bits of trivial locations (which are not required
         // to be zero).
         locations.clearBits(bits, opVal);
@@ -974,7 +982,8 @@ void MemoryLifetimeVerifier::verify() {
 
 } // anonymous namespace
 
-void SILFunction::verifyMemoryLifetime(CalleeCache *calleeCache) {
-  MemoryLifetimeVerifier verifier(this, calleeCache);
+void SILFunction::verifyMemoryLifetime(CalleeCache *calleeCache,
+                                       DeadEndBlocks *deadEndBlocks) {
+  MemoryLifetimeVerifier verifier(this, calleeCache, deadEndBlocks);
   verifier.verify();
 }

lib/SIL/Verifier/SILVerifier.cpp

@@ -7385,7 +7385,7 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
 
     if (F->hasOwnership() && F->shouldVerifyOwnership() &&
         !mod.getASTContext().hadError()) {
-      F->verifyMemoryLifetime(calleeCache);
+      F->verifyMemoryLifetime(calleeCache, &getDeadEndBlocks());
     }
   }
 

lib/SILOptimizer/Transforms/AllocBoxToStack.cpp

@@ -25,6 +25,8 @@
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILCloner.h"
+#include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
+#include "swift/SILOptimizer/Analysis/LoopAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
@@ -601,7 +603,9 @@ static void hoistMarkUnresolvedNonCopyableValueInsts(
 
 /// rewriteAllocBoxAsAllocStack - Replace uses of the alloc_box with a
 /// new alloc_stack, but do not delete the alloc_box yet.
-static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI) {
+static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI,
+                                        DeadEndBlocksAnalysis &deba,
+                                        SILLoopAnalysis &la) {
   LLVM_DEBUG(llvm::dbgs() << "*** Promoting alloc_box to stack: " << *ABI);
 
   SILValue HeapBox = ABI;
@@ -693,9 +697,31 @@ static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI) {
                          ABI->getBoxType(), ABI->getModule().Types, 0));
   auto Loc = CleanupLocation(ABI->getLoc());
 
+  auto *deb = deba.get(ABI->getFunction());
   for (auto LastRelease : FinalReleases) {
+    auto *dbi = dyn_cast<DeallocBoxInst>(LastRelease);
+    if (!dbi && deb->isDeadEnd(LastRelease->getParent()) &&
+        !la.get(ABI->getFunction())->getLoopFor(LastRelease->getParent())) {
+      // "Last" releases in dead-end regions may not actually destroy the box
+      // and consequently may not actually release the stored value.  That's
+      // because values (including boxes) may be leaked along paths into
+      // dead-end regions.  Thus it is invalid to lower such final releases of
+      // the box to destroy_addr's/dealloc_box's of the stack-promoted storage.
+      //
+      // There is one exception: if the alloc_box is in a dead-end loop.  In
+      // that case SIL invariants require that the final releases actually
+      // destroy the box; otherwise, a box would leak once per loop.  To check
+      // for this, it is sufficient check that the LastRelease is in a dead-end
+      // loop: if the alloc_box is not in that loop, then the entire loop is in
+      // the live range, so no release within the loop would be a "final
+      // release".
+      //
+      // None of this applies to dealloc_box instructions which always destroy
+      // the box.
+      continue;
+    }
     SILBuilderWithScope Builder(LastRelease);
-    if (!isa<DeallocBoxInst>(LastRelease)&& !Lowering.isTrivial()) {
+    if (!dbi && !Lowering.isTrivial()) {
       // If we have a mark_unresolved_non_copyable_value use of our stack box,
       // we want to destroy that.
       SILValue valueToDestroy = StackBox;
@@ -709,7 +735,6 @@ static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI) {
       // instruction we found that isn't an explicit dealloc_box.
       Builder.emitDestroyAddrAndFold(Loc, valueToDestroy);
     }
-    auto *dbi = dyn_cast<DeallocBoxInst>(LastRelease);
     if (dbi && dbi->isDeadEnd()) {
       // Don't bother to create dealloc_stack instructions in dead-ends.
       continue;
@@ -1265,7 +1290,9 @@ static void rewriteApplySites(AllocBoxToStackState &pass) {
 
 /// Clone closure bodies and rewrite partial applies. Returns the number of
 /// alloc_box allocations promoted.
-static unsigned rewritePromotedBoxes(AllocBoxToStackState &pass) {
+static unsigned rewritePromotedBoxes(AllocBoxToStackState &pass,
+                                     DeadEndBlocksAnalysis &deba,
+                                     SILLoopAnalysis &la) {
   // First we'll rewrite any ApplySite that we can to remove
   // the box container pointer from the operands.
   rewriteApplySites(pass);
@@ -1274,7 +1301,7 @@ static unsigned rewritePromotedBoxes(AllocBoxToStackState &pass) {
   auto rend = pass.Promotable.rend();
   for (auto I = pass.Promotable.rbegin(); I != rend; ++I) {
     auto *ABI = *I;
-    if (rewriteAllocBoxAsAllocStack(ABI)) {
+    if (rewriteAllocBoxAsAllocStack(ABI, deba, la)) {
       ++Count;
       ABI->eraseFromParent();
     }
@@ -1299,7 +1326,9 @@ class AllocBoxToStack : public SILFunctionTransform {
     }
 
     if (!pass.Promotable.empty()) {
-      auto Count = rewritePromotedBoxes(pass);
+      auto *deba = getAnalysis<DeadEndBlocksAnalysis>();
+      auto *la = getAnalysis<SILLoopAnalysis>();
+      auto Count = rewritePromotedBoxes(pass, *deba, *la);
       NumStackPromoted += Count;
       if (Count) {
         if (StackNesting::fixNesting(getFunction()) == StackNesting::Changes::CFG)