Add new apis to ForwardingOperation

include/swift/SIL/InstWrappers.h

@@ -252,44 +252,34 @@ class ForwardingOperation {
   ValueOwnershipKind getForwardingOwnershipKind();
   bool preservesOwnership();
 
-  // FIXME: Find a better name. Even unary instructions like struct_extract
-  // forward "all" operands.
-  bool canForwardAllOperands() const {
+  Operand *getSingleForwardingOperand() const {
     switch (forwardingInst->getKind()) {
     case SILInstructionKind::StructInst:
     case SILInstructionKind::TupleInst:
     case SILInstructionKind::LinearFunctionInst:
     case SILInstructionKind::DifferentiableFunctionInst:
-      return true;
+      return nullptr;
     default:
-      return false;
+      if (forwardingInst->getNumRealOperands() == 0) {
+        // This can happen with enum instructions that have no payload.
+        return nullptr;
+      }
+      return &forwardingInst->getOperandRef(0);
     }
   }
 
-  // FIXME: Find a better name. Even instructions that forward all operands can
-  // forward the first operand.
-  bool canForwardFirstOperandOnly() const {
-    return !canForwardAllOperands() && forwardingInst->getNumRealOperands() > 0;
-  }
-
   ArrayRef<Operand> getForwardedOperands() const {
-    if (canForwardAllOperands()) {
-      return forwardingInst->getAllOperands();
+    if (auto *singleForwardingOp = getSingleForwardingOperand()) {
+      return *singleForwardingOp;
     }
-    if (canForwardFirstOperandOnly()) {
-      return forwardingInst->getOperandRef(0);
-    }
-    return {};
+    return forwardingInst->getAllOperands();
   }
 
   MutableArrayRef<Operand> getForwardedOperands() {
-    if (canForwardAllOperands()) {
-      return forwardingInst->getAllOperands();
-    }
-    if (canForwardFirstOperandOnly()) {
-      return forwardingInst->getOperandRef(0);
+    if (auto *singleForwardingOp = getSingleForwardingOperand()) {
+      return *singleForwardingOp;
     }
-    return {};
+    return forwardingInst->getAllOperands();
   }
 
   bool canForwardOwnedCompatibleValuesOnly() {
@@ -321,6 +311,10 @@ class ForwardingOperation {
   /// Return true if the forwarded value is address-only either before or after
   /// forwarding.
   bool isAddressOnly() const;
+
+  // Call \p visitor on all forwarded results of the current forwarding
+  // operation.
+  bool visitForwardedValues(function_ref<bool(SILValue)> visitor);
 };
 } // end namespace swift
 

lib/SIL/Utils/InstWrappers.cpp

@@ -59,14 +59,39 @@ bool ForwardingOperation::hasSameRepresentation() const {
 }
 
 bool ForwardingOperation::isAddressOnly() const {
-  if (canForwardAllOperands()) {
-    // All ForwardingInstructions that forward all operands are currently a
-    // single value instruction.
-    auto *aggregate =
-        cast<OwnershipForwardingSingleValueInstruction>(forwardingInst);
-    // If any of the operands are address-only, then the aggregate must be.
-    return aggregate->getType().isAddressOnly(*forwardingInst->getFunction());
+  if (auto *singleForwardingOp = getSingleForwardingOperand()) {
+    return singleForwardingOp->get()->getType().isAddressOnly(
+        *forwardingInst->getFunction());
   }
-  return forwardingInst->getOperand(0)->getType().isAddressOnly(
-      *forwardingInst->getFunction());
+  // All ForwardingInstructions that forward all operands are currently a
+  // single value instruction.
+  auto *aggregate =
+      cast<OwnershipForwardingSingleValueInstruction>(forwardingInst);
+  // If any of the operands are address-only, then the aggregate must be.
+  return aggregate->getType().isAddressOnly(*forwardingInst->getFunction());
 }
+
+bool ForwardingOperation::visitForwardedValues(
+    function_ref<bool(SILValue)> visitor) {
+  if (auto *svi = dyn_cast<SingleValueInstruction>(forwardingInst)) {
+    return visitor(svi);
+  }
+  if (auto *mvri = dyn_cast<MultipleValueInstruction>(forwardingInst)) {
+    return llvm::all_of(mvri->getResults(), [&](SILValue value) {
+      if (value->getOwnershipKind() == OwnershipKind::None)
+        return true;
+      return visitor(value);
+    });
+  }
+  auto *ti = cast<TermInst>(forwardingInst);
+  assert(ti->mayHaveTerminatorResult());
+  return llvm::all_of(ti->getSuccessorBlocks(), [&](SILBasicBlock *succBlock) {
+    // If we do not have any arguments, then continue.
+    if (succBlock->args_empty())
+      return true;
+
+    auto args = succBlock->getSILPhiArguments();
+    assert(args.size() == 1 && "Transforming terminator with multiple args?!");
+    return visitor(args[0]);
+  });
+}

lib/SIL/Utils/MemAccessUtils.cpp

@@ -902,12 +902,8 @@ SILValue swift::findOwnershipReferenceAggregate(SILValue ref) {
 
     if (auto *inst = root->getDefiningInstruction()) {
       if (auto fwdOp = ForwardingOperation(inst)) {
-        if (fwdOp.canForwardFirstOperandOnly()) {
-          // The `enum` instruction can have no operand.
-          if (inst->getNumOperands() == 0)
-            return root;
-
-          root = inst->getOperand(0);
+        if (auto *singleForwardingOp = fwdOp.getSingleForwardingOperand()) {
+          root = singleForwardingOp->get();
           continue;
         }
       }

lib/SILOptimizer/SILCombiner/SILCombine.cpp

@@ -441,7 +441,7 @@ bool SILCombiner::doOneIteration(SILFunction &F, unsigned Iteration) {
       // block.
       if (auto *svi = dyn_cast<SingleValueInstruction>(I)) {
         if (auto fwdOp = ForwardingOperation(svi)) {
-          if (fwdOp.canForwardFirstOperandOnly() &&
+          if (fwdOp.getSingleForwardingOperand() &&
               SILValue(svi)->getOwnershipKind() == OwnershipKind::Owned) {
             // Try to sink the value. If we sank the value and deleted it,
             // continue. If we didn't optimize or sank but we are still able to

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -1806,7 +1806,11 @@ SILValue swift::makeValueAvailable(SILValue value, SILBasicBlock *inBlock) {
 bool swift::tryEliminateOnlyOwnershipUsedForwardingInst(
     SingleValueInstruction *forwardingInst, InstModCallbacks &callbacks) {
   auto fwdOp = ForwardingOperation(forwardingInst);
-  if (!fwdOp || !fwdOp.canForwardFirstOperandOnly()) {
+  if (!fwdOp) {
+    return false;
+  }
+  auto *singleFwdOp = fwdOp.getSingleForwardingOperand();
+  if (!singleFwdOp) {
     return false;
   }
 
@@ -1831,7 +1835,7 @@ bool swift::tryEliminateOnlyOwnershipUsedForwardingInst(
 
   // 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);
+  auto newValue = singleFwdOp->get();
   while (!forwardingInst->use_empty()) {
     auto *use = *(forwardingInst->use_begin());
     use->set(newValue);