SILOptimizer: correctly handle unreachable blocks in StackNesting.

Instead of some special treatment of unreachable blocks, model unreachable as implicitly deallocating all alive stack locations at that point.
This requires an additional forward-dataflow pass. But it now correctly models the problem and fixes a compiler crash.

rdar://problem/47402694

Author: eeckstein

include/swift/SILOptimizer/Utils/StackNesting.h

@@ -84,12 +84,8 @@ class StackNesting {
     /// The bit-set of alive stack locations at the block entry.
     BitVector AliveStackLocsAtEntry;
 
-    /// True if there is a path from this block to a function-exit.
-    ///
-    /// In other words: this block does not end in an unreachable-instruction.
-    /// This flag is only used for verifying that the lifetime of a stack
-    /// location does not end at the end of a block.
-    bool ExitReachable = false;
+    /// The bit-set of alive stack locations at the block exit.
+    BitVector AliveStackLocsAtExit;
 
     BlockInfo(SILBasicBlock *Block) : Block(Block) { }
   };
@@ -165,6 +161,19 @@ class StackNesting {
                       SILInstruction *InsertionPoint,
                       Optional<SILLocation> Location);
 
+  /// Reeturns the location bit number for a stack allocation instruction.
+  int bitNumberForAlloc(SILInstruction *AllocInst) {
+    assert(AllocInst->isAllocatingStack());
+    return StackLoc2BitNumbers[cast<SingleValueInstruction>(AllocInst)];
+  }
+
+  /// Reeturns the location bit number for a stack deallocation instruction.
+  int bitNumberForDealloc(SILInstruction *DeallocInst) {
+    assert(DeallocInst->isDeallocatingStack());
+    auto *AllocInst = cast<SingleValueInstruction>(DeallocInst->getOperand(0));
+    return bitNumberForAlloc(AllocInst);
+  }
+
   /// Modifies the SIL to end up with a correct stack nesting.
   ///
   /// Returns the status of what changes were made.

lib/SILOptimizer/Utils/StackNesting.cpp

@@ -134,9 +134,6 @@ void StackNesting::setup(SILFunction *F) {
         }
       }
     }
-    if (BI->Block->getTerminator()->isFunctionExiting())
-      BI->ExitReachable = true;
-
     for (auto *SuccBB : BI->Block->getSuccessorBlocks()) {
       BlockInfo *&SuccBI = BlockMapping[SuccBB];
       if (!SuccBI) {
@@ -165,33 +162,78 @@ bool StackNesting::solve() {
   bool isNested = false;
   BitVector Bits(StackLocs.size());
 
-  // Iterate until we reach a fixed-point.
+  // Initialize all bit fields to 1s, expect 0s for the entry block.
+  bool initVal = false;
+  for (BlockInfo &BI : BlockInfos) {
+    BI.AliveStackLocsAtEntry.resize(StackLocs.size(), initVal);
+    initVal = false;
+  }
+
+  // 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.
+  do {
+    changed = false;
+
+    for (BlockInfo &BI : BlockInfos) {
+      Bits = BI.AliveStackLocsAtEntry;
+      for (SILInstruction *StackInst : BI.StackInsts) {
+        if (StackInst->isAllocatingStack()) {
+          Bits.set(bitNumberForAlloc(StackInst));
+        } else if (StackInst->isDeallocatingStack()) {
+          Bits.reset(bitNumberForDealloc(StackInst));
+        }
+      }
+      if (Bits != BI.AliveStackLocsAtExit) {
+        BI.AliveStackLocsAtExit = Bits;
+        assert(!(BI.Block->getTerminator()->isFunctionExiting() && Bits.any())
+               && "stack location is missing dealloc");
+        changed = true;
+      }
+      // Merge the bits into the successors.
+      for (BlockInfo *SuccBI : BI.Successors) {
+        SuccBI->AliveStackLocsAtEntry &= Bits;
+      }
+    }
+  } while (changed);
+
+  // Second step: do a backward dataflow analysis to extend the lifetimes of
+  // no properly nested allocations.
   do {
     changed = false;
 
     // It's a backward dataflow problem.
     for (BlockInfo &BI : reversed(BlockInfos)) {
       // Collect the alive-bits (at the block exit) from the successor blocks.
-      Bits.reset();
       for (BlockInfo *SuccBI : BI.Successors) {
-        Bits |= SuccBI->AliveStackLocsAtEntry;
-
-        // Also get the ExitReachable flag from the successor blocks.
-        if (!BI.ExitReachable && SuccBI->ExitReachable) {
-          BI.ExitReachable = true;
-          changed = true;
+        BI.AliveStackLocsAtExit |= SuccBI->AliveStackLocsAtEntry;
+      }
+      Bits = BI.AliveStackLocsAtExit;
+
+      if (isa<UnreachableInst>(BI.Block->getTerminator())) {
+        // We treat unreachable as an implicit deallocation for all locations
+        // which are still alive at this point.
+        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).
+          Bits |= StackLocs[BitNr].AliveLocs;
         }
+        BI.AliveStackLocsAtExit = Bits;
       }
       for (SILInstruction *StackInst : reversed(BI.StackInsts)) {
         if (StackInst->isAllocatingStack()) {
-          auto AllocInst = cast<SingleValueInstruction>(StackInst);
-          int BitNr = StackLoc2BitNumbers[AllocInst];
+          int BitNr = bitNumberForAlloc(StackInst);
           if (Bits != StackLocs[BitNr].AliveLocs) {
             // More locations are alive around the StackInst's location.
             // Update the AlivaLocs bitset, which contains all those alive
             // locations.
-            assert((Bits.test(BitNr) || (!BI.ExitReachable && !Bits.any()))
-                   && "no dealloc found for alloc stack");
+            assert(Bits.test(BitNr) && "no dealloc found for alloc stack");
             StackLocs[BitNr].AliveLocs = Bits;
             changed = true;
             isNested = true;
@@ -203,10 +245,7 @@ bool StackNesting::solve() {
           // A stack deallocation begins the lifetime of its location (in
           // reverse order). And it also begins the lifetime of all other
           // locations which are alive at the allocation point.
-          auto *AllocInst =
-            cast<SingleValueInstruction>(StackInst->getOperand(0));
-          int BitNr = StackLoc2BitNumbers[AllocInst];
-          Bits |= StackLocs[BitNr].AliveLocs;
+          Bits |= StackLocs[bitNumberForDealloc(StackInst)].AliveLocs;
         }
       }
       if (Bits != BI.AliveStackLocsAtEntry) {
@@ -271,10 +310,7 @@ StackNesting::Changes StackNesting::adaptDeallocs() {
   // the same order as in solve().
   for (const BlockInfo &BI : reversed(BlockInfos)) {
     // Collect the alive-bits (at the block exit) from the successor blocks.
-    Bits.reset();
-    for (BlockInfo *SuccBI : BI.Successors) {
-      Bits |= SuccBI->AliveStackLocsAtEntry;
-    }
+    Bits = BI.AliveStackLocsAtExit;
 
     // Insert deallocations at block boundaries.
     // This can be necessary for unreachable blocks. Example:
@@ -293,12 +329,6 @@ StackNesting::Changes StackNesting::adaptDeallocs() {
          SuccIdx < NumSuccs; ++ SuccIdx) {
       BlockInfo *SuccBI = BI.Successors[SuccIdx];
 
-      // It's acceptable to not deallocate alive locations in unreachable
-      // blocks - as long as the nesting is not violated. So if there are no
-      // alive locations at the unreachable successor block, we can ignore it.
-      if (!SuccBI->ExitReachable && !SuccBI->AliveStackLocsAtEntry.any())
-        continue;
-
       if (SuccBI->AliveStackLocsAtEntry == Bits)
         continue;
 
@@ -320,16 +350,14 @@ StackNesting::Changes StackNesting::adaptDeallocs() {
     for (SILInstruction *StackInst : reversed(BI.StackInsts)) {
       if (StackInst->isAllocatingStack()) {
         // For allocations we just update the bit-set.
-        auto AllocInst = cast<SingleValueInstruction>(StackInst);
-        int BitNr = StackLoc2BitNumbers.lookup(AllocInst);
+        int BitNr = bitNumberForAlloc(StackInst);
         assert(Bits == StackLocs[BitNr].AliveLocs &&
                "dataflow didn't converge");
         Bits.reset(BitNr);
       } else if (StackInst->isDeallocatingStack()) {
         // Handle deallocations.
-        auto *AllocInst = cast<SingleValueInstruction>(StackInst->getOperand(0));
         SILLocation Loc = StackInst->getLoc();
-        int BitNr = StackLoc2BitNumbers.lookup(AllocInst);
+        int BitNr = bitNumberForDealloc(StackInst);
         SILInstruction *InsertionPoint = &*std::next(StackInst->getIterator());
         if (Bits.test(BitNr)) {
           // The location of StackInst is alive after StackInst. So we have to
@@ -372,17 +400,18 @@ void StackNesting::dump() const {
       continue;
 
     llvm::dbgs() << "Block " << BI.Block->getDebugID();
-    if (!BI.ExitReachable)
-      llvm::dbgs() << " (unreachable exit)";
-    llvm::dbgs() << ": bits=";
+    llvm::dbgs() << ": entry-bits=";
     dumpBits(BI.AliveStackLocsAtEntry);
+    llvm::dbgs() << ": exit-bits=";
+    dumpBits(BI.AliveStackLocsAtExit);
+    llvm::dbgs() << '\n';
     for (SILInstruction *StackInst : BI.StackInsts) {
       if (StackInst->isAllocatingStack()) {
         auto AllocInst = cast<SingleValueInstruction>(StackInst);
         int BitNr = StackLoc2BitNumbers.lookup(AllocInst);
         llvm::dbgs() << "  alloc #" << BitNr << ": alive=";
         dumpBits(StackLocs[BitNr].AliveLocs);
-        llvm::dbgs() << "    " << *StackInst;
+        llvm::dbgs() << ",     " << *StackInst;
       } else if (StackInst->isDeallocatingStack()) {
         auto *AllocInst = cast<SingleValueInstruction>(StackInst->getOperand(0));
         int BitNr = StackLoc2BitNumbers.lookup(AllocInst);
@@ -405,6 +434,5 @@ void StackNesting::dumpBits(const BitVector &Bits) {
     llvm::dbgs() << separator << Bit;
     separator = ",";
   }
-  llvm::dbgs() << ">\n";
+  llvm::dbgs() << '>';
 }
-

test/SILOptimizer/allocbox_to_stack.sil

@@ -819,6 +819,8 @@ bb0(%0 : $Int):
 // CHECK-NEXT:   [[STACK1:%[0-9]+]] = alloc_stack $Bool
 // CHECK-NEXT:   [[BOX:%[0-9]+]] = alloc_stack $Int
 // CHECK:      bb1:
+// CHECK-NEXT:   dealloc_stack [[BOX]]
+// CHECK-NEXT:   dealloc_stack [[STACK1]]
 // CHECK-NEXT:   [[STACK2:%[0-9]+]] = alloc_stack $Bool
 // CHECK-NEXT:   dealloc_stack [[STACK2]]
 // CHECK-NEXT:   unreachable
@@ -952,6 +954,7 @@ bb2:
 // CHECK:      bb0(%0 : $Int):
 // CHECK-NEXT:   [[BOX:%[0-9]+]] = alloc_stack $Int
 // CHECK:      bb1:
+// CHECK-NEXT:   dealloc_stack [[BOX]]
 // CHECK-NEXT:   {{.*}} = alloc_stack $Bool
 // CHECK-NEXT:   {{.*}} = alloc_stack $Bool
 // CHECK-NEXT:   unreachable
@@ -979,6 +982,38 @@ bb2:
   return %r : $()
 }
 
+// CHECK-LABEL: sil @nesting_and_unreachable6
+// CHECK:      bb0(%0 : $Int):
+// CHECK-NEXT:   [[BOX:%[0-9]+]] = alloc_stack $Int
+// CHECK:      bb1:
+// CHECK-NEXT:   dealloc_stack [[BOX]]
+// CHECK-NEXT:   br bb3
+// CHECK:      bb2:
+// CHECK-NEXT:   alloc_stack
+// CHECK-NEXT:   unreachable
+// CHECK:      bb3:
+// CHECK-NEXT:   return
+sil @nesting_and_unreachable6 : $@convention(thin) (Int) -> Int {
+bb0(%0 : $Int):
+  %1 = alloc_box ${ var Int }
+  %1a = project_box %1 : ${ var Int }, 0
+  store %0 to %1a : $*Int
+  %3 = load %1a : $*Int
+  cond_br undef, bb1, bb2
+
+bb1:
+  strong_release %1 : ${ var Int }
+  br bb3
+
+bb2:
+  %241 = alloc_stack $Int
+  strong_release %1 : ${ var Int }
+  unreachable
+
+bb3:
+  return %3 : $Int
+}
+
 // CHECK-LABEL: sil @nesting_and_unreachable_critical_edge
 // CHECK:      bb0(%0 : $Int):
 // CHECK-NEXT:   [[BOX:%[0-9]+]] = alloc_stack $Int

test/SILOptimizer/allocbox_to_stack_ownership.sil

@@ -814,6 +814,8 @@ bb0(%0 : $Int):
 // CHECK-NEXT:   [[STACK1:%[0-9]+]] = alloc_stack $Bool
 // CHECK-NEXT:   [[BOX:%[0-9]+]] = alloc_stack $Int
 // CHECK:      bb1:
+// CHECK-NEXT:   dealloc_stack [[BOX]]
+// CHECK-NEXT:   dealloc_stack [[STACK1]]
 // CHECK-NEXT:   [[STACK2:%[0-9]+]] = alloc_stack $Bool
 // CHECK-NEXT:   dealloc_stack [[STACK2]]
 // CHECK-NEXT:   unreachable
@@ -947,6 +949,7 @@ bb2:
 // CHECK:      bb0(%0 : $Int):
 // CHECK-NEXT:   [[BOX:%[0-9]+]] = alloc_stack $Int
 // CHECK:      bb1:
+// CHECK-NEXT:   dealloc_stack [[BOX]]
 // CHECK-NEXT:   {{.*}} = alloc_stack $Bool
 // CHECK-NEXT:   {{.*}} = alloc_stack $Bool
 // CHECK-NEXT:   unreachable