SIL Optimizer: handle dead-end infinite loops in StackNesting

include/swift/SILOptimizer/Utils/StackNesting.h

@@ -77,6 +77,11 @@ class StackNesting {
 
     /// Used in the setup function to walk over the CFG.
     bool visited = false;
+
+    /// True for dead-end blocks, i.e. blocks from which there is no path to
+    /// a function exit, e.g. blocks which end with `unreachable` or an
+    /// infinite loop.
+    bool isDeadEnd = false;
   };
 
   /// Data stored for each stack location (= allocation).

lib/SILOptimizer/Utils/StackNesting.cpp

@@ -82,20 +82,38 @@ bool StackNesting::solve() {
   bool isNested = false;
   BitVector Bits(StackLocs.size());
 
+  StackList<SILBasicBlock *> deadEndWorklist(BlockInfos.getFunction());
+
   // Initialize all bit fields to 1s, expect 0s for the entry block.
   bool initVal = false;
   for (auto bd : BlockInfos) {
     bd.data.AliveStackLocsAtEntry.resize(StackLocs.size(), initVal);
     initVal = true;
+
+    bd.data.isDeadEnd = !bd.block.getTerminator()->isFunctionExiting();
+    if (!bd.data.isDeadEnd)
+      deadEndWorklist.push_back(&bd.block);
+  }
+
+  // Calculate the isDeadEnd block flags.
+  while (!deadEndWorklist.empty()) {
+    SILBasicBlock *b = deadEndWorklist.pop_back_val();
+    for (SILBasicBlock *pred : b->getPredecessorBlocks()) {
+      BlockInfo &bi = BlockInfos[pred];
+      if (bi.isDeadEnd) {
+        bi.isDeadEnd = false;
+        deadEndWorklist.push_back(pred);
+      }
+    }
   }
 
   // First step: do a forward dataflow analysis to get the live stack locations
   // at the block exits.
-  // This is necessary to get the live locations at blocks which end in
-  // unreachable instructions (otherwise the backward data flow would be
-  // sufficient). The special thing about unreachable-blocks is that it's
-  // okay to have alive locations at that point, i.e. locations which are never
-  // dealloced. We cannot get such locations with a purly backward dataflow.
+  // This is necessary to get the live locations at dead-end blocks (otherwise
+  // the backward data flow would be sufficient).
+  // The special thing about dead-end blocks is that it's okay to have alive
+  // locations at that point (e.g. at an `unreachable`) i.e. locations which are
+  // never dealloced. We cannot get such locations with a purly backward dataflow.
   do {
     changed = false;
 
@@ -124,7 +142,7 @@ bool StackNesting::solve() {
   do {
     changed = false;
 
-   for (auto bd : llvm::reverse(BlockInfos)) {
+    for (auto bd : llvm::reverse(BlockInfos)) {
       // Collect the alive-bits (at the block exit) from the successor blocks.
       for (SILBasicBlock *SuccBB : bd.block.getSuccessorBlocks()) {
         bd.data.AliveStackLocsAtExit |= BlockInfos[SuccBB].AliveStackLocsAtEntry;
@@ -134,14 +152,18 @@ bool StackNesting::solve() {
                && Bits.any())
              && "stack location is missing dealloc");
 
-      if (isa<UnreachableInst>(bd.block.getTerminator())) {
-        // We treat unreachable as an implicit deallocation for all locations
-        // which are still alive at this point.
+      if (bd.data.isDeadEnd) {
+        // We treat `unreachable` as an implicit deallocation for all locations
+        // which are still alive at this point. The same is true for dead-end
+        // CFG regions due to an infinite loop.
         for (int BitNr = Bits.find_first(); BitNr >= 0;
              BitNr = Bits.find_next(BitNr)) {
           // For each alive location extend the lifetime of all locations which
           // are alive at the allocation point. This is the same as we do for
           // a "real" deallocation instruction (see below).
+          // In dead-end CFG regions we have to do that for all blocks (because
+          // of potential infinite loops), whereas in "normal" CFG regions it's
+          // sufficient to do it at deallocation instructions.
           Bits |= StackLocs[BitNr].AliveLocs;
         }
         bd.data.AliveStackLocsAtExit = Bits;
@@ -336,6 +358,8 @@ StackNesting::Changes StackNesting::fixNesting(SILFunction *F) {
 void StackNesting::dump() const {
   for (auto bd : BlockInfos) {
     llvm::dbgs() << "Block " << bd.block.getDebugID();
+    if (bd.data.isDeadEnd)
+      llvm::dbgs() << "(deadend)";
     llvm::dbgs() << ": entry-bits=";
     dumpBits(bd.data.AliveStackLocsAtEntry);
     llvm::dbgs() << ": exit-bits=";

test/SILOptimizer/stack_promotion.sil

@@ -1156,3 +1156,24 @@ bb0(%0 : $UnownedLink):
   %r = tuple ()
   return %r : $()
 }
+
+// CHECK-LABEL: sil @dont_crash_with_wrong_stacknesting_with_infinite_loop
+// CHECK:         alloc_ref [stack] $XX
+// CHECK-LABEL: } // end sil function 'dont_crash_with_wrong_stacknesting_with_infinite_loop'
+sil @dont_crash_with_wrong_stacknesting_with_infinite_loop : $@convention(thin) () -> () {
+bb0:
+  %0 = alloc_stack $Int
+  %1 = alloc_ref $XX
+  dealloc_stack %0 : $*Int
+  cond_br undef, bb1, bb2
+bb1:
+  strong_release %1 : $XX
+  %4 = tuple ()
+  return %4 : $()
+
+bb2:
+  br bb3
+bb3:
+  br bb3
+}
+