Merge pull request #39740 from atrick/oou-borrow

OSSA utility incremental redesign and SILCombine fixes

Author: atrick

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -148,11 +148,10 @@ struct OwnershipFixupContext {
   InstModCallbacks &callbacks;
   DeadEndBlocks &deBlocks;
 
-
-  // FIXME: remove these two vectors once BorrowedLifetimeExtender is used
-  // everywhere.
-  SmallVector<Operand *, 8> transitiveBorrowedUses;
-  SmallVector<PhiOperand, 8> recursiveReborrows;
+  // Cache the use-points for the lifetime of an inner guaranteed value (which
+  // does not introduce a borrow scope) after checking validity. These will be
+  // used again to extend the lifetime of the replacement value.
+  SmallVector<Operand *, 8> guaranteedUsePoints;
 
   /// Extra state initialized by OwnershipRAUWFixupHelper::get() that we use
   /// when RAUWing addresses. This ensures we do not need to recompute this
@@ -162,6 +161,8 @@ struct OwnershipFixupContext {
     /// compute all transitive address uses of oldValue. If we find that we do
     /// need this fixed up, then we will copy our interior pointer base value
     /// and use this to seed that new lifetime.
+    ///
+    /// FIXME: shouldn't these already be covered by guaranteedUsePoints?
     SmallVector<Operand *, 8> allAddressUsesFromOldValue;
 
     /// This is the interior pointer (e.g. ref_element_addr)
@@ -184,8 +185,7 @@ struct OwnershipFixupContext {
       : callbacks(callbacks), deBlocks(deBlocks) {}
 
   void clear() {
-    transitiveBorrowedUses.clear();
-    recursiveReborrows.clear();
+    guaranteedUsePoints.clear();
     extraAddressFixupInfo.allAddressUsesFromOldValue.clear();
     extraAddressFixupInfo.base = AccessBase();
   }
@@ -212,6 +212,8 @@ class OwnershipRAUWHelper {
 private:
   OwnershipFixupContext *ctx;
   SILValue oldValue;
+  // newValue is the aspirational replacement. It might not be the actual
+  // replacement after SILCombine fixups (like type casting) and OSSA fixups.
   SILValue newValue;
 
 public:
@@ -253,21 +255,27 @@ class OwnershipRAUWHelper {
   /// Perform OSSA fixup on newValue and return a fixed-up value based that can
   /// be used to replace all uses of oldValue.
   ///
-  /// This is so that we can avoid creating "forwarding" transformation
-  /// instructions before we know if we can perform the RAUW. Any such
-  /// "forwarding" transformation must be performed upon \p newValue at \p
-  /// oldValue's insertion point so that we can then here RAUW the transformed
-  /// \p newValue.
+  /// This is only used by clients that must transform \p newValue, such as
+  /// adding type casts, before it can be used to replace \p oldValue.
+  ///
+  /// \p rewrittenNewValue is only passed when the client needs to regenerate
+  /// newValue after checking its RAUW validity, but before performing OSSA
+  /// fixup on it.
+  SILValue prepareReplacement(SILValue rewrittenNewValue = SILValue());
+
+  /// Perform the actual RAUW--replace all uses if \p oldValue.
+  ///
+  /// Precondition: \p replacementValue is either invalid or has the same type
+  /// as \p oldValue and is a valid OSSA replacement.
   SILBasicBlock::iterator
-  perform(SingleValueInstruction *maybeTransformedNewValue = nullptr);
+  perform(SILValue replacementValue = SILValue());
 
 private:
-  SILBasicBlock::iterator replaceAddressUses(SingleValueInstruction *oldValue,
-                                             SILValue newValue);
-
   void invalidate() {
     ctx = nullptr;
   }
+
+  SILValue getReplacementAddress();
 };
 
 /// A utility composed ontop of OwnershipFixupContext that knows how to replace

lib/SILOptimizer/SILCombiner/SILCombinerCastVisitors.cpp

@@ -237,13 +237,17 @@ visitPointerToAddressInst(PointerToAddressInst *PTAI) {
     // handle issues around interior pointers and expanding borrow scopes.
     if (auto *ATPI = dyn_cast<AddressToPointerInst>(PTAI->getOperand())) {
       if (!hasOwnership()) {
-        return Builder.createUncheckedAddrCast(PTAI->getLoc(), ATPI->getOperand(),
+        return Builder.createUncheckedAddrCast(PTAI->getLoc(),
+                                               ATPI->getOperand(),
                                                PTAI->getType());
       }
 
-      OwnershipRAUWHelper helper(ownershipFixupContext, PTAI, ATPI->getOperand());
+      OwnershipRAUWHelper helper(ownershipFixupContext, PTAI,
+                                 ATPI->getOperand());
       if (helper) {
-        auto *newInst = Builder.createUncheckedAddrCast(PTAI->getLoc(), ATPI->getOperand(),
+        auto replacement = helper.prepareReplacement();
+        auto *newInst = Builder.createUncheckedAddrCast(PTAI->getLoc(),
+                                                        replacement,
                                                         PTAI->getType());
         helper.perform(newInst);
         return nullptr;
@@ -726,8 +730,9 @@ SILCombiner::visitRawPointerToRefInst(RawPointerToRefInst *rawToRef) {
         // Since we are using std::next, we use getAutogeneratedLocation to
         // avoid any issues if our next insertion point is a terminator.
         auto loc = RegularLocation::getAutoGeneratedLocation();
+        auto replacement = helper.prepareReplacement();
         auto *newInst = localBuilder.createUncheckedRefCast(
-            loc, originalRef, rawToRef->getType());
+            loc, replacement, rawToRef->getType());
         // If we have an operand with ownership, we need to change our
         // unchecked_ref_cast to produce an unowned value. This is because
         // otherwise, our unchecked_ref_cast will consume the underlying owned
@@ -800,8 +805,9 @@ visitUncheckedBitwiseCastInst(UncheckedBitwiseCastInst *UBCI) {
 
     OwnershipRAUWHelper helper(ownershipFixupContext, UBCI, Oper);
     if (helper) {
+      auto replacement = helper.prepareReplacement();
       auto *transformedOper = Builder.createUncheckedBitwiseCast(
-          UBCI->getLoc(), Oper, UBCI->getType());
+          UBCI->getLoc(), replacement, UBCI->getType());
       helper.perform(transformedOper);
       return nullptr;
     }
@@ -817,31 +823,21 @@ visitUncheckedBitwiseCastInst(UncheckedBitwiseCastInst *UBCI) {
                            Builder.getModule()))
     return nullptr;
 
-  if (!Builder.hasOwnership()) {
-    return Builder.createUncheckedRefCast(UBCI->getLoc(), UBCI->getOperand(),
-                                          UBCI->getType());
+  // Normally, OwnershipRAUWHelper needs to be called to handle ownership of
+  // UBCI->getOperand(). However, we know that UBCI->getOperand() is already
+  // available at the point of the cast, and by forcing the cast to be Unowned,
+  // we ensure that no ownership adjustment is needed. So we can skip
+  // prepareReplacement completely and just drop in the replacement. That avoids
+  // an extra copy in the case that UBCI->getOperand() is Owned.
+  auto *refCast = Builder.createUncheckedRefCast(
+      UBCI->getLoc(), UBCI->getOperand(), UBCI->getType());
+  if (Builder.hasOwnership()) {
+    // A bitwise cast is always unowned, so we can safely force the reference
+    // cast to forward as unowned and no ownership adjustment is needed.
+    assert(UBCI->getOwnershipKind() == OwnershipKind::Unowned);
+    refCast->setForwardingOwnershipKind(OwnershipKind::Unowned);
   }
-
-  {
-    OwnershipRAUWHelper helper(ownershipFixupContext, UBCI, UBCI->getOperand());
-    if (helper) {
-      auto *newInst = Builder.createUncheckedRefCast(UBCI->getLoc(), UBCI->getOperand(),
-                                                     UBCI->getType());
-      // If we have an operand with owned ownership, we change our
-      // unchecked_ref_cast to explicitly pass the owned operand as an unowned
-      // value. This is because otherwise, we would consume the owned value
-      // creating breaking OSSA. In contrast, if we have a guaranteed value, we
-      // are going to be replacing an UnownedInstantaneousUse with an
-      // InstantaneousUse which is always safe for a guaranteed value.
-      if (newInst->getForwardingOwnershipKind() == OwnershipKind::Owned) {
-        newInst->setForwardingOwnershipKind(OwnershipKind::Unowned);
-      }
-      helper.perform(newInst);
-      return nullptr;
-    }
-  }
-
-  return nullptr;
+  return refCast;
 }
 
 SILInstruction *
@@ -1064,10 +1060,11 @@ SILCombiner::visitConvertFunctionInst(ConvertFunctionInst *cfi) {
         continue;
       }
 
-      OwnershipRAUWHelper helper(ownershipFixupContext, pa,
-                                 cfi->getConverted());
-      if (!helper)
+      OwnershipRAUWHelper checkRAUW(ownershipFixupContext, pa,
+                                    cfi->getConverted());
+      if (!checkRAUW)
         continue;
+
       SmallVector<SILValue, 4> args(pa->getArguments().begin(),
                                     pa->getArguments().end());
       auto newValue = makeCopiedValueAvailable(cfi->getConverted(),
@@ -1085,7 +1082,12 @@ SILCombiner::visitConvertFunctionInst(ConvertFunctionInst *cfi) {
       // We need to end the lifetime of the convert_function/partial_apply since
       // the helper assumes that ossa is correct upon input.
       localBuilder.emitDestroyValueOperation(pa->getLoc(), newConvert);
-      helper.perform(newConvert);
+      // 'newConvert' may have different ownership than then 'cfi'. newConvert
+      // is always owned, while 'cfi' may have been guaranteed. OSSA-RAUW
+      // validity depends on the ownership kind. Reinstantiate
+      // OwnershipRAUWHelper to verify that it is still valid
+      // (a very fast check in this case).
+      OwnershipRAUWHelper(ownershipFixupContext, pa, newConvert).perform();
     }
   }