ClosureLifetimeFixup: Fix closure insertion point in deallocating dei…

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -66,6 +66,27 @@ static SILBasicBlock *getOptionalDiamondSuccessor(SwitchEnumInst *SEI) {
    return nullptr;
 }
 
+/// Find a safe insertion point for closure destruction. We might create a
+/// closure that captures self in deinit of self. In this situation it is not
+/// safe to destroy the closure after we called super deinit. We have to place
+/// the closure destruction before that call.
+///
+///  %deinit = objc_super_method %0 : $C, #A.deinit!deallocator.foreign
+///  %super = upcast %0 : $C to $A
+///  apply %deinit(%super) : $@convention(objc_method) (A) -> ()
+///  end_lifetime %super : $A
+static SILInstruction *getDeinitSafeClosureDestructionPoint(TermInst *Term) {
+  for (auto It = Term->getParent()->rbegin(), E = Term->getParent()->rend();
+       It != E; ++It) {
+    if (auto *EndLifetime = dyn_cast<EndLifetimeInst>(&*It)) {
+      auto *SuperInstance = EndLifetime->getOperand()->getDefiningInstruction();
+      assert(SuperInstance && "Expected an instruction");
+      return SuperInstance;
+    }
+  }
+  return Term;
+}
+
 /// Extend the lifetime of the convert_escape_to_noescape's operand to the end
 /// of the function.
 static void extendLifetimeToEndOfFunction(SILFunction &Fn,
@@ -111,10 +132,11 @@ static void extendLifetimeToEndOfFunction(SILFunction &Fn,
   Fn.findExitingBlocks(ExitingBlocks);
   for (auto *Exit : ExitingBlocks) {
     auto *Term = Exit->getTerminator();
-    SILBuilderWithScope B(Term);
-    B.setInsertionPoint(Term);
+    auto *SafeClosureDestructionPt = getDeinitSafeClosureDestructionPoint(Term);
+    SILBuilderWithScope B(SafeClosureDestructionPt);
     B.createDestroyAddr(loc, Slot);
-    B.createDeallocStack(loc, Slot);
+    SILBuilderWithScope B2(Term);
+    B2.createDeallocStack(loc, Slot);
   }
 }
 

test/SILOptimizer/closure_lifetime_fixup_objc.swift

@@ -79,3 +79,24 @@ public func dontCrash() {
     queue.sync { }
   }
 }
+
+@_silgen_name("getDispatchQueue")
+func getDispatchQueue() -> DispatchQueue
+
+// We must not release the closure after calling super.deinit.
+// CHECK: sil hidden @$S27closure_lifetime_fixup_objc1CCfD : $@convention(method) (@owned C) -> () {
+// CHECK: bb0([[SELF:%.*]] : $C):
+// CHECK:   [[F:%.*]] = function_ref @$S27closure_lifetime_fixup_objc1CCfdyyXEfU_
+// CHECK:   [[PA:%.*]] = partial_apply [callee_guaranteed] [[F]](%0)
+// CHECK:   [[OPT:%.*]] = enum $Optional<@callee_guaranteed () -> ()>, #Optional.some!enumelt.1, [[PA]]
+// CHECK:   [[DEINIT:%.*]] = objc_super_method [[SELF]] : $C, #NSObject.deinit!deallocator.foreign
+// CHECK:   release_value [[OPT]] : $Optional<@callee_guaranteed () -> ()>
+// CHECK:   [[SUPER:%.*]] = upcast [[SELF]] : $C to $NSObject               // user: %34
+// CHECK-NEXT:   apply [[DEINIT]]([[SUPER]]) : $@convention(objc_method) (NSObject) -> ()
+// CHECK-NEXT:   [[T:%.*]] = tuple ()
+// CHECK-NEXT:   return [[T]] : $()
+class C: NSObject {
+  deinit {
+    getDispatchQueue().sync(execute: { _ = self })
+  }
+}