Fix OSSA support for mark_dependence [nonescaping]

Handle escapes to function results.

Author: atrick

SwiftCompilerSources/Sources/Optimizer/Utilities/OwnershipLiveness.swift

@@ -307,8 +307,8 @@ extension OwnershipUseVisitor {
     }
     // Otherwise, directly visit the scope ending uses.
     //
-    // TODO: change visitScopeEndingOperands to take a non-escaping
-    // closure and call ownershipLeafUse directly.
+    // TODO: remove this stack by changign visitScopeEndingOperands to
+    // take a non-escaping closure that can call ownershipLeafUse.
     var stack = Stack<Operand>(context)
     defer { stack.deinitialize() }
     _ = borrowInst.visitScopeEndingOperands(context) {
@@ -448,9 +448,6 @@ extension OwnershipUseVisitor {
 /// MoveValueInst, and Allocation. Then this visitor should assert
 /// that the forward-extended lifetime introducer has no pointer
 /// escaping uses.
-///
-/// TODO: Change the operandOwnership of MarkDependenceInst base operand.
-/// It should be a borrowing operand, not a pointer escape.
 struct InteriorUseWalker {
   let functionContext: FunctionPassContext
   var context: Context { functionContext }

lib/SIL/IR/OperandOwnership.cpp

@@ -658,11 +658,16 @@ OperandOwnershipClassifier::visitMarkDependenceInst(MarkDependenceInst *mdi) {
     return getOwnershipKind().getForwardingOperandOwnership(
       /*allowUnowned*/true);
   }
-  if (getOperandIndex() == MarkDependenceInst::Base && mdi->isNonEscaping()) {
+  if (mdi->isNonEscaping()) {
     // This creates a "dependent value", just like on-stack partial_apply, which
     // we treat like a borrow.
     return OperandOwnership::Borrow;
   }
+  if (mdi->hasUnresolvedEscape()) {
+    // This creates a dependent value that may extend beyond the parent's
+    // lifetime.
+    return OperandOwnership::UnownedInstantaneousUse;
+  }
   // FIXME: Add an end_dependence instruction so we can treat mark_dependence as
   // a borrow of the base (mark_dependence %base -> end_dependence is analogous
   // to a borrow scope).

lib/SIL/IR/SILInstruction.cpp

@@ -1805,6 +1805,15 @@ static bool visitRecursivelyLifetimeEndingUses(
       }
       continue;
     }
+    if (auto *ret = dyn_cast<ReturnInst>(use->getUser())) {
+      auto fnTy = ret->getFunction()->getLoweredFunctionType();
+      assert(!fnTy->getLifetimeDependenceInfo().empty());
+      if (!func(use)) {
+        return false;
+      }
+      continue;
+    }
+    // FIXME: Handle store to indirect result
 
     // There shouldn't be any dead-end consumptions of a nonescaping
     // partial_apply that don't forward it along, aside from destroy_value.