Add findOwnershipReferenceAggregate().

Simple convenience on top of findOwnershipReferenceRoot to handle
guaranteed values forwarded through struct/tuple extract.

Note: eventually this should be handled by a ReferenceRoot abstraction
that stores the component path.

Author: atrick

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);
 
+/// Find the aggregate containing the first owned root of the
+/// reference. Identical to findOwnershipReferenceRoot, but looks through
+/// struct_extract, tuple_extract, etc.
+SILValue findOwnershipReferenceAggregate(SILValue ref);
+
 /// Return true if \p address points to a let-variable.
 ///
 /// let-variables are only written during let-variable initialization, which is

include/swift/SIL/SILInstruction.h

@@ -4937,6 +4937,8 @@ class ConversionInst : public SingleValueInstruction {
 
 /// A conversion inst that produces a static OwnershipKind set upon the
 /// instruction's construction.
+///
+/// The first operand is the ownership equivalent source.
 class OwnershipForwardingConversionInst : public ConversionInst,
                                           public OwnershipForwardingMixin {
 protected:
@@ -7877,6 +7879,8 @@ class TermInst : public NonValueInstruction {
   TermKind getTermKind() const { return TermKind(getKind()); }
 
   /// Returns true if this is a transformation terminator.
+  ///
+  /// The first operand is the transformed source.
   bool isTransformationTerminator() const {
     switch (getTermKind()) {
     case TermKind::UnwindInst:
@@ -9396,6 +9400,7 @@ SILFunction *ApplyInstBase<Impl, Base, false>::getCalleeFunction() const {
   }
 }
 
+/// The first operand is the ownership equivalent source.
 class OwnershipForwardingMultipleValueInstruction
     : public MultipleValueInstruction,
       public OwnershipForwardingMixin {

lib/SIL/Utils/MemAccessUtils.cpp

@@ -799,6 +799,38 @@ SILValue swift::findOwnershipReferenceRoot(SILValue ref) {
   return FindReferenceRoot(true /*preserveOwnership*/).findRoot(ref);
 }
 
+/// Find the first owned aggregate containing the reference, or simply the
+/// reference root if no aggregate is found.
+///
+/// TODO: Add a component path to a ReferenceRoot abstraction and handle
+/// that within FindReferenceRoot.
+SILValue swift::findOwnershipReferenceAggregate(SILValue ref) {
+  SILValue root = ref;
+  while(true) {
+    root = findOwnershipReferenceRoot(root);
+    if (!root)
+      return root;
+    if (isa<FirstArgOwnershipForwardingSingleValueInst>(root)
+        || isa<OwnershipForwardingConversionInst>(root)
+        || isa<OwnershipForwardingSelectEnumInstBase>(root)
+        || isa<OwnershipForwardingMultipleValueInstruction>(root)) {
+      root = root->getDefiningInstruction()->getOperand(0);
+      continue;
+    }
+    if (auto *arg = dyn_cast<SILArgument>(root)) {
+      if (auto *term = arg->getSingleTerminator()) {
+        if (term->isTransformationTerminator()) {
+          assert(OwnershipForwardingTermInst::isa(term));
+          root = term->getOperand(0);
+          continue;
+        }
+      }
+    }
+    break;
+  }
+  return root;
+}
+
 //===----------------------------------------------------------------------===//
 //                            MARK: AccessStorage
 //===----------------------------------------------------------------------===//