Merge pull request #27893 from atrick/fix-temprvalue-destroy

Author: swift-ci

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;

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 : $()
+}