SILOptimizer: handle copy_addr [take] in TempRValueOpt

Author: eeckstein

lib/SILOptimizer/Transforms/CopyForwarding.cpp

@@ -1604,10 +1604,7 @@ void TempRValueOptPass::run() {
       ++II;
 
       // Remove identity copies which are a result of this optimization.
-      if (CopyInst && CopyInst->getSrc() == CopyInst->getDest() &&
-          // Identity copies cannot take the source. This check is just here
-          // to be on the safe side.
-          !CopyInst->isTakeOfSrc()) {
+      if (CopyInst && CopyInst->getSrc() == CopyInst->getDest()) {
         // This is either the CopyInst which just got optimized or it is a
         // follow-up from an earlier iteration, where another copy_addr copied
         // the temporary back to the source location.
@@ -1725,9 +1722,8 @@ bool TempRValueOptPass::collectLoads(
 
   case SILInstructionKind::CopyAddrInst: {
     // copy_addr which read from the temporary are like loads.
-    // TODO: Handle copy_addr [take]. But this doesn't seem to be important.
     auto *copyFromTmp = cast<CopyAddrInst>(user);
-    if (copyFromTmp->getDest() == address || copyFromTmp->isTakeOfSrc()) {
+    if (copyFromTmp->getDest() == address) {
       LLVM_DEBUG(llvm::dbgs() << "  Temp written or taken" << *user);
       return false;
     }
@@ -1774,7 +1770,7 @@ bool TempRValueOptPass::checkNoSourceModification(CopyAddrInst *copyInst,
 
 /// Tries to perform the temporary rvalue copy elimination for \p copyInst
 bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
-  if (copyInst->isTakeOfSrc() || !copyInst->isInitializationOfDest())
+  if (!copyInst->isInitializationOfDest())
     return false;
 
   auto *tempObj = dyn_cast<AllocStackInst>(copyInst->getDest());
@@ -1818,10 +1814,25 @@ bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
     SILInstruction *user = use->getUser();
     switch (user->getKind()) {
     case SILInstructionKind::DestroyAddrInst:
+      if (copyInst->isTakeOfSrc()) {
+        use->set(copyInst->getSrc());
+      } else {
+        user->eraseFromParent();
+      }
+      break;
     case SILInstructionKind::DeallocStackInst:
       user->eraseFromParent();
       break;
-    case SILInstructionKind::CopyAddrInst:
+    case SILInstructionKind::CopyAddrInst: {
+      auto *CAI = cast<CopyAddrInst>(user);
+      if (CAI != copyInst) {
+        assert(CAI->getSrc() == tempObj);
+        if (CAI->isTakeOfSrc() && !copyInst->isTakeOfSrc())
+          CAI->setIsTakeOfSrc(IsNotTake);
+      }
+      use->set(copyInst->getSrc());
+      break;
+    }
     case SILInstructionKind::StructElementAddrInst:
     case SILInstructionKind::TupleElementAddrInst:
     case SILInstructionKind::LoadInst:

test/SILOptimizer/temp_rvalue_opt.sil

@@ -102,14 +102,10 @@ bb0(%0 : $*GS<B>, %1 : $*GS<B>):
   return %9999 : $()
 }
 
-// Currently not supported
 // CHECK-LABEL: sil @take_from_temp
 // CHECK:      bb0(%0 : $*B, %1 : $*GS<B>):
-// CHECK-NEXT:   alloc_stack
-// CHECK-NEXT:   copy_addr
 // CHECK-NEXT:   struct_element_addr
-// CHECK-NEXT:   copy_addr
-// CHECK-NEXT:   dealloc_stack
+// CHECK-NEXT:   copy_addr [take]
 // CHECK-NEXT:   tuple
 // CHECK-NEXT:   return
 sil @take_from_temp : $@convention(thin) <B> (@inout B, @inout GS<B>) -> () {
@@ -416,13 +412,17 @@ bb0(%0 : $*Klass):
   return %9 : $()
 }
 
-// Make sure that we do not optimize this case.
-//
 // CHECK-LABEL: sil @non_overlapping_lifetime : $@convention(thin) (@in Klass) -> () {
-// CHECK: destroy_addr
-// CHECK: destroy_addr
-// CHECK: destroy_addr
-// CHECK: } // end sil function 'non_overlapping_lifetime'
+// CHECK:      bb0([[ARG:%.*]] : $*Klass):
+// CHECK-NEXT:   [[TMP:%.*]] = alloc_stack $Klass
+// CHECK-NEXT:   copy_addr [[ARG]] to [initialization] [[TMP]]
+// CHECK-NEXT:   destroy_addr [[ARG]]
+// CHECK:        apply %{{[0-9]*}}([[TMP]])
+// CHECK-NEXT:   destroy_addr [[TMP]]
+// CHECK-NEXT:   dealloc_stack [[TMP]]
+// CHECK-NEXT:   tuple
+// CHECK-NEXT:   return
+// CHECK-NEXT: } // end sil function 'non_overlapping_lifetime'
 sil @non_overlapping_lifetime : $@convention(thin) (@in Klass) -> () {
 bb0(%0 : $*Klass):
   %1a = alloc_stack $Klass