From b961812efd922ed4492a9cbf29320c0fc71f03c7 Mon Sep 17 00:00:00 2001
From: Svyatoslav Scherbina <svyatoslav.scherbina@jetbrains.com>
Date: Fri, 18 Jun 2021 18:37:40 +0300
Subject: [PATCH] Native: fix thread state when calling ObjC completion for
 Kotlin suspend

---
 .../tests/objcexport/coroutines.swift         | 35 +++++++++++++++++++
 .../src/main/cpp/ObjCExportCoroutines.mm      |  6 +++-
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/kotlin-native/backend.native/tests/objcexport/coroutines.swift b/kotlin-native/backend.native/tests/objcexport/coroutines.swift
index eab64e65d9a..eebee70e555 100644
--- a/kotlin-native/backend.native/tests/objcexport/coroutines.swift
+++ b/kotlin-native/backend.native/tests/objcexport/coroutines.swift
@@ -95,6 +95,40 @@ private func testCall() throws {
     try testSuspendFuncAsync(doThrow: true)
 }
 
+private func testCallSuspendFunChain(doSuspend: Bool, doThrow: Bool) throws {
+    class C {}
+    let expectedResult = C()
+
+    var completionCalled = 0
+    var result: AnyObject? = nil
+    var error: Error? = nil
+
+    CoroutinesKt.suspendFun(result: expectedResult, doSuspend: doSuspend, doThrow: doThrow) { _resultOuter, _errorOuter in
+        CoroutinesKt.suspendFun(result: expectedResult, doSuspend: doSuspend, doThrow: doThrow) { _result, _error in
+            completionCalled += 1
+            result = _result as AnyObject?
+            error = _error
+        }
+    }
+
+    try assertEquals(actual: completionCalled, expected: 1)
+
+    if doThrow {
+        try assertNil(result)
+        try assertTrue(error?.kotlinException is CoroutineException)
+    } else {
+        try assertSame(actual: result, expected: expectedResult)
+        try assertNil(error)
+    }
+}
+
+private func testCallChain() throws {
+    try testCallSuspendFunChain(doSuspend: true, doThrow: false)
+    try testCallSuspendFunChain(doSuspend: false, doThrow: false)
+    try testCallSuspendFunChain(doSuspend: true, doThrow: true)
+    try testCallSuspendFunChain(doSuspend: false, doThrow: true)
+}
+
 private class SuspendFunImpl : SuspendFun {
     class E : Error {}
 
@@ -336,6 +370,7 @@ class CoroutinesTests : SimpleTestProvider {
 
         test("TestCallSimple", testCallSimple)
         test("TestCall", testCall)
+        test("TestCallChain", testCallChain)
         test("TestOverride", testOverride)
         test("TestBridges", testBridges)
         test("TestImplicitThrows1", testImplicitThrows1)
diff --git a/kotlin-native/runtime/src/main/cpp/ObjCExportCoroutines.mm b/kotlin-native/runtime/src/main/cpp/ObjCExportCoroutines.mm
index d041ccf78d1..f7022ab87eb 100644
--- a/kotlin-native/runtime/src/main/cpp/ObjCExportCoroutines.mm
+++ b/kotlin-native/runtime/src/main/cpp/ObjCExportCoroutines.mm
@@ -9,6 +9,7 @@
 #import <Foundation/NSException.h>
 #import <Foundation/NSObject.h>
 
+#import "Memory.h"
 #import "ObjCExport.h"
 #import "ObjCExportErrors.h"
 
@@ -16,7 +17,9 @@ typedef void (^Completion)(id _Nullable, NSError* _Nullable);
 
 extern "C" void Kotlin_ObjCExport_runCompletionSuccess(KRef completionHolder, KRef result) {
   Completion completion = (Completion)GetAssociatedObject(completionHolder);
-  completion(Kotlin_ObjCExport_refToObjC(result), nullptr);
+  id objCResult = Kotlin_ObjCExport_refToObjC(result);
+  kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
+  completion(objCResult, nullptr);
 }
 
 extern "C" void Kotlin_ObjCExport_runCompletionFailure(
@@ -26,6 +29,7 @@ extern "C" void Kotlin_ObjCExport_runCompletionFailure(
 ) {
   id error = Kotlin_ObjCExport_ExceptionAsNSError(exception, exceptionTypes);
   Completion completion = (Completion)GetAssociatedObject(completionHolder);
+  kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
   completion(nullptr, error);
 }