ObjCExport: test that completion is properly retained by continuation

Add a test that checks that Kotlin Continuation wrapping Objective-C
completion handler has strong reference to it, i.e. completion handler
doesn't get reclaimed too early.

Follow-up to 4dcfd38.

kotlin-native/backend.native/tests/objcexport/coroutines.kt

@@ -45,14 +45,16 @@ suspend fun unitSuspendFun(doSuspend: Boolean, doThrow: Boolean) {
 }
 
 class ContinuationHolder<T> {
-    internal lateinit var continuation: Continuation<T>
+    internal var continuation: Continuation<T>? = null
 
     fun resume(value: T) {
-        continuation.resume(value)
+        continuation!!.resume(value)
+        continuation = null
     }
 
     fun resumeWithException(exception: Throwable) {
-        continuation.resumeWithException(exception)
+        continuation!!.resumeWithException(exception)
+        continuation = null
     }
 }
 
@@ -219,3 +221,5 @@ suspend fun invoke1(block: suspend (Any?) -> Any?, argument: Any?): Any? = block
 
 fun getKSuspendCallableReference0(): KSuspendFunction0<String> = ::suspendCallableReference0Target
 fun getKSuspendCallableReference1(): KSuspendFunction1<String, String> = ::suspendCallableReference1Target
+
+fun gc() = kotlin.native.internal.GC.collect()

kotlin-native/backend.native/tests/objcexport/coroutines.swift

@@ -106,22 +106,54 @@ private func testCallUnitSuspendFun(doSuspend: Bool, doThrow: Bool) throws {
     }
 }
 
+private class WeakRefHolder {
+    weak var value: AnyObject? = nil
+}
+
+#if NO_GENERICS
+typealias AnyContinuationHolder = ContinuationHolder
+#else
+typealias AnyContinuationHolder = ContinuationHolder<AnyObject>
+#endif
+
+// This code is extracted to a function just to ensure that all local variables get released at the end.
+private func callSuspendFunAsync(
+    weakRefToObjectCapturedByCompletion: WeakRefHolder,
+    continuationHolder: AnyContinuationHolder,
+    completionHandler: @escaping (Any?, Error?) -> Void
+) throws {
+    class C {}
+    let capturedByCompletion = C()
+    weakRefToObjectCapturedByCompletion.value = capturedByCompletion
+
+    CoroutinesKt.suspendFunAsync(result: nil, continuationHolder: continuationHolder) { _result, _error in
+        try! assertSame(actual: capturedByCompletion, expected: weakRefToObjectCapturedByCompletion.value)
+        completionHandler(_result, _error)
+    }
+}
+
 private func testSuspendFuncAsync(doThrow: Bool) throws {
     var completionCalled = 0
     var result: AnyObject? = nil
     var error: Error? = nil
 
-#if NO_GENERICS
-    let continuationHolder = ContinuationHolder()
-#else
-    let continuationHolder = ContinuationHolder<AnyObject>()
-#endif
+    let continuationHolder = AnyContinuationHolder()
 
-    CoroutinesKt.suspendFunAsync(result: nil, continuationHolder: continuationHolder) { _result, _error in
-        completionCalled += 1
-        result = _result as AnyObject?
-        error = _error
+    let weakRefToObjectCapturedByCompletion = WeakRefHolder()
+    try assertTrue(weakRefToObjectCapturedByCompletion.value === nil)
+    try autoreleasepool {
+        try callSuspendFunAsync(
+            weakRefToObjectCapturedByCompletion: weakRefToObjectCapturedByCompletion,
+            continuationHolder: continuationHolder
+        ) { _result, _error in
+            completionCalled += 1
+            result = _result as AnyObject?
+            error = _error
+        }
     }
+    CoroutinesKt.gc()
+    // This assert checks that suspendFunAsync retains the completion handler:
+    try assertFalse(weakRefToObjectCapturedByCompletion.value === nil)
 
     try assertEquals(actual: completionCalled, expected: 0)
 
@@ -143,31 +175,62 @@ private func testSuspendFuncAsync(doThrow: Bool) throws {
         try assertSame(actual: result, expected: expectedResult)
         try assertNil(error)
     }
+
+#if !NOOP_GC
+    CoroutinesKt.gc()
+    // This assert checks that the completion handler gets properly released after all:
+    try assertTrue(weakRefToObjectCapturedByCompletion.value === nil)
+#endif
+}
+
+#if NO_GENERICS
+typealias UnitContinuationHolder = ContinuationHolder
+#else
+typealias UnitContinuationHolder = ContinuationHolder<KotlinUnit>
+#endif
+
+// This code is extracted to a function just to ensure that all local variables get released at the end.
+private func callUnitSuspendFunAsync(
+    weakRefToObjectCapturedByCompletion: WeakRefHolder,
+    continuationHolder: UnitContinuationHolder,
+    completionHandler: @escaping (Error?) -> Void
+) throws {
+    class C {}
+    let capturedByCompletion = C()
+    weakRefToObjectCapturedByCompletion.value = capturedByCompletion
+#if LEGACY_SUSPEND_UNIT_FUNCTION_EXPORT
+    CoroutinesKt.unitSuspendFunAsync(continuationHolder: continuationHolder) { _result, _error in
+        try! assertSame(actual: capturedByCompletion, expected: weakRefToObjectCapturedByCompletion.value)
+        completionHandler(_error)
+    }
+#else
+    CoroutinesKt.unitSuspendFunAsync(continuationHolder: continuationHolder) { _error in
+        try! assertSame(actual: capturedByCompletion, expected: weakRefToObjectCapturedByCompletion.value)
+        completionHandler(_error)
+    }
+#endif
 }
 
 private func testUnitSuspendFuncAsync(doThrow: Bool) throws {
     var completionCalled = 0
-    var result: AnyObject? = nil
     var error: Error? = nil
 
-#if NO_GENERICS
-    let continuationHolder = ContinuationHolder()
-#else
-    let continuationHolder = ContinuationHolder<KotlinUnit>()
-#endif
+    let continuationHolder = UnitContinuationHolder()
 
-#if LEGACY_SUSPEND_UNIT_FUNCTION_EXPORT
-    CoroutinesKt.unitSuspendFunAsync(continuationHolder: continuationHolder) { _result, _error in
-        completionCalled += 1
-        result = _result as AnyObject?
-        error = _error
+    let weakRefToObjectCapturedByCompletion = WeakRefHolder()
+    try assertTrue(weakRefToObjectCapturedByCompletion.value === nil)
+    try autoreleasepool {
+        try callUnitSuspendFunAsync(
+            weakRefToObjectCapturedByCompletion: weakRefToObjectCapturedByCompletion,
+            continuationHolder: continuationHolder
+        ) { _error in
+            completionCalled += 1
+            error = _error
+        }
     }
-#else
-    CoroutinesKt.unitSuspendFunAsync(continuationHolder: continuationHolder) { _error in
-        completionCalled += 1
-        error = _error
-    }
-#endif
+    CoroutinesKt.gc()
+    // This assert checks that unitSuspendFunAsync retains the completion handler:
+    try assertFalse(weakRefToObjectCapturedByCompletion.value === nil)
 
     try assertEquals(actual: completionCalled, expected: 0)
 
@@ -177,20 +240,20 @@ private func testUnitSuspendFuncAsync(doThrow: Bool) throws {
 
         try assertEquals(actual: completionCalled, expected: 1)
 
-#if LEGACY_SUSPEND_UNIT_FUNCTION_EXPORT
-        try assertNil(result)
-#endif
         try assertSame(actual: error?.kotlinException as AnyObject?, expected: exception)
     } else {
         continuationHolder.resume(value: KotlinUnit.shared)
 
         try assertEquals(actual: completionCalled, expected: 1)
 
-#if LEGACY_SUSPEND_UNIT_FUNCTION_EXPORT
-        try assertSame(actual: result, expected: KotlinUnit.shared)
-#endif
         try assertNil(error)
     }
+
+#if !NOOP_GC
+    CoroutinesKt.gc()
+    // This assert checks that the completion handler gets properly released after all:
+    try assertTrue(weakRefToObjectCapturedByCompletion.value === nil)
+#endif
 }
 
 private func testCall() throws {

kotlin-native/backend.native/tests/objcexport/expectedLazy.h

@@ -382,6 +382,7 @@ __attribute__((swift_name("CoroutinesKt")))
 + (void)invoke1Block:(id<KtKotlinSuspendFunction1>)block argument:(id _Nullable)argument completionHandler:(void (^)(id _Nullable_result, NSError * _Nullable))completionHandler __attribute__((swift_name("invoke1(block:argument:completionHandler:)")));
 + (id<KtKotlinKSuspendFunction0>)getKSuspendCallableReference0 __attribute__((swift_name("getKSuspendCallableReference0()")));
 + (id<KtKotlinKSuspendFunction1>)getKSuspendCallableReference1 __attribute__((swift_name("getKSuspendCallableReference1()")));
++ (void)gc __attribute__((swift_name("gc()")));
 @end;
 
 __attribute__((swift_name("DeallocRetainBase")))

kotlin-native/backend.native/tests/objcexport/expectedLazyLegacySuspendUnit.h

@@ -382,6 +382,7 @@ __attribute__((swift_name("CoroutinesKt")))
 + (void)invoke1Block:(id<KtKotlinSuspendFunction1>)block argument:(id _Nullable)argument completionHandler:(void (^)(id _Nullable_result, NSError * _Nullable))completionHandler __attribute__((swift_name("invoke1(block:argument:completionHandler:)")));
 + (id<KtKotlinKSuspendFunction0>)getKSuspendCallableReference0 __attribute__((swift_name("getKSuspendCallableReference0()")));
 + (id<KtKotlinKSuspendFunction1>)getKSuspendCallableReference1 __attribute__((swift_name("getKSuspendCallableReference1()")));
++ (void)gc __attribute__((swift_name("gc()")));
 @end;
 
 __attribute__((swift_name("DeallocRetainBase")))

kotlin-native/backend.native/tests/objcexport/expectedLazyNoGenerics.h

@@ -382,6 +382,7 @@ __attribute__((swift_name("CoroutinesKt")))
 + (void)invoke1Block:(id<KtKotlinSuspendFunction1>)block argument:(id _Nullable)argument completionHandler:(void (^)(id _Nullable_result, NSError * _Nullable))completionHandler __attribute__((swift_name("invoke1(block:argument:completionHandler:)")));
 + (id<KtKotlinKSuspendFunction0>)getKSuspendCallableReference0 __attribute__((swift_name("getKSuspendCallableReference0()")));
 + (id<KtKotlinKSuspendFunction1>)getKSuspendCallableReference1 __attribute__((swift_name("getKSuspendCallableReference1()")));
++ (void)gc __attribute__((swift_name("gc()")));
 @end;
 
 __attribute__((swift_name("DeallocRetainBase")))