Add visitForwardedGuaranteedOperands utility.

Factors a mess of code in MemAccessUtils to handle forwarding
instruction types into a simpler utility. This utility is also needed
for ownership APIs, which need to be extended to handle these cases.

Author: atrick

include/swift/SIL/MemAccessUtils.h

@@ -206,6 +206,10 @@ SILValue findOwnershipReferenceRoot(SILValue ref);
 
 /// Look through all ownership forwarding instructions to find the values which
 /// were originally borrowed.
+///
+/// Note: This treats guaranteed forwarding phis like roots even though they do
+/// not introduce the borrow scope. This ensures that all roots dominate \p
+/// reference Value. But the client will need to handle forwarding phis.
 void findGuaranteedReferenceRoots(SILValue referenceValue,
                                   bool lookThroughNestedBorrows,
                                   SmallVectorImpl<SILValue> &roots);

include/swift/SIL/OwnershipUtils.h

@@ -1299,6 +1299,18 @@ void visitExtendedGuaranteedForwardingPhiBaseValuePairs(
     BorrowedValue borrow, function_ref<void(SILPhiArgument *, SILValue)>
                               visitGuaranteedForwardingPhiBaseValuePair);
 
+/// If \p value is a guaranteed non-phi value forwarded from it's instruction's
+/// operands, visit each forwarded operand.
+///
+/// Returns true if \p visitOperand was called (for convenience).
+///
+/// Precondition: \p value is not a phi. The client must handle phis first by
+/// checking if they are reborrows (using BorrowedValue). Reborrows have no
+/// forwarded operands. Guaranteed forwarding need to be handled by recursing
+/// through the phi operands.
+bool visitForwardedGuaranteedOperands(
+  SILValue value, function_ref<void(Operand *)> visitOperand);
+
 /// Visit the phis in the same block as \p phi which are reborrows of a borrow
 /// of one of the values reaching \p phi.
 ///

include/swift/SIL/SILInstruction.h

@@ -599,6 +599,10 @@ class SILInstruction : public llvm::ilist_node<SILInstruction> {
     return getTypeDependentOperands().size();
   }
 
+  unsigned getNumRealOperands() const {
+    return getAllOperands().size() - getNumTypeDependentOperands();
+  }
+
   bool isTypeDependentOperand(unsigned i) const {
     return i >= getNumOperands() - getNumTypeDependentOperands();
   }

lib/SIL/Utils/MemAccessUtils.cpp

@@ -16,6 +16,8 @@
 #include "swift/Basic/GraphNodeWorklist.h"
 #include "swift/SIL/Consumption.h"
 #include "swift/SIL/DynamicCasts.h"
+#include "swift/SIL/NodeDatastructures.h"
+#include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/SILBridging.h"
 #include "swift/SIL/SILBridgingUtils.h"
 #include "swift/SIL/SILInstruction.h"
@@ -863,75 +865,39 @@ SILValue swift::findOwnershipReferenceRoot(SILValue ref) {
 void swift::findGuaranteedReferenceRoots(SILValue value,
                                          bool lookThroughNestedBorrows,
                                          SmallVectorImpl<SILValue> &roots) {
-  GraphNodeWorklist<SILValue, 4> worklist;
-  auto addAllOperandsToWorklist = [&worklist](SILInstruction *inst) -> bool {
-    if (inst->getNumOperands() > 0) {
-      for (auto operand : inst->getOperandValues()) {
-        worklist.insert(operand);
-      }
-      return true;
-    }
-    return false;
-  };
-  worklist.initialize(value);
+  ValueWorklist worklist(referenceValue->getFunction());
+  worklist.pushIfNotVisited(referenceValue);
   while (auto value = worklist.pop()) {
-    if (auto *result = SILArgument::isTerminatorResult(value)) {
-      if (auto *forwardedOper = result->forwardedTerminatorResultOperand()) {
-        worklist.insert(forwardedOper->get());
-        continue;
-      }
-    } else if (auto *inst = value->getDefiningInstruction()) {
-      if (auto *bbi = dyn_cast<BeginBorrowInst>(inst)) {
+    // Instructions may forwarded None ownership to guaranteed.
+    if (value->getOwnershipKind() != OwnershipKind::Guaranteed)
+      continue;
+
+    if (SILArgument::asPhi(value)) {
+      roots.push_back(value);
+      continue;
+    }
+
+    if (visitForwardedGuaranteedOperands(value, [&](Operand *operand) {
+        worklist.pushIfNotVisited(operand->get());
+      })) {
+      // This instruction is not a root if any operands were forwarded,
+      // regardless of whether they were already visited.
+      continue;
+    }
+    // Found a potential root.
+    if (lookThroughNestedBorrows) {
+      if (auto *bbi = dyn_cast<BeginBorrowInst>(value)) {
         auto borrowee = bbi->getOperand();
-        if (lookThroughNestedBorrows &&
-            borrowee->getOwnershipKind() == OwnershipKind::Guaranteed) {
+        if (borrowee->getOwnershipKind() == OwnershipKind::Guaranteed) {
           // A nested borrow, the root guaranteed earlier in the use-def chain.
-          worklist.insert(borrowee);
-        }
-        // The borrowee isn't guaranteed or we aren't looking through nested
-        // borrows.  Fall through to add the begin_borrow to roots.
-      } else if (auto *result =
-              dyn_cast<FirstArgOwnershipForwardingSingleValueInst>(inst)) {
-        if (result->getNumOperands() > 0) {
-          worklist.insert(result->getOperand(0));
-          continue;
-        }
-      } else if (auto *result =
-                     dyn_cast<AllArgOwnershipForwardingSingleValueInst>(inst)) {
-        if (addAllOperandsToWorklist(result)) {
-          continue;
-        }
-      } else if (auto *result = dyn_cast<OwnershipForwardingTermInst>(inst)) {
-        assert(false && "value defined by a terminator?!");
-      } else if (auto *result =
-                     dyn_cast<OwnershipForwardingConversionInst>(inst)) {
-        worklist.insert(result->getConverted());
-        continue;
-      } else if (auto *result =
-                     dyn_cast<OwnershipForwardingSelectEnumInstBase>(inst)) {
-        if (addAllOperandsToWorklist(result)) {
-          continue;
-        }
-      } else if (auto *result =
-                     dyn_cast<OwnershipForwardingMultipleValueInstruction>(
-                         inst)) {
-        if (addAllOperandsToWorklist(result)) {
-          continue;
-        }
-      } else if (auto *mvi =
-                     dyn_cast<MoveOnlyWrapperToCopyableValueInst>(inst)) {
-        if (addAllOperandsToWorklist(mvi)) {
-          continue;
-        }
-      } else if (auto *c = dyn_cast<CopyableToMoveOnlyWrapperValueInst>(inst)) {
-        if (addAllOperandsToWorklist(c)) {
+          worklist.pushIfNotVisited(borrowee);
           continue;
         }
+        // The borrowee isn't guaranteed or we aren't looking through nested
+        // borrows. Fall through to add the begin_borrow to roots.
       }
     }
-
-    if (value->getOwnershipKind() == OwnershipKind::Guaranteed)
-      roots.push_back(value);
+    roots.push_back(value);
   }
 }
 

lib/SIL/Utils/OwnershipUtils.cpp

@@ -1721,6 +1721,63 @@ void swift::visitExtendedGuaranteedForwardingPhiBaseValuePairs(
   }
 }
 
+/// If \p instruction forwards guaranteed values to its results, visit each
+/// forwarded operand. The visitor must check whether the forwarded value is
+/// guaranteed.
+///
+/// Return true \p visitOperand was called at least once.
+///
+/// \p visitOperand should always recheck for Guaranteed owernship if it
+/// matters, in case a cast forwards a trivial type to a nontrivial type.
+///
+/// This intentionally does not handle phis, which require recursive traversal
+/// to determine `isGuaranteedForwardingPhi`.
+bool swift::visitForwardedGuaranteedOperands(
+  SILValue value, function_ref<void(Operand *)> visitOperand) {
+
+  assert(!SILArgument::asPhi(value) && "phis are handled separately");
+
+  if (auto *termResult = SILArgument::isTerminatorResult(value)) {
+    if (auto *oper = termResult->forwardedTerminatorResultOperand()) {
+      visitOperand(oper);
+      return true;
+    }
+    return false;
+  }
+  auto *inst = value->getDefiningInstruction();
+  if (!inst)
+    return false;
+
+  // Bypass conversions that produce a guarantee value out of thin air.
+  if (inst->getNumRealOperands() == 0) {
+    return false;
+  }
+  if (isa<FirstArgOwnershipForwardingSingleValueInst>(inst)
+      || isa<OwnershipForwardingConversionInst>(inst)
+      || isa<OwnershipForwardingSelectEnumInstBase>(inst)
+      || isa<OwnershipForwardingMultipleValueInstruction>(inst)
+      || isa<MoveOnlyWrapperToCopyableValueInst>(inst)
+      || isa<CopyableToMoveOnlyWrapperValueInst>(inst)) {
+    assert(inst->getNumRealOperands() == 1
+           && "forwarding instructions must have a single real operand");
+    assert(!isa<SingleValueInstruction>(inst)
+           || !BorrowedValue(cast<SingleValueInstruction>(inst))
+                  && "forwarded operand cannot begin a borrow scope");
+    visitOperand(&inst->getOperandRef(0));
+    return true;
+  }
+  if (isa<AllArgOwnershipForwardingSingleValueInst>(inst)) {
+    assert(inst->getNumOperands() > 0 && "checked above");
+    assert(inst->getNumOperands() == inst->getNumRealOperands() &&
+           "mixin expects all readl operands");
+    for (auto &operand : inst->getAllOperands()) {
+      visitOperand(&operand);
+    }
+    return true;
+  }
+  return false;
+}
+
 /// Visit the phis in the same block as \p phi which are reborrows of a borrow
 /// of one of the values reaching \p phi.
 ///