From f494f370a55baef26a6c97e3f16b7f0b9f5380e6 Mon Sep 17 00:00:00 2001
From: Erik Eckstein <eeckstein@apple.com>
Date: Wed, 24 Mar 2021 12:19:14 +0100
Subject: [PATCH] TempRValueOpt: Support yield

This is one step of eliminating a copy_addr in the Array's subscript read coroutine

rdar://53996328
---
 .../Transforms/TempRValueElimination.cpp      | 13 +++++++++-
 test/SILOptimizer/temp_rvalue_opt_ossa.sil    | 24 +++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/lib/SILOptimizer/Transforms/TempRValueElimination.cpp b/lib/SILOptimizer/Transforms/TempRValueElimination.cpp
index becb48225bb9c..57190b2a7f03d 100644
--- a/lib/SILOptimizer/Transforms/TempRValueElimination.cpp
+++ b/lib/SILOptimizer/Transforms/TempRValueElimination.cpp
@@ -215,6 +215,15 @@ collectLoads(Operand *addressUse, CopyAddrInst *originalCopy,
     }
     return true;
   }
+  case SILInstructionKind::YieldInst: {
+    auto *yield = cast<YieldInst>(user);
+    auto convention = yield->getArgumentConventionForOperand(*addressUse);
+    if (!convention.isGuaranteedConvention())
+      return false;
+
+    loadInsts.insert(user);
+    return true;
+  }
   case SILInstructionKind::OpenExistentialAddrInst: {
     // We only support open existential addr if the access is immutable.
     auto *oeai = cast<OpenExistentialAddrInst>(user);
@@ -331,8 +340,10 @@ SILInstruction *TempRValueOptPass::getLastUseWhileSourceIsNotModified(
     if (numLoadsFound == useInsts.size()) {
       // Function calls are an exception: in a called function a potential
       // modification of copySrc could occur _before_ the read of the temporary.
-      if (FullApplySite::isa(inst) && aa->mayWriteToMemory(inst, copySrc))
+      if ((FullApplySite::isa(inst) || isa<YieldInst>(inst)) &&
+          aa->mayWriteToMemory(inst, copySrc)) {
         return nullptr;
+      }
 
       return inst;
     }
diff --git a/test/SILOptimizer/temp_rvalue_opt_ossa.sil b/test/SILOptimizer/temp_rvalue_opt_ossa.sil
index 64ef98b3370a2..0a76accfbf81d 100644
--- a/test/SILOptimizer/temp_rvalue_opt_ossa.sil
+++ b/test/SILOptimizer/temp_rvalue_opt_ossa.sil
@@ -1386,3 +1386,27 @@ bb0(%0 : @owned $Klass):
   return %105 : $Klass
 }
 
+// CHECK-LABEL: sil [ossa] @test_yield
+// CHECK:   [[TA:%[0-9]+]] = ref_tail_addr
+// CHECK-NOT: copy_addr
+// CHECK:   yield [[TA]]
+// CHECK: } // end sil function 'test_yield'
+sil [ossa] @test_yield : $@yield_once @convention(thin) (@guaranteed Klass) -> @yields @in_guaranteed Two {
+bb0(%0 : @guaranteed $Klass):
+  %1 = alloc_stack $Two
+  %2 = ref_tail_addr [immutable] %0 : $Klass, $Two
+  copy_addr %2 to [initialization] %1 : $*Two
+  yield %1 : $*Two, resume bb1, unwind bb2
+
+bb1:
+  destroy_addr %1 : $*Two
+  dealloc_stack %1 : $*Two
+  %90 = tuple ()
+  return %90 : $()
+
+bb2:
+  destroy_addr %1 : $*Two
+  dealloc_stack %1 : $*Two
+  unwind
+}
+