[sil-combine] Improve dead closures elimination

Cover more corner cases, e.g. unbalanced retain/release pairs due to an unreachable at the end of the block.
This significantly improves the performance of some closure-heavy test-cases.

rdar://23691514

Author: swiftix

lib/SILAnalysis/ARCAnalysis.cpp

@@ -777,21 +777,19 @@ static bool successorHasLiveIn(SILBasicBlock *BB,
   return false;
 }
 
-// Walk backwards in BB looking for strong_release or dealloc_box of
-// the given value, and add it to Releases.
-static bool addLastRelease(SILValue V, SILBasicBlock *BB,
-                           ReleaseTracker &Tracker) {
+// Walk backwards in BB looking for the last use of a given
+// value, and add it to the set of release points.
+static bool addLastUse(SILValue V, SILBasicBlock *BB,
+                       ReleaseTracker &Tracker) {
   for (auto I = BB->rbegin(); I != BB->rend(); ++I) {
-    if (isa<StrongReleaseInst>(*I) || isa<DeallocBoxInst>(*I) ||
-        isa<ReleaseValueInst>(*I)) {
-      if (I->getOperand(0) != V)
-        continue;
-
-      Tracker.trackLastRelease(&*I);
-      return true;
-    }
+    for (auto &Op : I->getAllOperands())
+      if (Op.get().getDef() == V.getDef()) {
+        Tracker.trackLastRelease(&*I);
+        return true;
+      }
   }
 
+  llvm_unreachable("BB is expected to have a use of a closure");
   return false;
 }
 
@@ -850,7 +848,7 @@ bool swift::getFinalReleasesForValue(SILValue V, ReleaseTracker &Tracker) {
   // release/dealloc.
   for (auto *BB : UseBlocks)
     if (!successorHasLiveIn(BB, LiveIn))
-      if (!addLastRelease(V, BB, Tracker))
+      if (!addLastUse(V, BB, Tracker))
         return false;
 
   return true;

test/SILPasses/sil_combine.sil

@@ -736,18 +736,19 @@ bb1:
   return %5 : $()
 }
 
-sil @unbalanced_closure : $@convention(thin) (Builtin.Int32) -> ()
+sil @unbalanced_closure : $@convention(thin) (@guaranteed B) -> ()
 
 // CHECK-LABEL: sil @partial_apply_unbalanced_retain_release
-sil @partial_apply_unbalanced_retain_release : $@convention(thin) (Builtin.Int32) -> () {
+sil @partial_apply_unbalanced_retain_release : $@convention(thin) (@owned B) -> () {
 // CHECK: bb0
-bb0(%0 : $Builtin.Int32):
-  %1 = function_ref @unbalanced_closure : $@convention(thin) (Builtin.Int32) -> ()
-// CHECK: partial_apply
-  %2 = partial_apply %1(%0) : $@convention(thin) (Builtin.Int32) -> ()
+bb0(%0 : $B):
+  %1 = function_ref @unbalanced_closure : $@convention(thin) (@guaranteed B) -> ()
+// CHECK-NOT: partial_apply
+  %2 = partial_apply %1(%0) : $@convention(thin) (@guaranteed B) -> ()
+// Check that the arguments of the closure are released after its last use
+// CHECK-NEXT: strong_release %0 : $B
   strong_retain %2 : $@callee_owned () -> ()
-  %4 = builtin "int_trap"() : $()
-// CHECK: unreachable
+// CHECK-NEXT: unreachable
   unreachable
 }