Merge pull request swiftlang#81341 from eeckstein/closure-optimizations

SILCombine: handle `mark_dependence` in some closure optimizations

Author: eeckstein

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -1520,8 +1520,14 @@ SILInstruction *SILCombiner::legacyVisitApplyInst(ApplyInst *AI) {
     return nullptr;
 
   SILValue callee = AI->getCallee();
-  if (auto *cee = dyn_cast<ConvertEscapeToNoEscapeInst>(callee)) {
-    callee = cee->getOperand();
+  for (;;) {
+    if (auto *cee = dyn_cast<ConvertEscapeToNoEscapeInst>(callee)) {
+      callee = cee->getOperand();
+    } else if (auto *mdi = dyn_cast<MarkDependenceInst>(callee)) {
+      callee = mdi->getValue();
+    } else {
+      break;
+    }
   }
   if (auto *CFI = dyn_cast<ConvertFunctionInst>(callee))
     return optimizeApplyOfConvertFunctionInst(AI, CFI);

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -963,6 +963,9 @@ static bool useHasTransitiveOwnership(const SILInstruction *inst) {
   if (isa<ConvertEscapeToNoEscapeInst>(inst))
     return true;
 
+  if (isa<MarkDependenceInst>(inst))
+    return true;
+
   // Look through copy_value, begin_borrow, move_value. They are inert for our
   // purposes, but we need to look through it.
   return isa<CopyValueInst>(inst) || isa<BeginBorrowInst>(inst) ||
@@ -1095,13 +1098,16 @@ void swift::getConsumedPartialApplyArgs(PartialApplyInst *pai,
   }
 }
 
-bool swift::collectDestroys(SingleValueInstruction *inst,
-                            SmallVectorImpl<Operand *> &destroys) {
+static bool collectDestroysRecursively(SingleValueInstruction *inst,
+                                       SmallVectorImpl<Operand *> &destroys,
+                                       InstructionSet &visited) {
   bool isDead = true;
   for (Operand *use : inst->getUses()) {
     SILInstruction *user = use->getUser();
+    if (!visited.insert(user))
+      continue;
     if (useHasTransitiveOwnership(user)) {
-      if (!collectDestroys(cast<SingleValueInstruction>(user), destroys))
+      if (!collectDestroysRecursively(cast<SingleValueInstruction>(user), destroys, visited))
         isDead = false;
       destroys.push_back(use);
     } else if (useDoesNotKeepValueAlive(user)) {
@@ -1113,6 +1119,12 @@ bool swift::collectDestroys(SingleValueInstruction *inst,
   return isDead;
 }
 
+bool swift::collectDestroys(SingleValueInstruction *inst,
+                            SmallVectorImpl<Operand *> &destroys) {
+  InstructionSet visited(inst->getFunction());
+  return collectDestroysRecursively(inst, destroys, visited);
+}
+
 /// Move the original arguments of the partial_apply into newly created
 /// temporaries to extend the lifetime of the arguments until the partial_apply
 /// is finally destroyed.

test/SILOptimizer/sil_combine1.swift

@@ -36,3 +36,9 @@ public func test_compose_closure() -> Int32 {
   return gs
 }
 
+// CHECK-LABEL: sil @$s12sil_combine122remove_partial_applies_3key5valueySDySSSiGSgz_SSSitF :
+// CHECK-NOT:     partial_apply
+// CHECK:       } // end sil function '$s12sil_combine122remove_partial_applies_3key5valueySDySSSiGSgz_SSSitF'
+public func remove_partial_applies(_ dict: inout [String: Int]?, key: String, value: Int) {
+  dict?[key, default: value] = value 
+}

test/SILOptimizer/sil_combine_apply_ossa.sil

@@ -1172,6 +1172,37 @@ bb0(%0 : $Int):
   return %r : $()
 }
 
+// CHECK-LABEL: sil [ossa] @test_mark_dependence_of_closure :
+// CHECK-NOT:     partial_apply
+// CHECK-LABEL: } // end sil function 'test_mark_dependence_of_closure'
+sil [ossa] @test_mark_dependence_of_closure : $@convention(thin) (Int) -> () {
+bb0(%0 : $Int):
+  %3 = function_ref @closure2 : $@convention(thin) (Int) -> ()
+  %4 = partial_apply %3(%0) : $@convention(thin) (Int) -> ()
+  %5 = convert_function %4 to $@callee_owned () -> ()
+  %6 = convert_escape_to_noescape %5 to $@noescape @callee_owned () -> ()
+  %7 = mark_dependence %6 on %5
+  apply %7() : $@noescape @callee_owned () -> ()
+  destroy_value %5
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_dead_closure_with_mark_dependence :
+// CHECK-NOT:     partial_apply
+// CHECK-LABEL: } // end sil function 'test_dead_closure_with_mark_dependence'
+sil [ossa] @test_dead_closure_with_mark_dependence : $@convention(thin) (Int) -> () {
+bb0(%0 : $Int):
+  %3 = function_ref @closure2 : $@convention(thin) (Int) -> ()
+  %4 = partial_apply %3(%0) : $@convention(thin) (Int) -> ()
+  %6 = convert_escape_to_noescape %4 to $@noescape @callee_owned () -> ()
+  %7 = mark_dependence %6 on %4
+  destroy_value %4
+  destroy_value %7
+  %r = tuple ()
+  return %r : $()
+}
+
 sil [ossa] @closure2 : $@convention(thin) (Int) -> ()
 
 sil [ossa] @yield1 : $@yield_once(Float) -> (@yields Float) {