[ownership] Implement getSingleOwnedValueIntroducer and getAllOwnedValueIntroducers().

These enable robust walking up the def-use ossa graph to find the value that
introduces an independent underlying owned value, looking through forwarding
insts.

Author: gottesmm

include/swift/SIL/OwnershipUtils.h

@@ -58,6 +58,14 @@ bool isOwnedForwardingValueKind(SILNodeKind kind);
 /// forward guaranteed ownership.
 bool isOwnedForwardingInstruction(SILInstruction *inst);
 
+/// Does this value 'forward' owned ownership, but may not be able to forward
+/// guaranteed ownership.
+///
+/// This will be either a multiple value instruction resuilt, a single value
+/// instruction that forwards or an argument that forwards the ownership from a
+/// previous terminator.
+bool isOwnedForwardingValue(SILValue value);
+
 struct BorrowScopeOperandKind {
   enum Kind {
     BeginBorrow,
@@ -607,6 +615,17 @@ struct OwnedValueIntroducer {
       : value(value), kind(kind) {}
 };
 
+/// Look up the def-use graph starting at use \p inputOperand, recording any
+/// values that act as "owned" introducers.
+///
+/// NOTE: This may return multiple owned introducers in cases where there are
+/// phi-like nodes in the IR like any true phi block arguments or aggregate
+/// literal instructions (struct, tuple, enum, etc.).
+bool getAllOwnedValueIntroducers(SILValue value,
+                                 SmallVectorImpl<OwnedValueIntroducer> &out);
+
+Optional<OwnedValueIntroducer> getSingleOwnedValueIntroducer(SILValue value);
+
 } // namespace swift
 
 #endif

lib/SIL/OwnershipUtils.cpp

@@ -85,6 +85,17 @@ bool swift::isOwnedForwardingInstruction(SILInstruction *inst) {
   }
 }
 
+bool swift::isOwnedForwardingValue(SILValue value) {
+  switch (value->getKind()) {
+  // Phi arguments always forward ownership.
+  case ValueKind::SILPhiArgument:
+    return true;
+  default:
+    return isOwnedForwardingValueKind(
+        value->getKindOfRepresentativeSILNodeInObject());
+  }
+}
+
 bool swift::isGuaranteedForwardingValue(SILValue value) {
   // If we have an argument from a transforming terminator, we can forward
   // guaranteed.
@@ -95,6 +106,7 @@ bool swift::isGuaranteedForwardingValue(SILValue value) {
       }
     }
   }
+
   return isGuaranteedForwardingValueKind(
       value->getKindOfRepresentativeSILNodeInObject());
 }
@@ -443,6 +455,49 @@ bool BorrowScopeIntroducingValue::visitInteriorPointerOperands(
   return true;
 }
 
+//===----------------------------------------------------------------------===//
+//                          Owned Value Introducers
+//===----------------------------------------------------------------------===//
+
+void OwnedValueIntroducerKind::print(llvm::raw_ostream &os) const {
+  switch (value) {
+  case OwnedValueIntroducerKind::Apply:
+    os << "Apply";
+    return;
+  case OwnedValueIntroducerKind::BeginApply:
+    os << "BeginApply";
+    return;
+  case OwnedValueIntroducerKind::TryApply:
+    os << "TryApply";
+    return;
+  case OwnedValueIntroducerKind::Copy:
+    os << "Copy";
+    return;
+  case OwnedValueIntroducerKind::LoadCopy:
+    os << "LoadCopy";
+    return;
+  case OwnedValueIntroducerKind::LoadTake:
+    os << "LoadTake";
+    return;
+  case OwnedValueIntroducerKind::Phi:
+    os << "Phi";
+    return;
+  case OwnedValueIntroducerKind::FunctionArgument:
+    os << "FunctionArgument";
+    return;
+  case OwnedValueIntroducerKind::PartialApplyInit:
+    os << "PartialApplyInit";
+    return;
+  case OwnedValueIntroducerKind::AllocBoxInit:
+    os << "AllocBoxInit";
+    return;
+  case OwnedValueIntroducerKind::AllocRefInit:
+    os << "AllocRefInit";
+    return;
+  }
+  llvm_unreachable("Covered switch isn't covered");
+}
+
 //===----------------------------------------------------------------------===//
 //                       Introducer Searching Routines
 //===----------------------------------------------------------------------===//
@@ -546,3 +601,103 @@ swift::getSingleBorrowIntroducingValue(SILValue inputValue) {
 
   llvm_unreachable("Should never hit this");
 }
+
+bool swift::getAllOwnedValueIntroducers(
+    SILValue inputValue, SmallVectorImpl<OwnedValueIntroducer> &out) {
+  if (inputValue.getOwnershipKind() != ValueOwnershipKind::Owned)
+    return false;
+
+  SmallVector<SILValue, 32> worklist;
+  worklist.emplace_back(inputValue);
+
+  while (!worklist.empty()) {
+    SILValue value = worklist.pop_back_val();
+
+    // First check if v is an introducer. If so, stash it and continue.
+    if (auto introducer = OwnedValueIntroducer::get(value)) {
+      out.push_back(*introducer);
+      continue;
+    }
+
+    // If v produces .none ownership, then we can ignore it. It is important
+    // that we put this before checking for guaranteed forwarding instructions,
+    // since we want to ignore guaranteed forwarding instructions that in this
+    // specific case produce a .none value.
+    if (value.getOwnershipKind() == ValueOwnershipKind::None)
+      continue;
+
+    // Otherwise if v is an ownership forwarding value, add its defining
+    // instruction
+    if (isOwnedForwardingValue(value)) {
+      if (auto *i = value->getDefiningInstruction()) {
+        llvm::copy(i->getOperandValues(true /*skip type dependent ops*/),
+                   std::back_inserter(worklist));
+        continue;
+      }
+
+      // Otherwise, we should have a block argument that is defined by a single
+      // predecessor terminator.
+      auto *arg = cast<SILPhiArgument>(value);
+      auto *termInst = arg->getSingleTerminator();
+      assert(termInst && termInst->isTransformationTerminator());
+      assert(termInst->getNumOperands() == 1 &&
+             "Transforming terminators should always have a single operand");
+      worklist.push_back(termInst->getAllOperands()[0].get());
+      continue;
+    }
+
+    // Otherwise, this is an introducer we do not understand. Bail and return
+    // false.
+    return false;
+  }
+
+  return true;
+}
+
+Optional<OwnedValueIntroducer>
+swift::getSingleOwnedValueIntroducer(SILValue inputValue) {
+  if (inputValue.getOwnershipKind() != ValueOwnershipKind::Owned)
+    return None;
+
+  SILValue currentValue = inputValue;
+  while (true) {
+    // First check if our initial value is an introducer. If we have one, just
+    // return it.
+    if (auto introducer = OwnedValueIntroducer::get(currentValue)) {
+      return introducer;
+    }
+
+    // Otherwise if v is an ownership forwarding value, add its defining
+    // instruction
+    if (isOwnedForwardingValue(currentValue)) {
+      if (auto *i = currentValue->getDefiningInstruction()) {
+        auto instOps = i->getOperandValues(true /*ignore type dependent ops*/);
+        // If we have multiple incoming values, return .None. We can't handle
+        // this.
+        auto begin = instOps.begin();
+        if (std::next(begin) != instOps.end()) {
+          return None;
+        }
+        // Otherwise, set currentOp to the single operand and continue.
+        currentValue = *begin;
+        continue;
+      }
+
+      // Otherwise, we should have a block argument that is defined by a single
+      // predecessor terminator.
+      auto *arg = cast<SILPhiArgument>(currentValue);
+      auto *termInst = arg->getSingleTerminator();
+      assert(termInst && termInst->isTransformationTerminator());
+      assert(termInst->getNumOperands() == 1 &&
+             "Transformation terminators should only have single operands");
+      currentValue = termInst->getAllOperands()[0].get();
+      continue;
+    }
+
+    // Otherwise, this is an introducer we do not understand. Bail and return
+    // None.
+    return None;
+  }
+
+  llvm_unreachable("Should never hit this");
+}