Fix TempRValue: Add checkTempObjectDestroy bail out logic.

This avoids use-after-free bugs that can be introduced by removing a
copy without replacing all corresponding destroys.

Add more descriptive comments since multiple bugs have been introduced
in this pass.

None of this will be relevant once the pass is converted to OSSA.

Author: atrick

Diff for: lib/SILOptimizer/Transforms/CopyForwarding.cpp

@@ -69,6 +69,7 @@
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/CFGOptUtils.h"
+#include "swift/SILOptimizer/Utils/ValueLifetime.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Support/CommandLine.h"
@@ -1569,6 +1570,8 @@ class TempRValueOptPass : public SILFunctionTransform {
   bool checkNoSourceModification(CopyAddrInst *copyInst,
                        const llvm::SmallPtrSetImpl<SILInstruction *> &useInsts);
 
+  bool checkTempObjectDestroy(AllocStackInst *tempObj, CopyAddrInst *copyInst);
+
   bool tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst);
 
   void run() override;
@@ -1620,6 +1623,17 @@ void TempRValueOptPass::run() {
 
 /// Transitively explore all data flow uses of the given \p address until
 /// reaching a load or returning false.
+///
+/// Any user opcode recognized by collectLoads must be replaced correctly later
+/// during tryOptimizeCopyIntoTemp. If it is possible for any use to destroy the
+/// value in \p address, then that use must be removed or made non-destructive
+/// after the copy is removed and its operand is replaced.
+///
+/// Warning: To preserve the original object lifetime, tryOptimizeCopyIntoTemp
+/// must assume that there are no holes in lifetime of the temporary stack
+/// location at \address. The temporary must be initialized by the original copy
+/// and never written to again. Therefore, collectLoads disallows any operation
+/// that may write to memory at \p address.
 bool TempRValueOptPass::collectLoads(
     Operand *userOp, SILInstruction *user, SingleValueInstruction *address,
     SILValue srcObject,
@@ -1768,6 +1782,73 @@ bool TempRValueOptPass::checkNoSourceModification(CopyAddrInst *copyInst,
   return false;
 }
 
+/// Return true if the \p tempObj, which is initialized by \p copyInst, is
+/// destroyed in an orthodox way.
+///
+/// When tryOptimizeCopyIntoTemp replaces all of tempObj's uses, it assumes that
+/// the object is initialized by the original copy and directly destroyed on all
+/// paths by one of the recognized 'destroy_addr' or 'copy_addr [take]'
+/// operations. This assumption must be checked. For example, in non-OSSA,
+/// it is legal to destroy an in-memory object by loading the value and
+/// releasing it. Rather than detecting unbalanced load releases, simply check
+/// that tempObj is destroyed directly on all paths.
+bool TempRValueOptPass::checkTempObjectDestroy(AllocStackInst *tempObj,
+                                               CopyAddrInst *copyInst) {
+  // If the original copy was a take, then replacing all uses cannot affect
+  // the lifetime.
+  if (copyInst->isTakeOfSrc())
+    return true;
+
+  // ValueLifetimeAnalysis is not normally used for address types. It does not
+  // reason about the lifetime of the in-memory object. However the utility can
+  // be abused here to check that the address is directly destroyed on all
+  // paths. collectLoads has already guaranteed that tempObj's lifetime has no
+  // holes/reinitializations.
+  SmallVector<SILInstruction *, 8> users;
+  for (auto result : tempObj->getResults()) {
+    for (Operand *operand : result->getUses()) {
+      SILInstruction *user = operand->getUser();
+      if (user == copyInst)
+        continue;
+      if (isa<DeallocStackInst>(user))
+        continue;
+      users.push_back(user);
+    }
+  }
+  // Find the boundary of tempObj's address lifetime, starting at copyInst.
+  ValueLifetimeAnalysis vla(copyInst, users);
+  ValueLifetimeAnalysis::Frontier tempAddressFrontier;
+  if (!vla.computeFrontier(tempAddressFrontier,
+                           ValueLifetimeAnalysis::DontModifyCFG)) {
+    return false;
+  }
+  // Check that the lifetime boundary ends at direct destroy points.
+  for (SILInstruction *frontierInst : tempAddressFrontier) {
+    auto pos = frontierInst->getIterator();
+    // If the frontier is at the head of a block, then either it is an
+    // unexpected lifetime exit, or the lifetime ended at a
+    // terminator. TempRValueOptPass does not handle either case.
+    if (pos == frontierInst->getParent()->begin())
+      return false;
+
+    // Look for a known destroy point as described in the funciton level
+    // comment. This whitelist can be expanded as more cases are handled in
+    // tryOptimizeCopyIntoTemp during copy replacement.
+    SILInstruction *lastUser = &*std::prev(pos);
+    if (isa<DestroyAddrInst>(lastUser))
+      continue;
+
+    if (auto *cai = dyn_cast<CopyAddrInst>(lastUser)) {
+      assert(cai->getSrc() == tempObj && "collectLoads checks for writes");
+      assert(!copyInst->isTakeOfSrc() && "checked above");
+      if (cai->isTakeOfSrc())
+        continue;
+    }
+    return false;
+  }
+  return true;
+}
+
 /// Tries to perform the temporary rvalue copy elimination for \p copyInst
 bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
   if (!copyInst->isInitializationOfDest())
@@ -1803,6 +1884,9 @@ bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
   if (!checkNoSourceModification(copyInst, loadInsts))
     return false;
 
+  if (!checkTempObjectDestroy(tempObj, copyInst))
+    return false;
+
   LLVM_DEBUG(llvm::dbgs() << "  Success: replace temp" << *tempObj);
 
   // Do a "replaceAllUses" by either deleting the users or replacing them with
@@ -1833,6 +1917,9 @@ bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
       use->set(copyInst->getSrc());
       break;
     }
+    // ASSUMPTION: no operations that may be handled by this default clause can
+    // destroy tempObj. This includes operations that load the value from memory
+    // and release it.
     default:
       use->set(copyInst->getSrc());
       break;

Diff for: test/SILOptimizer/temp_rvalue_opt.sil

@@ -106,8 +106,11 @@ bb0(%0 : $*GS<B>, %1 : $*GS<B>):
 
 // CHECK-LABEL: sil @take_from_temp
 // CHECK:      bb0(%0 : $*B, %1 : $*GS<B>):
-// CHECK-NEXT:   struct_element_addr
-// CHECK-NEXT:   copy_addr [take]
+// CHECK-NEXT:   [[STK:%.*]] = alloc_stack
+// CHECK-NEXT:   copy_addr %1 to [initialization] [[STK]]
+// CHECK-NEXT:   [[INNER:%.*]] = struct_element_addr
+// CHECK-NEXT:   copy_addr [take] [[INNER]]
+// CHECK-NEXT:   dealloc_stack
 // CHECK-NEXT:   tuple
 // CHECK-NEXT:   return
 sil @take_from_temp : $@convention(thin) <B> (@inout B, @inout GS<B>) -> () {
@@ -600,3 +603,55 @@ bb3:
   dealloc_stack %2 : $*Optional<P>
   br bb2
 }
+
+///////////////////////////////////////////////////////////////////////////////
+// Test checkTempObjectDestroy
+// <rdar://problem/56393491> Use-after free crashing an XCTest.
+
+sil @takeGuaranteedObj : $@convention(thin) (@guaranteed Builtin.NativeObject) -> ()
+
+// Do not remove a copy that is released via a load (because
+// TempRValueOpt is an address-based optimization does not know how to
+// remove releases, and trying to do that would reduce to ARC
+// optimization).
+// CHECK-LABEL: sil @copyWithLoadRelease : $@convention(thin) (@in_guaranteed Builtin.NativeObject) -> () {
+// CHECK: bb0(%0 : $*Builtin.NativeObject):
+// CHECK:   [[STK:%.*]] = alloc_stack $Builtin.NativeObject
+// CHECK:   copy_addr %0 to [initialization] [[STK]] : $*Builtin.NativeObject
+// CHECK:   [[VAL:%.*]] = load [[STK]] : $*Builtin.NativeObject
+// CHECK:   apply %{{.*}}([[VAL]]) : $@convention(thin) (@guaranteed Builtin.NativeObject) -> ()
+// CHECK:   release_value [[VAL]] : $Builtin.NativeObject
+// CHECK:   dealloc_stack [[STK]] : $*Builtin.NativeObject
+// CHECK-LABEL: } // end sil function 'copyWithLoadRelease'
+sil @copyWithLoadRelease : $@convention(thin) (@in_guaranteed Builtin.NativeObject) -> () {
+bb0(%0 : $*Builtin.NativeObject):
+  %stk = alloc_stack $Builtin.NativeObject
+  copy_addr %0 to [initialization] %stk : $*Builtin.NativeObject
+  %obj = load %stk : $*Builtin.NativeObject
+  %f = function_ref @takeGuaranteedObj : $@convention(thin) (@guaranteed Builtin.NativeObject) -> ()
+  %call = apply %f(%obj) : $@convention(thin) (@guaranteed Builtin.NativeObject) -> ()
+  release_value %obj : $Builtin.NativeObject
+  dealloc_stack %stk : $*Builtin.NativeObject
+  %v = tuple ()
+  return %v : $()
+}
+
+// Remove a copy that is released via a load as long as it was a copy [take].
+// CHECK-LABEL: sil @takeWithLoadRelease : $@convention(thin) (@in Builtin.NativeObject) -> () {
+// CHECK: bb0(%0 : $*Builtin.NativeObject):
+// CHECK:   [[V:%.*]] = load %0 : $*Builtin.NativeObject
+// CHECK:   apply %{{.*}}([[V]]) : $@convention(thin) (@guaranteed Builtin.NativeObject) -> ()
+// CHECK:   release_value [[V]] : $Builtin.NativeObject
+// CHECK-LABEL: } // end sil function 'takeWithLoadRelease'
+sil @takeWithLoadRelease : $@convention(thin) (@in Builtin.NativeObject) -> () {
+bb0(%0 : $*Builtin.NativeObject):
+  %stk = alloc_stack $Builtin.NativeObject
+  copy_addr [take] %0 to [initialization] %stk : $*Builtin.NativeObject
+  %obj = load %stk : $*Builtin.NativeObject
+  %f = function_ref @takeGuaranteedObj : $@convention(thin) (@guaranteed Builtin.NativeObject) -> ()
+  %call = apply %f(%obj) : $@convention(thin) (@guaranteed Builtin.NativeObject) -> ()
+  release_value %obj : $Builtin.NativeObject
+  dealloc_stack %stk : $*Builtin.NativeObject
+  %v = tuple ()
+  return %v : $()
+}