Improve copy elimination of noncopyable captures in ClosureLifetimeFixup

For a mark_dependence user of the stack, look through the mark_dependence chain
to match the partial_apply

Author: meg-gupta

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -554,6 +554,18 @@ collectStackClosureLifetimeEnds(SmallVectorImpl<SILInstruction *> &lifetimeEnds,
   }
 }
 
+static bool lookThroughMarkDependenceChainForValue(MarkDependenceInst *mark,
+                                                   PartialApplyInst *pai) {
+  if (mark->getValue() == pai) {
+    return true;
+  }
+  auto *markChain = dyn_cast<MarkDependenceInst>(mark->getValue());
+  if (!markChain) {
+    return false;
+  }
+  return lookThroughMarkDependenceChainForValue(markChain, pai);
+}
+
 /// Rewrite a partial_apply convert_escape_to_noescape sequence with a single
 /// apply/try_apply user to a partial_apply [stack] terminated with a
 /// dealloc_stack placed after the apply.
@@ -736,6 +748,7 @@ static SILValue tryRewriteToPartialApplyStack(
         newPA->getOrigCalleeType()->getNumParameters() - newPA->getNumArguments();
 
   MarkDependenceInst *markDepChain = nullptr;
+
   for (unsigned i : indices(newPA->getArgumentOperands())) {
     auto &arg = newPA->getArgumentOperands()[i];
     SILValue copy = arg.get();
@@ -798,14 +811,14 @@ static SILValue tryRewriteToPartialApplyStack(
           continue;
         }
 
-        // If we're marking dependence of the current partial_apply on this
-        // stack slot, that's fine.
-        if (mark->getValue() != newPA
-            || mark->getBase() != stack) {
+        // If we're marking dependence on this stack slot for the current
+        // partial_apply or it's chain of mark_dependences, that's fine.
+        if (!lookThroughMarkDependenceChainForValue(mark, newPA) ||
+            mark->getBase() != stack) {
           LLVM_DEBUG(llvm::dbgs() << "-- had unexpected mark_dependence use\n";
                      use->getUser()->print(llvm::dbgs());
                      llvm::dbgs() << "\n");
-          
+          initialization = nullptr;
           break;
         }
         markDep = mark;

test/SILOptimizer/moveonly_closure_lifetime_fixup_address_only_borrowed_captures.swift

@@ -8,12 +8,27 @@ struct E<T>: ~Copyable {
     var t: T
 }
 
+struct C<T> {
+  var t: T
+}
+
 var escaper: () -> () = {}
 
 func nonescaping(_ x: () -> ()) { }
 func escaping(_ x: @escaping () -> ()) {  escaper = x }
 
 func borrow<T>(_: borrowing E<T>) {}
+func borrow<T>(_: borrowing C<T>) {}
+
+func testMultiCapture<T>(_ e: borrowing E<T>, _ c: C<T>) {
+    // CHECK: [[C_STK:%.*]] = alloc_stack $C<T>
+    // CHECK: copy_addr %1 to [init] [[C_STK]] : $*C<T>
+    // CHECK: partial_apply {{.*}}[on_stack] {{.*}}([[C_STK]], %0) :
+    nonescaping {
+        borrow(c)
+        borrow(e)
+    }
+}
 
 // CHECK-LABEL: sil {{.*}}16testNeverEscaped
 func testNeverEscaped<T>(_ e: borrowing E<T>) {