[AddressLowering] Fix indirect coro yield handling

The begin_apply instruction introduces storage which is only valid until
the coroutine is ended via end_apply/abort_apply.  Previously,
AddressLowering assumed the memory was valid "as long as it needed to
be": either as an argument to the function being lowered (so valid for
the whole lifetime) or as a new alloc_stack (whose lifetime is
determined by AddressLowering itself).  The result was that the storage
made available by the begin_apply was used after the
end_apply/abort_apply when there were uses of [a copy of] [a projection
of] the yield after the end_apply/abort_apply.

Here, the behavior is fixed.  Now, the storage is used "as soon as
possible".  There are two cases:

(1) If an indirect value's ownership is transferred into the
caller (i.e. its convention is `@in` or `@in_constant`), the value is
`copy_addr [take]`n into caller-local storage when lowering the
begin_apply.

(2) If an indirect value's ownership is only lent to the caller (i.e.
its convention is `@in_guaranteed`), no equivalent operation could be
done: there's no `copy_addr [borrow]`; on the other hand, there's no
need to do it--uses of the value must have occurred within the
coroutine's range.  Instead, it's necessary to disable the store-copy
optimization in this case: storage must be allocated for a copy of [a
projection of] the yielded value.

Author: nate-chandler

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -141,6 +141,7 @@
 #include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/DebugUtils.h"
 #include "swift/SIL/DynamicCasts.h"
+#include "swift/SIL/MemAccessUtils.h"
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/PrettyStackTrace.h"
 #include "swift/SIL/PrunedLiveness.h"
@@ -156,6 +157,7 @@
 #include "swift/SILOptimizer/Utils/InstructionDeleter.h"
 #include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
@@ -344,6 +346,22 @@ static bool isStoreCopy(SILValue value) {
   if (!copyInst->hasOneUse())
     return false;
 
+  auto source = copyInst->getOperand();
+  if (source->getOwnershipKind() == OwnershipKind::Guaranteed) {
+    // [in_guaranteed_begin_apply_results] If any root of the source is a
+    // begin_apply, we can't rely on projecting from the (rewritten) source:
+    // The store may not be in the coroutine's range. The result would be
+    // attempting to access invalid storage.
+    SmallVector<SILValue, 4> roots;
+    findGuaranteedReferenceRoots(source, /*lookThroughNestedBorrows=*/true,
+                                 roots);
+    if (llvm::any_of(roots, [](SILValue root) {
+          return isa<BeginApplyInst>(root->getDefiningInstruction());
+        })) {
+      return false;
+    }
+  }
+
   auto *user = value->getSingleUse()->getUser();
   return isa<StoreInst>(user);
 }
@@ -928,8 +946,19 @@ static bool doesNotNeedStackAllocation(SILValue value) {
   auto *defInst = value->getDefiningInstruction();
   if (!defInst)
     return false;
-
-  if (isa<LoadBorrowInst>(defInst) || isa<BeginApplyInst>(defInst))
+  // [in_guaranteed_begin_apply_results] OSSA ensures that every use of a
+  // guaranteed value resulting from a begin_apply will occur in the
+  // coroutine's range (i.e. "before" the end_apply/abort apply).
+  // AddressLowering takes advantage of this lack of uses outside of the
+  // coroutine's range to directly use the storage that is yielded by the
+  // coroutine rather than moving it to local storage.
+  //
+  // It is, however, valid in OSSA to have uses of an owned value produced by a
+  // begin_apply outside of the coroutine range.  So in that case, it is
+  // necessary to introduce new storage and move to it.
+  if (isa<LoadBorrowInst>(defInst) ||
+      (isa<BeginApplyInst>(defInst) &&
+       value.getOwnershipKind() == OwnershipKind::Guaranteed))
     return true;
 
   return false;
@@ -2223,9 +2252,32 @@ void ApplyRewriter::convertBeginApplyWithOpaqueYield() {
       continue;
     }
     if (oldResult.getType().isAddressOnly(*pass.function)) {
-      // Remap storage when an address-only type is yielded as an opaque value.
-      pass.valueStorageMap.setStorageAddress(&oldResult, &newResult);
-      pass.valueStorageMap.getStorage(&oldResult).markRewritten();
+      auto info = newCall->getSubstCalleeConv().getYieldInfoForOperandIndex(i);
+      assert(info.isFormalIndirect());
+      if (info.isConsumed()) {
+        // Because it is legal to have uses of an owned value produced by a
+        // begin_apply after a coroutine's range, AddressLowering must move the
+        // value into local storage so that such out-of-coroutine-range uses can
+        // be rewritten in terms of that address (instead of being rewritten in
+        // terms of the yielded owned storage which is no longer valid beyond
+        // the coroutine's range).
+        auto &storage = pass.valueStorageMap.getStorage(&oldResult);
+        auto destAddr = addrMat.materializeAddress(&oldResult);
+        storage.storageAddress = destAddr;
+        storage.markRewritten();
+        resultBuilder.createCopyAddr(callLoc, &newResult, destAddr,
+                                     info.isConsumed() ? IsTake : IsNotTake,
+                                     IsInitialization);
+      } else {
+        // [in_guaranteed_begin_apply_results] Because OSSA ensure that all uses
+        // of a guaranteed value produced by a begin_apply are used within the
+        // coroutine's range, AddressLowering will not introduce uses of
+        // invalid memory by rewriting the uses of a yielded guaranteed opaque
+        // value as uses of yielded guaranteed storage.  However, it must
+        // allocate storage for copies of [projections of] such values.
+        pass.valueStorageMap.setStorageAddress(&oldResult, &newResult);
+        pass.valueStorageMap.getStorage(&oldResult).markRewritten();
+      }
       continue;
     }
     assert(oldResult.getType().isObject());

test/SILOptimizer/address_lowering.sil

@@ -1648,8 +1648,8 @@ entry:
 // CHECK:       {{bb[0-9]+}}({{%[^,]+}} : $*Klass, [[OPAQUE_OUT:%[^,]+]] : $*T):
 // CHECK:         ([[ADDR_1:%[^,]+]], [[ADDR_2:%[^,]+]], {{%[^,]+}}) = begin_apply
 // CHECK:         [[INSTANCE_1:%[^,]+]] = load [take] [[ADDR_1]]
-// CHECK:         begin_borrow [[INSTANCE_1]]
 // CHECK:         copy_addr [take] [[ADDR_2]] to [init] [[OPAQUE_OUT]]
+// CHECK:         begin_borrow [[INSTANCE_1]]
 // CHECK-LABEL: } // end sil function 'testBeginApplyCYield1LoadableNontrivialOwned1OpaqueOwned'
 sil [ossa] @testBeginApplyCYield1LoadableNontrivialOwned1OpaqueOwned : $@convention(thin) <T> () -> (@out Klass, @out T) {
 entry:
@@ -1710,6 +1710,123 @@ entry:
   return %retval : $T
 }
 
+// Verify that a copy_value of a @in_guaranteed yield generates storage.
+// CHECK-LABEL: sil [ossa] @testBeginApplyGCopyConsumeInGuaranteedValue : {{.*}} {
+// CHECK:       bb0:
+// CHECK:         [[COPY_STORAGE:%[^,]+]] = alloc_stack $T                             
+// CHECK:         ([[YIELD_STORAGE:%[^,]+]], [[TOKEN:%[^,]+]]) = begin_apply
+// CHECK:         copy_addr [[YIELD_STORAGE]] to [init] [[COPY_STORAGE]]
+// CHECK:         end_apply [[TOKEN]]                                    
+// CHECK:         apply undef<T>([[COPY_STORAGE]])
+// CHECK:         dealloc_stack [[COPY_STORAGE]]
+// CHECK-LABEL: } // end sil function 'testBeginApplyGCopyConsumeInGuaranteedValue'
+sil [ossa] @testBeginApplyGCopyConsumeInGuaranteedValue : $@convention(thin) <T> () -> () {
+entry:
+  (%yield, %token) = begin_apply undef<T>() : $@yield_once @convention(thin) <τ_0_0> () -> @yields @in_guaranteed τ_0_0
+  %copy = copy_value %yield : $T
+  end_apply %token
+  apply undef<T>(%copy) : $@convention(thin) <T> (@in T) -> ()
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// Verify that a copy_value that is a "copy-store" whose source is an
+// @in_guaranteed storage generates storage.
+// CHECK-LABEL: sil [ossa] @testBeginApplyH1CopyStoreInGuaranteedValue : {{.*}} {
+// CHECK:         [[COPY_STORAGE:%[^,]+]] = alloc_stack $T
+// CHECK:         ([[YIELD_STORAGE:%[^,]+]], [[TOKEN:%[^,]+]]) = begin_apply
+// CHECK:         copy_addr [[YIELD_STORAGE]] to [init] [[COPY_STORAGE]]
+// CHECK:         end_apply [[TOKEN]]
+// CHECK:         [[PTR:%[^,]+]] = apply undef
+// CHECK:         [[ADDR:%[^,]+]] = pointer_to_address [[PTR:%[^,]+]]
+// CHECK:         copy_addr [take] [[COPY_STORAGE]] to [[ADDR]]
+// CHECK-LABEL: } // end sil function 'testBeginApplyH1CopyStoreInGuaranteedValue'
+sil [ossa] @testBeginApplyH1CopyStoreInGuaranteedValue : $@convention(thin) <T> () -> () {
+entry:
+  (%yield, %token) = begin_apply undef<T>() : $@yield_once @convention(thin) <τ_0_0> () -> @yields @in_guaranteed τ_0_0
+  %copy = copy_value %yield : $T
+  end_apply %token
+  %ptr = apply undef<T>() : $@convention(method) <τ_0_0> () -> Builtin.RawPointer
+  %addr = pointer_to_address %ptr : $Builtin.RawPointer to [strict] $*T
+  store %copy to [assign] %addr : $*T
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// sil [ossa] @testBeginApplyIConsumeInValue : {{.*}} {
+// bb0:
+//   [[IN_STORAGE:%[^,]+]] = alloc_stack $T
+//   ([[YIELD_STORAGE:%[^,]+]], [[TOKEN:%[^,]+]]) = begin_apply undef<T>()
+//   copy_addr [take] [[YIELD_STORAGE]] to [init] [[IN_STORAGE]]
+//   end_apply [[TOKEN]]
+//   apply undef<T>([[IN_STORAGE]])
+// } // end sil function 'testBeginApplyIConsumeInValue'
+sil [ossa] @testBeginApplyIConsumeInValue : $@convention(thin) <T> () -> () {
+entry:
+  (%yield, %token) = begin_apply undef<T>() : $@yield_once @convention(thin) <τ_0_0> () -> @yields @in τ_0_0
+  end_apply %token
+  apply undef<T>(%yield) : $@convention(thin) <T> (@in T) -> ()
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// sil [ossa] @testBeginApplyJStoreInValue : {{.*}} {
+//   [[IN_STORAGE:%[^,]+]] = alloc_stack $T
+//   ([[YIELD_STORAGE:%[^,]+]], [[TOKEN:%[^,]+]]) = begin_apply
+//   copy_addr [take] [[YIELD_STORAGE]] to [init] [[IN_STORAGE]]
+//   end_apply [[TOKEN]]
+//   [[PTR:%[^,]+]] = apply
+//   [[ADDR:%[^,]+]] = pointer_to_address [[PTR:%[^,]+]]
+//   copy_addr [take] [[IN_STORAGE]] to [[ADDR]]
+// } // end sil function 'testBeginApplyJStoreInValue'
+sil [ossa] @testBeginApplyJStoreInValue : $@convention(thin) <T> () -> () {
+entry:
+  (%yield, %token) = begin_apply undef<T>() : $@yield_once @convention(thin) <τ_0_0> () -> @yields @in τ_0_0
+  end_apply %token
+  %ptr = apply undef<T>() : $@convention(method) <τ_0_0> () -> Builtin.RawPointer
+  %addr = pointer_to_address %ptr : $Builtin.RawPointer to [strict] $*T
+  store %yield to [assign] %addr : $*T
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// sil [ossa] @testBeginApplyKConsumeInConstantValue : {{.*}} {
+// bb0:
+//   [[IN_CONSTANT_STORAGE:%[^,]+]] = alloc_stack $T
+//   ([[YIELD_STORAGE:%[^,]+]], [[TOKEN:%[^,]+]]) = begin_apply undef<T>()
+//   copy_addr [take] [[YIELD_STORAGE]] to [init] [[IN_CONSTANT_STORAGE]]
+//   end_apply [[TOKEN]]
+//   apply undef<T>([[IN_CONSTANT_STORAGE]])
+// } // end sil function 'testBeginApplyKConsumeInConstantValue'
+sil [ossa] @testBeginApplyKConsumeInConstantValue : $@convention(thin) <T> () -> () {
+entry:
+  (%yield, %token) = begin_apply undef<T>() : $@yield_once @convention(thin) <τ_0_0> () -> @yields @in_constant τ_0_0
+  end_apply %token
+  apply undef<T>(%yield) : $@convention(thin) <T> (@in_constant T) -> ()
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// sil [ossa] @testBeginApplyLStoreInConstantValue : {{.*}} {
+//   [[IN_CONSTANT_STORAGE:%[^,]+]] = alloc_stack $T
+//   ([[YIELD_STORAGE:%[^,]+]], [[TOKEN:%[^,]+]]) = begin_apply
+//   copy_addr [take] [[YIELD_STORAGE]] to [init] [[IN_CONSTANT_STORAGE]]
+//   end_apply [[TOKEN]]
+//   [[PTR:%[^,]+]] = apply
+//   [[ADDR:%[^,]+]] = pointer_to_address [[PTR:%[^,]+]]
+//   copy_addr [take] [[IN_CONSTANT_STORAGE]] to [[ADDR]]
+// } // end sil function 'testBeginApplyLStoreInConstantValue'
+sil [ossa] @testBeginApplyLStoreInConstantValue : $@convention(thin) <T> () -> () {
+entry:
+  (%yield, %token) = begin_apply undef<T>() : $@yield_once @convention(thin) <τ_0_0> () -> @yields @in_constant τ_0_0
+  end_apply %token
+  %ptr = apply undef<T>() : $@convention(method) <τ_0_0> () -> Builtin.RawPointer
+  %addr = pointer_to_address %ptr : $Builtin.RawPointer to [strict] $*T
+  store %yield to [assign] %addr : $*T
+  %retval = tuple ()
+  return %retval : $()
+}
+
 // CHECK-LABEL: sil hidden [ossa] @testOpaqueYield :
 // CHECK: bb0(%0 : @guaranteed $TestGeneric<T>):
 // CHECK:   [[REF:%.*]] = ref_element_addr %0 : $TestGeneric<T>, #TestGeneric.borrowedGeneric