Fix SILCombine to use OSSA RAUW prepareReplacement

These SILCombine patterns produce new instructions that use the the
original value after it is updated for OSSA. The RAUW utility needs to
copy/borrow the original value before SILCombine can generate the new
replacement. Then SILCombine can pass that replacement back to RAUW to
perform the final replaceAllUsesWith.

1. PointerToAddress and RawPointerToRef

Fixes bugs where RAUW inserts borrows in the wrong place.

2. UncheckedBitwiseCast

The OSSA RAUW helper is incorrect here for producing a new
UncheckedRefCast. Simply don't use it. Since a bitwise cast is an
"Unowned" pointer escape, simply convert the unchecked_bitwise_cast
opcode to unchecked_ref_cast with Unowned forwarding ownership.

3. ConvertFunction

Make sure RAUW is never called with a replacement value that has
different ownership than the original.

Author: atrick

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();
     }
   }