Introduce ForwardingOperation wrapper type

APIs on ForwardingInstruction should be written as static taking in
a SILInstruction as a parameter making it awkward.

Introduce a ForwardingOperation wrapper type and move the apis from the
old "mixin" class to the wrapper type.

Add new api getForwardedOperands()

Author: meg-gupta

include/swift/SIL/InstWrappers.h

@@ -237,6 +237,90 @@ class SelectEnumOperation {
   }
 };
 
+class ForwardingOperation {
+  SILInstruction *forwardingInst = nullptr;
+
+public:
+  explicit ForwardingOperation(SILInstruction *inst);
+
+  operator bool() const { return bool(forwardingInst); }
+  const SILInstruction *operator->() const { return forwardingInst; }
+  SILInstruction *operator->() { return forwardingInst; }
+  const SILInstruction *operator*() const { return forwardingInst; }
+  SILInstruction *operator*() { return forwardingInst; }
+
+  ValueOwnershipKind getForwardingOwnershipKind();
+  bool preservesOwnership();
+
+  // FIXME: Find a better name. Even unary instructions like struct_extract
+  // forward "all" operands.
+  bool canForwardAllOperands() const {
+    switch (forwardingInst->getKind()) {
+    case SILInstructionKind::StructInst:
+    case SILInstructionKind::TupleInst:
+    case SILInstructionKind::LinearFunctionInst:
+    case SILInstructionKind::DifferentiableFunctionInst:
+      return true;
+    default:
+      return false;
+    }
+  }
+
+  // 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);
+    }
+    return {};
+  }
+
+  MutableArrayRef<Operand> getForwardedOperands() {
+    if (canForwardAllOperands()) {
+      return forwardingInst->getAllOperands();
+    }
+    if (canForwardFirstOperandOnly()) {
+      return forwardingInst->getOperandRef(0);
+    }
+    return {};
+  }
+
+  bool canForwardOwnedCompatibleValuesOnly() {
+    switch (forwardingInst->getKind()) {
+    case SILInstructionKind::MarkUninitializedInst:
+      return true;
+    default:
+      return false;
+    }
+  }
+
+  bool canForwardGuaranteedCompatibleValuesOnly() {
+    switch (forwardingInst->getKind()) {
+    case SILInstructionKind::TupleExtractInst:
+    case SILInstructionKind::StructExtractInst:
+    case SILInstructionKind::DifferentiableFunctionExtractInst:
+    case SILInstructionKind::LinearFunctionExtractInst:
+      return true;
+    default:
+      return false;
+    }
+  }
+
+  /// Return true if the forwarded value has the same representation. If true,
+  /// then the result can be mapped to the same storage without a move or copy.
+  bool hasSameRepresentation() const;
+
+  /// Return true if the forwarded value is address-only either before or after
+  /// forwarding.
+  bool isAddressOnly() const;
+};
 } // end namespace swift
 
 #endif

include/swift/SIL/SILInstruction.h

@@ -1239,46 +1239,6 @@ class ForwardingInstruction {
     ownershipKind = newKind;
   }
 
-  static bool canForwardAllOperands(SILInstruction *inst) {
-    switch (inst->getKind()) {
-    case SILInstructionKind::StructInst:
-    case SILInstructionKind::TupleInst:
-    case SILInstructionKind::LinearFunctionInst:
-    case SILInstructionKind::DifferentiableFunctionInst:
-      return true;
-    default:
-      return false;
-    }
-  }
-
-  static bool canForwardFirstOperandOnly(SILInstruction *inst) {
-    if (!ForwardingInstruction::isa(inst)) {
-      return false;
-    }
-    return !canForwardAllOperands(inst);
-  }
-
-  static bool canForwardOwnedCompatibleValuesOnly(SILInstruction *inst) {
-    switch (inst->getKind()) {
-    case SILInstructionKind::MarkUninitializedInst:
-      return true;
-    default:
-      return false;
-    }
-  }
-
-  static bool canForwardGuaranteedCompatibleValuesOnly(SILInstruction *inst) {
-    switch (inst->getKind()) {
-    case SILInstructionKind::TupleExtractInst:
-    case SILInstructionKind::StructExtractInst:
-    case SILInstructionKind::DifferentiableFunctionExtractInst:
-    case SILInstructionKind::LinearFunctionExtractInst:
-      return true;
-    default:
-      return false;
-    }
-  }
-
   /// Defined inline below due to forward declaration issues.
   static ForwardingInstruction *get(SILInstruction *inst);
   static bool isa(SILInstructionKind kind);
@@ -1288,18 +1248,6 @@ class ForwardingInstruction {
       return isa(i);
     return false;
   }
-
-  /// Return true if the forwarded value has the same representation. If true,
-  /// then the result can be mapped to the same storage without a move or copy.
-  ///
-  /// \p inst is an ForwardingInstruction
-  static bool hasSameRepresentation(SILInstruction *inst);
-
-  /// Return true if the forwarded value is address-only either before or after
-  /// forwarding.
-  ///
-  /// \p inst is an ForwardingInstruction
-  static bool isAddressOnly(SILInstruction *inst);
 };
 
 /// A single value inst that forwards a static ownership from its first operand.

lib/SIL/IR/SILInstructions.cpp

@@ -3141,47 +3141,6 @@ ReturnInst::ReturnInst(SILFunction &func, SILDebugLocation debugLoc,
          "result info?!");
 }
 
-bool ForwardingInstruction::hasSameRepresentation(SILInstruction *inst) {
-  switch (inst->getKind()) {
-  // Explicitly list instructions which definitely involve a representation
-  // change.
-  case SILInstructionKind::SwitchEnumInst:
-  default:
-    // Conservatively assume that a conversion changes representation.
-    // Operations can be added as needed to participate in SIL opaque values.
-    assert(ForwardingInstruction::isa(inst));
-    return false;
-
-  case SILInstructionKind::ConvertFunctionInst:
-  case SILInstructionKind::DestructureTupleInst:
-  case SILInstructionKind::DestructureStructInst:
-  case SILInstructionKind::InitExistentialRefInst:
-  case SILInstructionKind::ObjectInst:
-  case SILInstructionKind::OpenExistentialBoxValueInst:
-  case SILInstructionKind::OpenExistentialRefInst:
-  case SILInstructionKind::OpenExistentialValueInst:
-  case SILInstructionKind::MarkMustCheckInst:
-  case SILInstructionKind::MarkUninitializedInst:
-  case SILInstructionKind::SelectEnumInst:
-  case SILInstructionKind::StructExtractInst:
-  case SILInstructionKind::TupleExtractInst:
-    return true;
-  }
-}
-
-bool ForwardingInstruction::isAddressOnly(SILInstruction *inst) {
-  if (canForwardAllOperands(inst)) {
-    // All ForwardingInstructions that forward all operands are currently a
-    // single value instruction.
-    auto *aggregate = cast<OwnershipForwardingSingleValueInstruction>(inst);
-    // If any of the operands are address-only, then the aggregate must be.
-    return aggregate->getType().isAddressOnly(*inst->getFunction());
-  }
-  // All other forwarding instructions must forward their first operand.
-  assert(canForwardFirstOperandOnly(inst));
-  return inst->getOperand(0)->getType().isAddressOnly(*inst->getFunction());
-}
-
 // This may be called in an invalid SIL state. SILCombine creates new
 // terminators in non-terminator position and defers deleting the original
 // terminator until after all modification.

lib/SIL/Utils/CMakeLists.txt

@@ -8,6 +8,7 @@ target_sources(swiftSIL PRIVATE
   FieldSensitivePrunedLiveness.cpp
   GenericSpecializationMangler.cpp
   InstructionUtils.cpp
+  InstWrappers.cpp
   LoopInfo.cpp
   MemAccessUtils.cpp
   MemoryLocations.cpp
@@ -22,4 +23,5 @@ target_sources(swiftSIL PRIVATE
   SILBridging.cpp
   SILInstructionWorklist.cpp
   SILRemarkStreamer.cpp
-  ValueUtils.cpp)
+  ValueUtils.cpp
+)

lib/SIL/Utils/InstWrappers.cpp

@@ -0,0 +1,72 @@
+//===--- InstWrappers.cpp ------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "swift/SIL/InstWrappers.h"
+#include "swift/SIL/SILFunction.h"
+
+using namespace swift;
+
+ForwardingOperation::ForwardingOperation(SILInstruction *inst) {
+  if (ForwardingInstruction::isa(inst)) {
+    this->forwardingInst = inst;
+    return;
+  }
+}
+
+ValueOwnershipKind ForwardingOperation::getForwardingOwnershipKind() {
+  return ForwardingInstruction::get(forwardingInst)
+      ->getForwardingOwnershipKind();
+}
+
+bool ForwardingOperation::preservesOwnership() {
+  return ForwardingInstruction::get(forwardingInst)->preservesOwnership();
+}
+
+bool ForwardingOperation::hasSameRepresentation() const {
+  switch (forwardingInst->getKind()) {
+  // Explicitly list instructions which definitely involve a representation
+  // change.
+  case SILInstructionKind::SwitchEnumInst:
+  default:
+    // Conservatively assume that a conversion changes representation.
+    // Operations can be added as needed to participate in SIL opaque values.
+    return false;
+
+  case SILInstructionKind::ConvertFunctionInst:
+  case SILInstructionKind::DestructureTupleInst:
+  case SILInstructionKind::DestructureStructInst:
+  case SILInstructionKind::InitExistentialRefInst:
+  case SILInstructionKind::ObjectInst:
+  case SILInstructionKind::OpenExistentialBoxValueInst:
+  case SILInstructionKind::OpenExistentialRefInst:
+  case SILInstructionKind::OpenExistentialValueInst:
+  case SILInstructionKind::MarkMustCheckInst:
+  case SILInstructionKind::MarkUninitializedInst:
+  case SILInstructionKind::SelectEnumInst:
+  case SILInstructionKind::StructExtractInst:
+  case SILInstructionKind::TupleExtractInst:
+    return true;
+  }
+}
+
+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());
+  }
+  return forwardingInst->getOperand(0)->getType().isAddressOnly(
+      *forwardingInst->getFunction());
+}

lib/SIL/Utils/MemAccessUtils.cpp

@@ -900,13 +900,15 @@ SILValue swift::findOwnershipReferenceAggregate(SILValue ref) {
       return root;
 
     if (auto *inst = root->getDefiningInstruction()) {
-      if (ForwardingInstruction::canForwardFirstOperandOnly(inst)) {
-        // The `enum` instruction can have no operand.
-        if (inst->getNumOperands() == 0)
-          return root;
+      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);
-        continue;
+          root = inst->getOperand(0);
+          continue;
+        }
       }
     }
     if (auto *termResult = SILArgument::isTerminatorResult(root)) {