Add ForwardingOperation::getSingleForwardingOperand api

Author: meg-gupta

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 (canForwardFirstOperandOnly()) {
-      return forwardingInst->getOperandRef(0);
+    if (auto *singleForwardingOp = getSingleForwardingOperand()) {
+      return *singleForwardingOp;
     }
-    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() {

lib/SIL/Utils/InstWrappers.cpp

@@ -59,14 +59,14 @@ 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());
 }

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