[SILOpt] Add utility to find guaranteed roots.

The new utility looks through ownership forwarding instructions to find
the original values with guaranteed ownership.

Author: nate-chandler

include/swift/SIL/MemAccessUtils.h

@@ -204,6 +204,11 @@ SILValue findReferenceRoot(SILValue ref);
 /// Find the first owned root of the reference.
 SILValue findOwnershipReferenceRoot(SILValue ref);
 
+/// Look through all ownership forwarding instructions to find the values which
+/// were originally borrowed.
+void findGuaranteedReferenceRoots(SILValue value,
+                                  SmallVectorImpl<SILValue> &roots);
+
 /// Find the aggregate containing the first owned root of the
 /// reference. Identical to findOwnershipReferenceRoot, but looks through
 /// struct_extract, tuple_extract, etc.

lib/SIL/Utils/MemAccessUtils.cpp

@@ -13,11 +13,12 @@
 #define DEBUG_TYPE "sil-access-utils"
 
 #include "swift/SIL/MemAccessUtils.h"
-#include "swift/SIL/SILModule.h"
-#include "swift/SIL/SILUndef.h"
-#include "swift/SIL/DynamicCasts.h"
+#include "swift/Basic/GraphNodeWorklist.h"
 #include "swift/SIL/Consumption.h"
+#include "swift/SIL/DynamicCasts.h"
 #include "swift/SIL/SILInstruction.h"
+#include "swift/SIL/SILModule.h"
+#include "swift/SIL/SILUndef.h"
 #include "llvm/Support/Debug.h"
 
 using namespace swift;
@@ -800,6 +801,61 @@ SILValue swift::findOwnershipReferenceRoot(SILValue ref) {
   return ref;
 }
 
+void swift::findGuaranteedReferenceRoots(SILValue value,
+                                         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);
+  while (auto value = worklist.pop()) {
+    if (auto *arg = dyn_cast<SILPhiArgument>(value)) {
+      if (auto *terminator = arg->getSingleTerminator()) {
+        worklist.insert(terminator->getOperand(arg->getIndex()));
+        continue;
+      }
+    } else if (auto *inst = value->getDefiningInstruction()) {
+      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;
+        }
+      }
+    }
+    if (value.getOwnershipKind() == OwnershipKind::Guaranteed)
+      roots.push_back(value);
+  }
+}
+
 /// Find the first owned aggregate containing the reference, or simply the
 /// reference root if no aggregate is found.
 ///