Merge pull request #41555 from nate-chandler/mem_access_utils/unique_storage_use_visitor/count-coro-ends-as-uses

[MemAccessUtils] Regard end/abort_apply as use.

Author: swift-ci

lib/SIL/Utils/MemAccessUtils.cpp

@@ -1948,6 +1948,29 @@ bool UniqueStorageUseVisitor::findUses(UniqueStorageUseVisitor &visitor) {
   return visitAccessStorageUses(gather, visitor.storage, visitor.function);
 }
 
+static bool
+visitApplyOperand(Operand *use, UniqueStorageUseVisitor &visitor,
+                  bool (UniqueStorageUseVisitor::*visit)(Operand *)) {
+  auto *user = use->getUser();
+  if (auto *bai = dyn_cast<BeginApplyInst>(user)) {
+    if (!(visitor.*visit)(use))
+      return false;
+    SmallVector<Operand *, 2> endApplyUses;
+    SmallVector<Operand *, 2> abortApplyUses;
+    bai->getCoroutineEndPoints(endApplyUses, abortApplyUses);
+    for (auto *endApplyUse : endApplyUses) {
+      if (!(visitor.*visit)(endApplyUse))
+        return false;
+    }
+    for (auto *abortApplyUse : abortApplyUses) {
+      if (!(visitor.*visit)(abortApplyUse))
+        return false;
+    }
+    return true;
+  }
+  return (visitor.*visit)(use);
+}
+
 bool GatherUniqueStorageUses::visitUse(Operand *use, AccessUseType useTy) {
   unsigned operIdx = use->getOperandNumber();
   auto *user = use->getUser();
@@ -1960,16 +1983,19 @@ bool GatherUniqueStorageUses::visitUse(Operand *use, AccessUseType useTy) {
     case SILArgumentConvention::Indirect_Inout:
     case SILArgumentConvention::Indirect_InoutAliasable:
     case SILArgumentConvention::Indirect_Out:
-      return visitor.visitStore(use);
+      return visitApplyOperand(use, visitor,
+                               &UniqueStorageUseVisitor::visitStore);
     case SILArgumentConvention::Indirect_In_Guaranteed:
     case SILArgumentConvention::Indirect_In:
     case SILArgumentConvention::Indirect_In_Constant:
-      return visitor.visitLoad(use);
+      return visitApplyOperand(use, visitor,
+                               &UniqueStorageUseVisitor::visitLoad);
     case SILArgumentConvention::Direct_Unowned:
     case SILArgumentConvention::Direct_Owned:
     case SILArgumentConvention::Direct_Guaranteed:
       // most likely an escape of a box
-      return visitor.visitUnknownUse(use);
+      return visitApplyOperand(use, visitor,
+                               &UniqueStorageUseVisitor::visitUnknownUse);
     }
   }
   switch (user->getKind()) {

test/SILOptimizer/hoist_destroy_addr.sil

@@ -59,6 +59,7 @@ sil @f_out : $@convention(thin) <T> () -> @out T
 sil @f_bool : $@convention(thin) () -> Builtin.Int1
 sil [ossa] @take_trivial_struct : $@convention(thin) (TrivialStruct) -> ()
 sil [ossa] @get_change_out : $@convention(thin) () -> @out Change
+sil [ossa] @coro : $@yield_once @convention(thin) (@inout X) -> @yields ()
 
 // CHECK-LABEL: sil [ossa] @test_simple
 // CHECK:      bb0(%0 : $*S):
@@ -346,6 +347,24 @@ entry(%in_addr : $*X):
   return %retval : $X
 }
 
+// If a begin_apply uses the address, the end_appy and abort_apply should be
+// regarded as uses too.  Don't hoist over them.
+//
+// CHECK-LABEL: sil [ossa] @nohoist_over_end_apply_use : {{.*}} {
+// CHECK:         end_apply
+// CHECK:         destroy_addr
+// CHECK:         tuple
+// CHECK-LABEL: } // end sil function 'nohoist_over_end_apply_use'
+sil [ossa] @nohoist_over_end_apply_use : $@convention(thin) (@inout X, @owned X) -> () {
+entry(%addr : $*X, %instance : @owned $X):
+  %coro = function_ref @coro : $@yield_once @convention(thin) (@inout X) -> @yields ()
+  (%empty, %continuation) = begin_apply %coro(%addr) : $@yield_once @convention(thin) (@inout X) -> @yields ()
+  end_apply %continuation
+  %retval = tuple ()
+  store %instance to [assign] %addr : $*X
+  return %retval : $()
+}
+
 // Fold destroy_addr and a load [copy] into a load [take] even when that
 // load [take] is guarded by an access scope.
 //