From 9d911247fa2c97485299fca277cc66ff70ede758 Mon Sep 17 00:00:00 2001
From: Igor Chevdar <igor.chevdar@jetbrains.com>
Date: Thu, 19 Jan 2023 18:46:57 +0200
Subject: [PATCH] [K/N] Wrap continuation with ContinuationImpl on callable
 references

 #KT-55869
---
 .../intrinsicSemantics/intercepted.kt         |  2 +-
 .../intrinsicSemantics/releaseIntercepted.kt  |  2 +-
 ...outineUninterceptedOrReturnInterception.kt |  1 -
 .../coroutines/intrinsics/IntrinsicsNative.kt | 59 +++++++++++++++----
 4 files changed, 51 insertions(+), 13 deletions(-)

diff --git a/compiler/testData/codegen/box/coroutines/intrinsicSemantics/intercepted.kt b/compiler/testData/codegen/box/coroutines/intrinsicSemantics/intercepted.kt
index f76f9a726ec..544c15a3b5a 100644
--- a/compiler/testData/codegen/box/coroutines/intrinsicSemantics/intercepted.kt
+++ b/compiler/testData/codegen/box/coroutines/intrinsicSemantics/intercepted.kt
@@ -1,6 +1,6 @@
 // WITH_STDLIB
 // WITH_COROUTINES
-// IGNORE_BACKEND: JS_IR, JS_IR_ES6, NATIVE
+// IGNORE_BACKEND: JS_IR, JS_IR_ES6
 import helpers.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
diff --git a/compiler/testData/codegen/box/coroutines/intrinsicSemantics/releaseIntercepted.kt b/compiler/testData/codegen/box/coroutines/intrinsicSemantics/releaseIntercepted.kt
index 08cf443111a..8049e9f10f2 100644
--- a/compiler/testData/codegen/box/coroutines/intrinsicSemantics/releaseIntercepted.kt
+++ b/compiler/testData/codegen/box/coroutines/intrinsicSemantics/releaseIntercepted.kt
@@ -1,5 +1,5 @@
 // WITH_STDLIB
-// IGNORE_BACKEND: JS_IR, JS_IR_ES6, NATIVE
+// IGNORE_BACKEND: JS_IR, JS_IR_ES6
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
 
diff --git a/compiler/testData/codegen/box/coroutines/intrinsicSemantics/startCoroutineUninterceptedOrReturnInterception.kt b/compiler/testData/codegen/box/coroutines/intrinsicSemantics/startCoroutineUninterceptedOrReturnInterception.kt
index 038e072810d..08ff9455767 100644
--- a/compiler/testData/codegen/box/coroutines/intrinsicSemantics/startCoroutineUninterceptedOrReturnInterception.kt
+++ b/compiler/testData/codegen/box/coroutines/intrinsicSemantics/startCoroutineUninterceptedOrReturnInterception.kt
@@ -1,6 +1,5 @@
 // WITH_STDLIB
 // WITH_COROUTINES
-// IGNORE_BACKEND: NATIVE
 import helpers.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
diff --git a/kotlin-native/runtime/src/main/kotlin/kotlin/coroutines/intrinsics/IntrinsicsNative.kt b/kotlin-native/runtime/src/main/kotlin/kotlin/coroutines/intrinsics/IntrinsicsNative.kt
index ef48e866700..9fbe20dfaf1 100644
--- a/kotlin-native/runtime/src/main/kotlin/kotlin/coroutines/intrinsics/IntrinsicsNative.kt
+++ b/kotlin-native/runtime/src/main/kotlin/kotlin/coroutines/intrinsics/IntrinsicsNative.kt
@@ -27,10 +27,13 @@ public actual inline fun <T> (suspend () -> T).startCoroutineUninterceptedOrRetu
         completion: Continuation<T>
 ): Any? {
     val function = this as? Function1<Continuation<T>, Any?>
-    return if (function != null)
-        function.invoke(completion)
-    else
+    return if (function == null)
         startCoroutineUninterceptedOrReturnFallback(this, completion)
+    else {
+        function.invoke(
+                completion.takeIf { this is BaseContinuationImpl } ?: wrapWithContinuationImpl(completion)
+        )
+    }
 }
 
 @Suppress("UNCHECKED_CAST")
@@ -64,10 +67,13 @@ public actual inline fun <R, T> (suspend R.() -> T).startCoroutineUninterceptedO
         completion: Continuation<T>
 ): Any? {
     val function = this as? Function2<R, Continuation<T>, Any?>
-    return if (function != null)
-        function.invoke(receiver, completion)
-    else
+    return if (function == null)
         startCoroutineUninterceptedOrReturnFallback(this, receiver, completion)
+    else {
+        function.invoke(receiver,
+                completion.takeIf { this is BaseContinuationImpl } ?: wrapWithContinuationImpl(completion)
+        )
+    }
 }
 
 @Suppress("UNCHECKED_CAST")
@@ -91,10 +97,13 @@ internal actual inline fun <R, P, T> (suspend R.(P) -> T).startCoroutineUninterc
         completion: Continuation<T>
 ): Any? {
     val function = this as? Function3<R, P, Continuation<T>, Any?>
-    return if (function != null)
-        function.invoke(receiver, param, completion)
-    else
+    return if (function == null)
         startCoroutineUninterceptedOrReturnFallback(this, receiver, param, completion)
+    else {
+        function.invoke(receiver, param,
+                completion.takeIf { this is BaseContinuationImpl } ?: wrapWithContinuationImpl(completion)
+        )
+    }
 }
 
 @Suppress("UNCHECKED_CAST")
@@ -288,4 +297,34 @@ internal inline fun createContinuationArgumentFromCallback(
 
         return Unit // Not suspended.
     }
-}
\ No newline at end of file
+}
+
+@PublishedApi
+internal fun <T> wrapWithContinuationImpl(completion: Continuation<T>) =
+        createSimpleCoroutineForSuspendFunction(probeCoroutineCreated(completion))
+
+/**
+ * This function is used when [startCoroutineUninterceptedOrReturn] encounters suspending lambda that does not extend BaseContinuationImpl.
+ *
+ * It happens in two cases: callable reference to suspending function or tail-call lambdas.
+ *
+ * This function is the same as above, but does not run lambda itself - the caller is expected to call [invoke] manually.
+ */
+@Suppress("UNCHECKED_CAST")
+private fun <T> createSimpleCoroutineForSuspendFunction(
+        completion: Continuation<T>
+): Continuation<T> {
+    val context = completion.context
+    return if (context === EmptyCoroutineContext)
+        object : RestrictedContinuationImpl(completion as Continuation<Any?>) {
+            override fun invokeSuspend(result: Result<Any?>): Any? {
+                return result.getOrThrow()
+            }
+        }
+    else
+        object : ContinuationImpl(completion as Continuation<Any?>, context) {
+            override fun invokeSuspend(result: Result<Any?>): Any? {
+                return result.getOrThrow()
+            }
+        }
+}