Merge pull request #2264 from eeckstein/stack-promotion

Don't stack promote allocations in no-return blocks.

Author: eeckstein

lib/SILOptimizer/Transforms/StackPromotion.cpp

@@ -220,8 +220,35 @@ static bool isPromotableAllocInst(SILInstruction *I) {
 }
 
 StackPromoter::ChangeState StackPromoter::promote() {
+
+  llvm::SetVector<SILBasicBlock *> ReachableBlocks;
+
+  // First step: find blocks which end up in a no-return block (terminated by
+  // an unreachable instruction).
+  // Search for function-exiting blocks, i.e. return and throw.
+  for (SILBasicBlock &BB : *F) {
+    TermInst *TI = BB.getTerminator();
+    if (TI->isFunctionExiting())
+      ReachableBlocks.insert(&BB);
+  }
+  // Propagate the reachability up the control flow graph.
+  unsigned Idx = 0;
+  while (Idx < ReachableBlocks.size()) {
+    SILBasicBlock *BB = ReachableBlocks[Idx++];
+    for (SILBasicBlock *Pred : BB->getPreds())
+      ReachableBlocks.insert(Pred);
+  }
+
   // Search the whole function for stack promotable allocations.
   for (SILBasicBlock &BB : *F) {
+
+    // Don't stack promote any allocation inside a code region which ends up in
+    // a no-return block. Such allocations may missing their final release.
+    // We would insert the deallocation too early, which may result in a
+    // use-after-free problem.
+    if (ReachableBlocks.count(&BB) == 0)
+      continue;
+
     for (auto Iter = BB.begin(); Iter != BB.end();) {
       // The allocation instruction may be moved, so increment Iter prior to
       // doing the optimization.

test/SILOptimizer/stack_promotion.sil

@@ -39,6 +39,7 @@ bb0(%0 : $YY):
   return %t : $()
 }
 
+sil @consume_int : $@convention(thin) (Int32) -> ()
 
 // CHECK-LABEL: sil @simple_promote
 // CHECK: [[O:%[0-9]+]] = alloc_ref [stack] $XX
@@ -68,6 +69,29 @@ bb0:
   return %n1 : $XX
 }
 
+// CHECK-LABEL: sil @dont_promote_in_unreachable
+// CHECK: bb1:
+// CHECK-NEXT: alloc_ref $XX
+sil @dont_promote_in_unreachable : $@convention(thin) () -> () {
+bb0:
+  cond_br undef, bb1, bb2
+
+bb1:
+  %o1 = alloc_ref $XX
+  %f1 = function_ref @xx_init : $@convention(thin) (@guaranteed XX) -> XX
+  %n1 = apply %f1(%o1) : $@convention(thin) (@guaranteed XX) -> XX
+  %l1 = ref_element_addr %n1 : $XX, #XX.x
+  %l2 = load %l1 : $*Int32
+  %f2 = function_ref @consume_int : $@convention(thin) (Int32) -> ()
+  %a = apply %f2(%l2) : $@convention(thin) (Int32) -> ()
+  // An unreachable block may missing the final release.
+  unreachable
+
+bb2:
+  %r = tuple ()
+  return %r : $()
+}
+
 // CHECK-LABEL: sil @promote_nested
 // CHECK: [[X:%[0-9]+]] = alloc_ref [stack] $XX
 // CHECK: [[Y:%[0-9]+]] = alloc_ref [stack] $YY