SILCombine: fix a miscompile in the alloc_stack optimization which causes a use-after-free.

lib/SILOptimizer/SILCombiner/SILCombinerMiscVisitors.cpp

@@ -417,6 +417,36 @@ struct AllocStackAnalyzer : SILInstructionVisitor<AllocStackAnalyzer> {
 
 } // end anonymous namespace
 
+/// Returns true if there is a retain instruction between \p from and the
+/// destroy or deallocation of \p alloc.
+static bool somethingIsRetained(SILInstruction *from, AllocStackInst *alloc) {
+  llvm::SmallVector<SILInstruction *, 8> workList;
+  llvm::SmallPtrSet<SILBasicBlock *, 8> handled;
+  workList.push_back(from);
+  while (!workList.empty()) {
+    SILInstruction *start = workList.pop_back_val();
+    for (auto iter = start->getIterator(), end = start->getParent()->end();
+         iter != end;
+         ++iter) {
+      SILInstruction *inst = &*iter;
+      if (isa<RetainValueInst>(inst) || isa<StrongRetainInst>(inst)) {
+        return true;
+      }
+      if ((isa<DeallocStackInst>(inst) || isa<DestroyAddrInst>(inst)) &&
+          inst->getOperand(0) == alloc) {
+        break;
+      }
+      if (isa<TermInst>(inst)) {
+        for (SILBasicBlock *succ : start->getParent()->getSuccessors()) {
+          if (handled.insert(succ).second)
+            workList.push_back(&*succ->begin());
+        }
+      }
+    }
+  }
+  return false;
+}
+
 SILInstruction *SILCombiner::visitAllocStackInst(AllocStackInst *AS) {
   // If we are testing SILCombine and we are asked not to eliminate
   // alloc_stacks, just return.
@@ -477,15 +507,15 @@ SILInstruction *SILCombiner::visitAllocStackInst(AllocStackInst *AS) {
   //
   // TODO: Do we not remove purely dead live ranges here? Seems like we should.
   SmallPtrSet<SILInstruction *, 16> ToDelete;
+  SmallVector<CopyAddrInst *, 4> takingCopies;
 
   for (auto *Op : AS->getUses()) {
     // Replace a copy_addr [take] %src ... by a destroy_addr %src if %src is
     // no the alloc_stack.
     // Otherwise, just delete the copy_addr.
     if (auto *CopyAddr = dyn_cast<CopyAddrInst>(Op->getUser())) {
       if (CopyAddr->isTakeOfSrc() && CopyAddr->getSrc() != AS) {
-        Builder.setInsertionPoint(CopyAddr);
-        Builder.createDestroyAddr(CopyAddr->getLoc(), CopyAddr->getSrc());
+        takingCopies.push_back(CopyAddr);
       }
     }
 
@@ -506,6 +536,20 @@ SILInstruction *SILCombiner::visitAllocStackInst(AllocStackInst *AS) {
     ToDelete.insert(Op->getUser());
   }
 
+  // Check if a retain is moved after the copy_addr. If the retained object
+  // happens to be the source of the copy_addr it might be only kept alive by
+  // the stack location. This cannot happen with OSSA.
+  // TODO: remove this check once we have OSSA.
+  for (CopyAddrInst *copy : takingCopies) {
+    if (somethingIsRetained(copy, AS))
+      return nullptr;
+  }
+
+  for (CopyAddrInst *copy : takingCopies) {
+    SILBuilderWithScope destroyBuilder(copy, Builder.getBuilderContext());
+    destroyBuilder.createDestroyAddr(copy->getLoc(), copy->getSrc());
+  }
+
   // Erase the 'live-range'
   for (auto *Inst : ToDelete) {
     Inst->replaceAllUsesOfAllResultsWithUndef();

test/SILOptimizer/sil_combine.sil

@@ -2640,6 +2640,107 @@ bb0(%0 : $*B):
   return %4 : $()
 }
 
+// CHECK-LABEL: sil @moved_retain_prevents_alloc_stack_deletion1
+// CHECK: copy_addr
+// CHECK: strong_retain
+// CHECK: destroy_addr
+// CHECK: } // end sil function 'moved_retain_prevents_alloc_stack_deletion1'
+sil @moved_retain_prevents_alloc_stack_deletion1 : $@convention(thin) (@guaranteed B) -> () {
+bb0(%0 : $B):
+  %3 = alloc_stack $B
+  %4 = alloc_stack $B
+  store %0 to %4 : $*B
+  copy_addr [take] %4 to [initialization] %3 : $*B
+  dealloc_stack %4 : $*B
+  strong_retain %0 : $B
+  destroy_addr %3 : $*B
+  dealloc_stack %3 : $*B
+  %14 = tuple ()
+  return %14 : $()
+}
+
+// CHECK-LABEL: sil @moved_retain_prevents_alloc_stack_deletion2
+// CHECK:   copy_addr
+// CHECK: bb1:
+// CHECK:   strong_retain
+// CHECK:   destroy_addr
+// CHECK: bb2:
+// CHECK:   strong_retain
+// CHECK:   destroy_addr
+// CHECK: } // end sil function 'moved_retain_prevents_alloc_stack_deletion2'
+sil @moved_retain_prevents_alloc_stack_deletion2 : $@convention(thin) (@guaranteed B) -> () {
+bb0(%0 : $B):
+  %3 = alloc_stack $B
+  %4 = alloc_stack $B
+  store %0 to %4 : $*B
+  copy_addr [take] %4 to [initialization] %3 : $*B
+  dealloc_stack %4 : $*B
+  cond_br undef, bb1, bb2
+bb1:
+  strong_retain %0 : $B
+  destroy_addr %3 : $*B
+  dealloc_stack %3 : $*B
+  br bb3
+bb2:
+  strong_retain %0 : $B
+  destroy_addr %3 : $*B
+  dealloc_stack %3 : $*B
+  br bb3
+bb3:
+  %14 = tuple ()
+  return %14 : $()
+}
+
+// CHECK-LABEL: sil @retain_is_after_stack_location
+// CHECK-NOT:   copy_addr
+// CHECK: } // end sil function 'retain_is_after_stack_location'
+sil @retain_is_after_stack_location : $@convention(thin) (@guaranteed B) -> () {
+bb0(%0 : $B):
+  %3 = alloc_stack $B
+  %4 = alloc_stack $B
+  store %0 to %4 : $*B
+  copy_addr [take] %4 to [initialization] %3 : $*B
+  dealloc_stack %4 : $*B
+  cond_br undef, bb1, bb2
+bb1:
+  destroy_addr %3 : $*B
+  strong_retain %0 : $B
+  dealloc_stack %3 : $*B
+  br bb3
+bb2:
+  destroy_addr %3 : $*B
+  strong_retain %0 : $B
+  dealloc_stack %3 : $*B
+  br bb3
+bb3:
+  %14 = tuple ()
+  return %14 : $()
+}
+
+// CHECK-LABEL: sil @moved_retain_prevents_alloc_stack_deletion3
+// CHECK:   copy_addr
+// CHECK: bb2:
+// CHECK:   strong_retain
+// CHECK:   destroy_addr
+// CHECK: } // end sil function 'moved_retain_prevents_alloc_stack_deletion3'
+sil @moved_retain_prevents_alloc_stack_deletion3 : $@convention(thin) (@guaranteed B) -> () {
+bb0(%0 : $B):
+  %3 = alloc_stack $B
+  %4 = alloc_stack $B
+  store %0 to %4 : $*B
+  copy_addr [take] %4 to [initialization] %3 : $*B
+  dealloc_stack %4 : $*B
+  br bb1
+bb1:
+  cond_br undef, bb1, bb2
+bb2:
+  strong_retain %0 : $B
+  destroy_addr %3 : $*B
+  dealloc_stack %3 : $*B
+  %14 = tuple ()
+  return %14 : $()
+}
+
 protocol someProtocol {
     var i: Int { get }
 }

test/SILOptimizer/sil_combine_alloc_stack.swift

@@ -0,0 +1,43 @@
+// RUN: %empty-directory(%t) 
+// RUN: %target-build-swift -O %s -o %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s
+
+protocol E {
+  func f() -> Bool
+}
+
+final class K {
+  deinit {
+    print("deinit")
+  }
+}
+
+
+struct X : E {
+  var x: K
+  func f() -> Bool { return true }
+}
+
+func g<T>(_ x : T) -> Bool {
+  if let y = x as? E { return y.f() }
+  return false
+}
+
+// CHECK that there is no use-after-free in this function.
+@inline(never)
+func foo(_ x: X) -> Bool {
+  return g(x)
+}
+
+@inline(never)
+func testit() {
+  let x = X(x: K())
+  _ = foo(x)
+  print(x)
+}
+
+// CHECK: X(x: a.K)
+// CHECK: deinit
+testit()
+
+