Merge pull request #62442 from nate-chandler/opaque-values/1/20221201

[AddressLowering] Handle indirectly yielded values used outside of coroutine range.

Author: nate-chandler

include/swift/SIL/MemAccessUtils.h

@@ -206,7 +206,7 @@ SILValue findOwnershipReferenceRoot(SILValue ref);
 
 /// Look through all ownership forwarding instructions to find the values which
 /// were originally borrowed.
-void findGuaranteedReferenceRoots(SILValue value,
+void findGuaranteedReferenceRoots(SILValue value, bool lookThroughNestedBorrows,
                                   SmallVectorImpl<SILValue> &roots);
 
 /// Find the aggregate containing the first owned root of the

include/swift/SIL/SILArgumentConvention.h

@@ -94,12 +94,12 @@ struct SILArgumentConvention {
   bool isOwnedConvention() const {
     switch (Value) {
     case SILArgumentConvention::Indirect_In:
+    case SILArgumentConvention::Indirect_In_Constant:
     case SILArgumentConvention::Direct_Owned:
       return true;
     case SILArgumentConvention::Indirect_In_Guaranteed:
     case SILArgumentConvention::Direct_Guaranteed:
     case SILArgumentConvention::Indirect_Inout:
-    case SILArgumentConvention::Indirect_In_Constant:
     case SILArgumentConvention::Indirect_Out:
     case SILArgumentConvention::Indirect_InoutAliasable:
     case SILArgumentConvention::Direct_Unowned:

lib/SIL/Utils/MemAccessUtils.cpp

@@ -861,6 +861,7 @@ SILValue swift::findOwnershipReferenceRoot(SILValue ref) {
 }
 
 void swift::findGuaranteedReferenceRoots(SILValue value,
+                                         bool lookThroughNestedBorrows,
                                          SmallVectorImpl<SILValue> &roots) {
   GraphNodeWorklist<SILValue, 4> worklist;
   auto addAllOperandsToWorklist = [&worklist](SILInstruction *inst) -> bool {
@@ -882,7 +883,16 @@ void swift::findGuaranteedReferenceRoots(SILValue value,
         }
       }
     } else if (auto *inst = value->getDefiningInstruction()) {
-      if (auto *result =
+      if (auto *bbi = dyn_cast<BeginBorrowInst>(inst)) {
+        auto borrowee = bbi->getOperand();
+        if (lookThroughNestedBorrows &&
+            borrowee->getOwnershipKind() == OwnershipKind::Guaranteed) {
+          // A nested borrow, the root guaranteed earlier in the use-def chain.
+          worklist.insert(borrowee);
+        }
+        // The borrowee isn't guaranteed or we aren't looking through nested
+        // borrows.  Fall through to add the begin_borrow to roots.
+      } else if (auto *result =
               dyn_cast<FirstArgOwnershipForwardingSingleValueInst>(inst)) {
         if (result->getNumOperands() > 0) {
           worklist.insert(result->getOperand(0));

lib/SIL/Utils/MemoryLocations.cpp

@@ -173,6 +173,16 @@ void MemoryLocations::analyzeLocations(SILFunction *function) {
           }
         }
       }
+      if (auto *BAI = dyn_cast<BeginApplyInst>(&I)) {
+        auto convention = BAI->getSubstCalleeConv();
+        auto yields = convention.getYields();
+        auto yieldedValues = BAI->getYieldedValues();
+        for (auto index : indices(yields)) {
+          if (convention.isSILIndirect(yields[index])) {
+            analyzeLocation(yieldedValues[index]);
+          }
+        }
+      }
     }
   }
 }

lib/SIL/Verifier/MemoryLifetimeVerifier.cpp

@@ -401,6 +401,39 @@ void MemoryLifetimeVerifier::initDataflowInBlock(SILBasicBlock *block,
         }
         break;
       }
+      case SILInstructionKind::BeginApplyInst: {
+        auto *BAI = cast<BeginApplyInst>(&I);
+        auto yieldedValues = BAI->getYieldedValues();
+        for (auto index : indices(yieldedValues)) {
+          auto fnType = BAI->getSubstCalleeType();
+          SILArgumentConvention argConv(
+              fnType->getYields()[index].getConvention());
+          if (argConv.isIndirectConvention()) {
+            genBits(state, yieldedValues[index]);
+          }
+        }
+        break;
+      }
+      case SILInstructionKind::EndApplyInst:
+      case SILInstructionKind::AbortApplyInst: {
+        auto *BAI = [&]() {
+          if (auto *EAI = dyn_cast<EndApplyInst>(&I)) {
+            return EAI->getBeginApply();
+          }
+          auto *AAI = dyn_cast<AbortApplyInst>(&I);
+          return AAI->getBeginApply();
+        }();
+        auto yieldedValues = BAI->getYieldedValues();
+        for (auto index : indices(yieldedValues)) {
+          auto fnType = BAI->getSubstCalleeType();
+          SILArgumentConvention argConv(
+              fnType->getYields()[index].getConvention());
+          if (argConv.isIndirectConvention()) {
+            killBits(state, yieldedValues[index]);
+          }
+        }
+        break;
+      }
       case SILInstructionKind::YieldInst: {
         auto *YI = cast<YieldInst>(&I);
         for (Operand &op : YI->getAllOperands()) {
@@ -694,6 +727,48 @@ void MemoryLifetimeVerifier::checkBlock(SILBasicBlock *block, Bits &bits) {
         }
         break;
       }
+      case SILInstructionKind::BeginApplyInst: {
+        auto *BAI = cast<BeginApplyInst>(&I);
+        auto yieldedValues = BAI->getYieldedValues();
+        for (auto index : indices(yieldedValues)) {
+          auto fnType = BAI->getSubstCalleeType();
+          SILArgumentConvention argConv(
+              fnType->getYields()[index].getConvention());
+          if (argConv.isIndirectConvention()) {
+            requireBitsClear(bits, yieldedValues[index], &I);
+            locations.setBits(bits, yieldedValues[index]);
+          }
+        }
+        break;
+      }
+      case SILInstructionKind::EndApplyInst:
+      case SILInstructionKind::AbortApplyInst: {
+        auto *BAI = [&]() {
+          if (auto *EAI = dyn_cast<EndApplyInst>(&I)) {
+            return EAI->getBeginApply();
+          }
+          auto *AAI = dyn_cast<AbortApplyInst>(&I);
+          return AAI->getBeginApply();
+        }();
+        auto yieldedValues = BAI->getYieldedValues();
+        for (auto index : indices(yieldedValues)) {
+          auto fnType = BAI->getSubstCalleeType();
+          SILArgumentConvention argConv(
+              fnType->getYields()[index].getConvention());
+          if (argConv.isIndirectConvention()) {
+            if (argConv.isInoutConvention() ||
+                argConv.isGuaranteedConvention()) {
+              requireBitsSet(bits | ~nonTrivialLocations, yieldedValues[index],
+                             &I);
+            } else if (argConv.isOwnedConvention()) {
+              requireBitsClear(bits & nonTrivialLocations, yieldedValues[index],
+                               &I);
+            }
+            locations.clearBits(bits, yieldedValues[index]);
+          }
+        }
+        break;
+      }
       case SILInstructionKind::YieldInst: {
         auto *YI = cast<YieldInst>(&I);
         for (Operand &op : YI->getAllOperands()) {

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"
@@ -157,6 +158,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"
@@ -345,6 +347,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);
 }
@@ -933,8 +951,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;
@@ -2230,9 +2259,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());
@@ -3271,7 +3323,7 @@ static void emitEndBorrows(SILValue value, AddressLoweringState &pass) {
   SSAPrunedLiveness liveness(&discoveredBlocks);
   liveness.initializeDef(value);
   for (auto *use : usePoints) {
-    assert(!use->isLifetimeEnding());
+    assert(!use->isLifetimeEnding() || isa<EndBorrowInst>(use->getUser()));
     liveness.updateForUse(use->getUser(), /*lifetimeEnding*/ false);
   }
   PrunedLivenessBoundary guaranteedBoundary;

lib/SILOptimizer/Transforms/DeadCodeElimination.cpp

@@ -292,7 +292,9 @@ void DCE::markLive() {
           // Visit the end_borrows of all the borrow scopes that this
           // begin_borrow could be borrowing.
           SmallVector<SILValue, 4> roots;
-          findGuaranteedReferenceRoots(borrowInst->getOperand(), roots);
+          findGuaranteedReferenceRoots(borrowInst->getOperand(),
+                                       /*lookThroughNestedBorrows=*/false,
+                                       roots);
           for (auto root : roots) {
             visitTransitiveEndBorrows(root,
                                       [&](EndBorrowInst *endBorrow) {

test/SIL/memory_lifetime.sil

@@ -683,3 +683,46 @@ bb3:
   return %14 : $()
 }
 
+sil [ossa] @begin_apply_in_destroy : $@convention(thin) () -> () {
+entry:
+  (%instance, %token) = begin_apply undef<Inner>() : $@yield_once @convention(thin) <U> () -> @yields @in U
+  destroy_addr %instance : $*Inner
+  end_apply %token
+  %retval = tuple ()
+  return %retval : $()
+}
+
+sil [ossa] @begin_apply_in_constant_destroy : $@convention(thin) () -> () {
+entry:
+  (%instance, %token) = begin_apply undef<Inner>() : $@yield_once @convention(thin) <U> () -> @yields @in_constant U
+  destroy_addr %instance : $*Inner
+  end_apply %token
+  %retval = tuple ()
+  return %retval : $()
+}
+
+sil [ossa] @begin_apply_in_guaranteed_no_destroy : $@convention(thin) () -> () {
+entry:
+  (%instance, %token) = begin_apply undef<Inner>() : $@yield_once @convention(thin) <U> () -> @yields @in_guaranteed U
+  end_apply %token
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// CHECK: SIL memory lifetime failure in @begin_apply_inout_destroy: memory is not initialized, but should be
+sil [ossa] @begin_apply_inout_no_destroy : $@convention(thin) () -> () {
+entry:
+  (%instance, %token) = begin_apply undef<Inner>() : $@yield_once @convention(thin) <U> () -> @yields @inout U
+  end_apply %token
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// CHECK: SIL memory lifetime failure in @begin_apply_inout_aliasable_destroy: memory is not initialized, but should be
+sil [ossa] @begin_apply_inout_aliasable_no_destroy : $@convention(thin) () -> () {
+entry:
+  (%instance, %token) = begin_apply undef<Inner>() : $@yield_once @convention(thin) <U> () -> @yields @inout_aliasable U
+  end_apply %token
+  %retval = tuple ()
+  return %retval : $()
+}