[sil-combine] Update convert_function canonicalization for ownership.

Some notes:

1. I moved the identity round-trip case to InstSimplify since that is where
   optimizations like that are.

2. I did not update in this commit the code that eliminates convert_function
   when it is only destroyed. In a subsequent commit I am going to implement
   that in a general way and apply it to all forwarding instructions.

3. I implemented eliminating convert_function with ownership only uses in a
   utility so that I can reuse it for other similar optimizations in SILCombine.

Author: gottesmm

include/swift/SIL/ApplySite.h

@@ -116,6 +116,9 @@ class ApplySite {
 
   ApplySiteKind getKind() const { return ApplySiteKind(Inst->getKind()); }
 
+  SILInstruction *operator*() const { return Inst; }
+  SILInstruction *operator->() const { return Inst; }
+
   explicit operator bool() const { return Inst != nullptr; }
 
   SILInstruction *getInstruction() const { return Inst; }

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -705,6 +705,11 @@ SILValue
 makeNewValueAvailable(SILValue value, SILBasicBlock *inBlock,
                       JointPostDominanceSetComputer *jointPostDomComputer);
 
+/// Given a forwarding instruction, eliminate it if all of its users are debug
+/// instructions and ownership uses.
+bool tryEliminateOnlyOwnershipUsedForwardingInst(
+    SingleValueInstruction *forwardingInst, InstModCallbacks &callbacks);
+
 } // end namespace swift
 
 #endif

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -19,6 +19,7 @@
 #ifndef SWIFT_SILOPTIMIZER_UTILS_OWNERSHIPOPTUTILS_H
 #define SWIFT_SILOPTIMIZER_UTILS_OWNERSHIPOPTUTILS_H
 
+#include "swift/Basic/Defer.h"
 #include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/SILModule.h"
@@ -76,6 +77,17 @@ struct OwnershipFixupContext {
     extraAddressFixupInfo.intPtrOp = InteriorPointerOperand();
   }
 
+  /// Gets access to the joint post dominance computer and clears it after \p
+  /// callback.
+  template <typename ResultTy>
+  ResultTy withJointPostDomComputer(
+      function_ref<ResultTy(JointPostDominanceSetComputer &)> callback) {
+    // Make sure we clear the joint post dom computer after callback.
+    SWIFT_DEFER { jointPostDomSetComputer.clear(); };
+    // Then return callback passing in the computer.
+    return callback(jointPostDomSetComputer);
+  }
+
 private:
   /// Helper method called to determine if we discovered we needed interior
   /// pointer fixups while simplifying.
@@ -129,6 +141,37 @@ class OwnershipRAUWHelper {
                                              SILValue newValue);
 };
 
+/// A utility composed ontop of OwnershipFixupContext that knows how to replace
+/// a single use of a value with another value with a different ownership. We
+/// allow for the values to have different types.
+///
+/// NOTE: When not in OSSA, this just performs a normal set use, so this code is
+/// safe to use with all code.
+class OwnershipReplaceSingleUseHelper {
+  OwnershipFixupContext *ctx;
+  Operand *use;
+  SILValue newValue;
+
+public:
+  OwnershipReplaceSingleUseHelper()
+      : ctx(nullptr), use(nullptr), newValue(nullptr) {}
+
+  /// Return an instance of this class if we support replacing \p use->get()
+  /// with \p newValue.
+  ///
+  /// NOTE: For now we only support objects, not addresses so addresses will
+  /// always yield an invalid helper.
+  OwnershipReplaceSingleUseHelper(OwnershipFixupContext &ctx, Operand *use,
+                                  SILValue newValue);
+
+  /// Returns true if this helper was initialized into a valid state.
+  operator bool() const { return isValid(); }
+  bool isValid() const { return bool(ctx) && bool(use) && bool(newValue); }
+
+  /// Perform the actual RAUW.
+  SILBasicBlock::iterator perform();
+};
+
 } // namespace swift
 
 #endif

lib/SILOptimizer/Analysis/SimplifyInstruction.cpp

@@ -70,6 +70,7 @@ namespace {
     SILValue visitPointerToThinFunctionInst(PointerToThinFunctionInst *PTTFI);
     SILValue visitBeginAccessInst(BeginAccessInst *BAI);
     SILValue visitMetatypeInst(MetatypeInst *MTI);
+    SILValue visitConvertFunctionInst(ConvertFunctionInst *cfi);
 
     SILValue simplifyOverflowBuiltin(BuiltinInst *BI);
   };
@@ -482,6 +483,19 @@ SILValue InstSimplifier::visitBeginAccessInst(BeginAccessInst *BAI) {
   return SILValue();
 }
 
+SILValue InstSimplifier::visitConvertFunctionInst(ConvertFunctionInst *cfi) {
+  // Eliminate round trip convert_function. Non round-trip is performed in
+  // SILCombine.
+  //
+  // (convert_function Y->X (convert_function x X->Y)) -> x
+  SILValue convertedValue = lookThroughOwnershipInsts(cfi->getConverted());
+  if (auto *subCFI = dyn_cast<ConvertFunctionInst>(convertedValue))
+    if (subCFI->getConverted()->getType() == cfi->getType())
+      return lookThroughOwnershipInsts(subCFI->getConverted());
+
+  return SILValue();
+}
+
 SILValue InstSimplifier::visitMetatypeInst(MetatypeInst *MI) {
   auto metaType = MI->getType().castTo<MetatypeType>();
   auto instanceType = metaType.getInstanceType();

lib/SILOptimizer/SILCombiner/SILCombinerCastVisitors.cpp

@@ -978,10 +978,8 @@ SILInstruction *SILCombiner::visitConvertEscapeToNoEscapeInst(
       SILType::getPrimitiveObjectType(NewTy));
 }
 
-SILInstruction *SILCombiner::visitConvertFunctionInst(ConvertFunctionInst *CFI) {
-  if (CFI->getFunction()->hasOwnership())
-    return nullptr;
-
+SILInstruction *
+SILCombiner::visitConvertFunctionInst(ConvertFunctionInst *cfi) {
   // If this conversion only changes substitutions, then rewrite applications
   // of the converted function as applications of the original.
   //
@@ -990,84 +988,124 @@ SILInstruction *SILCombiner::visitConvertFunctionInst(ConvertFunctionInst *CFI)
   //
   // TODO: We could generalize this to handle other ABI-compatible cases, by
   // inserting the necessary casts around the arguments.
-  if (CFI->onlyConvertsSubstitutions()) {
-    auto usei = CFI->use_begin();
-    while (usei != CFI->use_end()) {
-      auto use = *usei++;
-      auto user = use->getUser();
-      if (isa<ApplySite>(user) && use->getOperandNumber() == 0) {
-        auto applySite = ApplySite(user);
-        // If this is a partial_apply, insert a convert_function back to the
-        // original result type.
-
-        if (auto pa = dyn_cast<PartialApplyInst>(user)) {
-          auto partialApplyTy = pa->getType();
-          Builder.setInsertionPoint(std::next(pa->getIterator()));
-          
-          SmallVector<SILValue, 4> args(pa->getArguments().begin(),
-                                        pa->getArguments().end());
-          
-          auto newPA = Builder.createPartialApply(pa->getLoc(),
-                                  CFI->getConverted(),
-                                  pa->getSubstitutionMap(),
-                                  args,
-                                  pa->getFunctionType()->getCalleeConvention());
-          auto newConvert = Builder.createConvertFunction(pa->getLoc(),
-                                                          newPA, partialApplyTy,
-                                                          false);
-          pa->replaceAllUsesWith(newConvert);
-          eraseInstFromFunction(*pa);
-          
-          continue;
-        }
-        
+  if (cfi->onlyConvertsSubstitutions()) {
+    SmallVector<Operand *, 32> worklist(cfi->getUses());
+    while (!worklist.empty()) {
+      auto *use = worklist.pop_back_val();
+      auto *user = use->getUser();
+
+      // Look through begin_borrow and copy_value.
+      if (isa<BeginBorrowInst>(user) || isa<CopyValueInst>(user)) {
+        for (auto result : user->getResults())
+          for (auto *resultUse : result->getUses())
+            worklist.push_back(resultUse);
+        continue;
+      }
+
+      if (!isa<ApplySite>(user) || use->getOperandNumber() != 0)
+        continue;
+
+      if (auto fas = FullApplySite::isa(user)) {
         // For full apply sites, we only need to replace the `convert_function`
         // with the original value.
-        use->set(CFI->getConverted());
-        applySite.setSubstCalleeType(
-                      CFI->getConverted()->getType().castTo<SILFunctionType>());
+        //
+        // OWNERSHIP DISCUSSION: We know that cfi is forwarding, so we know that
+        // if cfi is not owned, then we know that cfi->getConverted() must be
+        // valid at applySite and also that the applySite does not consume a
+        // value. In such a case, just perform the change and continue.
+        SILValue newValue = cfi->getConverted();
+        if (newValue.getOwnershipKind() != OwnershipKind::Owned &&
+            newValue.getOwnershipKind() != OwnershipKind::Guaranteed) {
+          instModCallbacks.setUseValue(use, newValue);
+          fas.setSubstCalleeType(newValue->getType().castTo<SILFunctionType>());
+          continue;
+        }
+
+        // Otherwise, we need to use the OwnershipReplaceSingleUseHelper since
+        // we have been looking through ownership forwarding insts and newValue
+        // may be a value with a different lifetime from our original value
+        // beyond the initial base value.
+        OwnershipReplaceSingleUseHelper helper(ownershipFixupContext, use,
+                                               newValue);
+        if (!helper)
+          continue;
+        helper.perform();
+        fas.setSubstCalleeType(newValue->getType().castTo<SILFunctionType>());
+        continue;
+      }
+
+      // If this is a partial_apply, insert a convert_function back to the
+      // original result type.
+      auto *pa = dyn_cast<PartialApplyInst>(user);
+      if (!pa)
+        continue;
+
+      auto partialApplyTy = pa->getType();
+      if (!hasOwnership()) {
+        SILBuilderWithScope localBuilder(std::next(pa->getIterator()), Builder);
+        SmallVector<SILValue, 4> args(pa->getArguments().begin(),
+                                      pa->getArguments().end());
+
+        auto newPA = Builder.createPartialApply(
+            pa->getLoc(), cfi->getConverted(), pa->getSubstitutionMap(), args,
+            pa->getFunctionType()->getCalleeConvention());
+        auto newConvert = Builder.createConvertFunction(pa->getLoc(), newPA,
+                                                        partialApplyTy, false);
+        replaceInstUsesWith(*pa, newConvert);
+        eraseInstFromFunction(*pa);
+        continue;
       }
+
+      OwnershipRAUWHelper helper(ownershipFixupContext, pa,
+                                 cfi->getConverted());
+      if (!helper)
+        continue;
+      SmallVector<SILValue, 4> args(pa->getArguments().begin(),
+                                    pa->getArguments().end());
+      auto newValue = withJointPostDomComputer<SILValue>([&](auto &j) {
+        return makeCopiedValueAvailable(cfi->getConverted(), pa->getParent(),
+                                        &j);
+      });
+
+      SILBuilderWithScope localBuilder(std::next(pa->getIterator()), Builder);
+      auto *newPA = localBuilder.createPartialApply(
+          pa->getLoc(), newValue, pa->getSubstitutionMap(), args,
+          pa->getFunctionType()->getCalleeConvention());
+      if (!use->isLifetimeEnding()) {
+        localBuilder.emitDestroyValueOperation(pa->getLoc(), newValue);
+      }
+      auto *newConvert = localBuilder.createConvertFunction(
+          pa->getLoc(), newPA, partialApplyTy, false);
+      // 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);
     }
   }
-  
+
   // (convert_function (convert_function x)) => (convert_function x)
-  if (auto subCFI = dyn_cast<ConvertFunctionInst>(CFI->getConverted())) {
-    // If we convert the function type back to itself, we can replace the
-    // conversion completely.
-    if (subCFI->getConverted()->getType() == CFI->getType()) {
-      CFI->replaceAllUsesWith(subCFI->getConverted());
-      eraseInstFromFunction(*CFI);
-      return nullptr;
+  if (auto *subCFI = dyn_cast<ConvertFunctionInst>(cfi->getConverted())) {
+    // We handle the case of an identity conversion in inst simplify, so if we
+    // see this pattern then we know that we don't have a round trip and thus
+    // should just bypass the intermediate conversion.
+    if (cfi->getOwnershipKind() != OwnershipKind::Owned) {
+      cfi->getOperandRef().set(subCFI->getConverted());
+      // Return cfi to show we changed it.
+      return cfi;
     }
-    
-    // Otherwise, we can still bypass the intermediate conversion.
-    CFI->getOperandRef().set(subCFI->getConverted());
-  }
-  
-  // Replace a convert_function that only has refcounting uses with its
-  // operand.
-  auto anyNonRefCountUse =
-    std::any_of(CFI->use_begin(),
-                CFI->use_end(),
-                [](Operand *Use) {
-                  return !isa<RefCountingInst>(Use->getUser());
-                });
-
-  if (anyNonRefCountUse)
-    return nullptr;
 
-  // Replace all retain/releases on convert_function by retain/releases on
-  // its argument. This is required to preserve the lifetime of its argument,
-  // which could be e.g. a partial_apply instruction capturing some further
-  // arguments.
-  auto Converted = CFI->getConverted();
-  while (!CFI->use_empty()) {
-    auto *Use = *(CFI->use_begin());
-    assert(Use->getUser()->getResults().empty() &&
-           "Did not expect user with a result!");
-    Use->set(Converted);
+    // If we have an owned value, we can only perform this optimization if the
+    // convert_function is in the same block to ensure that we know we will
+    // eliminate the convert_function. Otherwise we may be breaking up a
+    // forwarding chain in favor of additional ARC traffic which isn't
+    // canonical.
+    SingleBlockOwnedForwardingInstFolder folder(*this, cfi);
+    if (folder.add(subCFI))
+      return std::move(folder).optimizeWithSetValue(subCFI->getConverted());
   }
 
-  eraseInstFromFunction(*CFI);
+  // Replace a convert_function that only has refcounting uses with its
+  // operand.
+  tryEliminateOnlyOwnershipUsedForwardingInst(cfi, instModCallbacks);
   return nullptr;
 }

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -2024,3 +2024,40 @@ SILValue swift::makeNewValueAvailable(
 
   return controlEqCopy ? controlEqCopy : value;
 }
+
+bool swift::tryEliminateOnlyOwnershipUsedForwardingInst(
+    SingleValueInstruction *forwardingInst, InstModCallbacks &callbacks) {
+  if (!OwnershipForwardingMixin::isa(forwardingInst) ||
+      isa<AllArgOwnershipForwardingSingleValueInst>(forwardingInst))
+    return false;
+
+  SmallVector<Operand *, 32> worklist(getNonDebugUses(forwardingInst));
+  while (!worklist.empty()) {
+    auto *use = worklist.pop_back_val();
+    auto *user = use->getUser();
+
+    if (isa<EndBorrowInst>(user) || isa<DestroyValueInst>(user) ||
+        isa<RefCountingInst>(user))
+      continue;
+
+    if (isa<CopyValueInst>(user) || isa<BeginBorrowInst>(user)) {
+      for (auto result : user->getResults())
+        for (auto *resultUse : getNonDebugUses(result))
+          worklist.push_back(resultUse);
+      continue;
+    }
+
+    return false;
+  }
+
+  // Now that we know we can perform our transform, set all uses of
+  // forwardingInst to be used of its operand and then delete \p forwardingInst.
+  auto newValue = forwardingInst->getOperand(0);
+  while (!forwardingInst->use_empty()) {
+    auto *use = *(forwardingInst->use_begin());
+    use->set(newValue);
+  }
+
+  callbacks.deleteInst(forwardingInst);
+  return true;
+}