Wrap continuation with ContinuationImpl on callable references

in startCoroutineUninterceptedOrReturn. Otherwise, the coroutine will
not be interceptable later.

Add a test, which checks, that intercepted continuation is released.

 #KT-55869

libraries/stdlib/js-v1/src/kotlin/coroutines/intrinsics/IntrinsicsJs.kt

@@ -24,7 +24,7 @@ import kotlin.internal.InlineOnly
 @InlineOnly
 public actual inline fun <T> (suspend () -> T).startCoroutineUninterceptedOrReturn(
     completion: Continuation<T>
-): Any? = this.asDynamic()(completion, false)
+): Any? = this.asDynamic()(if (this !is CoroutineImpl) createSimpleCoroutineFromSuspendFunction(completion) else completion, false)
 
 /**
  * Starts an unintercepted coroutine with receiver type [R] and result type [T] and executes it until its first suspension.
@@ -43,14 +43,18 @@ public actual inline fun <T> (suspend () -> T).startCoroutineUninterceptedOrRetu
 public actual inline fun <R, T> (suspend R.() -> T).startCoroutineUninterceptedOrReturn(
     receiver: R,
     completion: Continuation<T>
-): Any? = this.asDynamic()(receiver, completion, false)
+): Any? = this.asDynamic()(
+    receiver, if (this !is CoroutineImpl) createSimpleCoroutineFromSuspendFunction(completion) else completion, false
+)
 
 @InlineOnly
 internal actual inline fun <R, P, T> (suspend R.(P) -> T).startCoroutineUninterceptedOrReturn(
     receiver: R,
     param: P,
     completion: Continuation<T>
-): Any? = this.asDynamic()(receiver, param, completion, false)
+): Any? = this.asDynamic()(
+    receiver, param, if (this !is CoroutineImpl) createSimpleCoroutineFromSuspendFunction(completion) else completion, false
+)
 
 /**
  * Creates unintercepted coroutine without receiver and with result type [T].
@@ -149,3 +153,13 @@ private inline fun <T> createCoroutineFromSuspendFunction(
         }
     }
 }
+
+@PublishedApi
+internal fun <T> createSimpleCoroutineFromSuspendFunction(
+    completion: Continuation<T>
+): CoroutineImpl = object : CoroutineImpl(completion as Continuation<Any?>) {
+    override fun doResume(): Any? {
+        if (exception != null) throw exception as Throwable
+        return result
+    }
+}

libraries/stdlib/jvm/src/kotlin/coroutines/intrinsics/IntrinsicsJvm.kt

@@ -29,7 +29,19 @@ import kotlin.internal.InlineOnly
 @InlineOnly
 public actual inline fun <T> (suspend () -> T).startCoroutineUninterceptedOrReturn(
     completion: Continuation<T>
-): Any? = (this as Function1<Continuation<T>, Any?>).invoke(completion)
+): Any? =
+    // Wrap with ContinuationImpl, otherwise the coroutine will not be interceptable. See KT-55869
+    if (this !is BaseContinuationImpl) wrapWithContinuationImpl(completion)
+    else (this as Function1<Continuation<T>, Any?>).invoke(completion)
+
+// Work around private and internal visibilities of functions used: [createCoroutineFromSuspendFunction] and [probeCoroutineCreated].
+@PublishedApi
+internal fun <T> (suspend () -> T).wrapWithContinuationImpl(
+    completion: Continuation<T>
+): Any? {
+    val newCompletion = createSimpleCoroutineForSuspendFunction(probeCoroutineCreated(completion))
+    return (this as Function1<Continuation<T>, Any?>).invoke(newCompletion)
+}
 
 /**
  * Starts an unintercepted coroutine with receiver type [R] and result type [T] and executes it until its first suspension.
@@ -48,14 +60,41 @@ public actual inline fun <T> (suspend () -> T).startCoroutineUninterceptedOrRetu
 public actual inline fun <R, T> (suspend R.() -> T).startCoroutineUninterceptedOrReturn(
     receiver: R,
     completion: Continuation<T>
-): Any? = (this as Function2<R, Continuation<T>, Any?>).invoke(receiver, completion)
+): Any? =
+    // Wrap with ContinuationImpl, otherwise the coroutine will not be interceptable. See KT-55869
+    if (this !is BaseContinuationImpl) wrapWithContinuationImpl(receiver, completion)
+    else (this as Function2<R, Continuation<T>, Any?>).invoke(receiver, completion)
+
+// Work around private and internal visibilities of functions used: [createCoroutineFromSuspendFunction] and [probeCoroutineCreated].
+@PublishedApi
+internal fun <R, T> (suspend R.() -> T).wrapWithContinuationImpl(
+    receiver: R,
+    completion: Continuation<T>
+): Any? {
+    val newCompletion = createSimpleCoroutineForSuspendFunction(probeCoroutineCreated(completion))
+    return (this as Function2<R, Continuation<T>, Any?>).invoke(receiver, newCompletion)
+}
 
 @InlineOnly
 internal actual inline fun <R, P, T> (suspend R.(P) -> T).startCoroutineUninterceptedOrReturn(
     receiver: R,
     param: P,
     completion: Continuation<T>
-): Any? = (this as Function3<R, P, Continuation<T>, Any?>).invoke(receiver, param, completion)
+): Any? =
+    // Wrap with ContinuationImpl, otherwise the coroutine will not be interceptable. See KT-55869
+    if (this !is BaseContinuationImpl) wrapWithContinuationImpl(receiver, param, completion)
+    else (this as Function3<R, P, Continuation<T>, Any?>).invoke(receiver, param, completion)
+
+// Work around private and internal visibilities of functions used: [createCoroutineFromSuspendFunction] and [probeCoroutineCreated].
+@PublishedApi
+internal fun <R, P, T> (suspend R.(P) -> T).wrapWithContinuationImpl(
+    receiver: R,
+    param: P,
+    completion: Continuation<T>
+): Any? {
+    val newCompletion = createSimpleCoroutineForSuspendFunction(probeCoroutineCreated(completion))
+    return (this as Function3<R, P, Continuation<T>, Any?>).invoke(receiver, param, newCompletion)
+}
 
 // JVM declarations
 
@@ -201,3 +240,30 @@ private inline fun <T> createCoroutineFromSuspendFunction(
                 }
         }
 }
+
+/**
+ * This function is used when [startCoroutineUninterceptedOrReturn] encounters suspending lambda that does not extend BaseContinuationImpl.
+ *
+ * It happens in two cases:
+ *   1. Callable reference to suspending function or tail-call lambdas,
+ *   2. Suspending function reference implemented by Java code.
+ *
+ * This function is the same as above, but does not run lambda itself - the caller is expected to call [invoke] manually.
+ */
+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()
+            }
+        }
+}

libraries/stdlib/wasm/src/kotlin/coroutines/CoroutinesIntrinsics.kt

@@ -24,7 +24,9 @@ import kotlin.wasm.internal.*
 @kotlin.internal.InlineOnly
 public actual inline fun <T> (suspend () -> T).startCoroutineUninterceptedOrReturn(
     completion: Continuation<T>
-): Any? = startCoroutineUninterceptedOrReturnIntrinsic0(this, completion)
+): Any? = startCoroutineUninterceptedOrReturnIntrinsic0(
+    this, if (this !is CoroutineImpl) createSimpleCoroutineFromSuspendFunction(completion) else completion
+)
 
 /**
  * Starts an unintercepted coroutine with receiver type [R] and result type [T] and executes it until its first suspension.
@@ -43,7 +45,9 @@ public actual inline fun <T> (suspend () -> T).startCoroutineUninterceptedOrRetu
 public actual inline fun <R, T> (suspend R.() -> T).startCoroutineUninterceptedOrReturn(
     receiver: R,
     completion: Continuation<T>
-): Any? = startCoroutineUninterceptedOrReturnIntrinsic1(this, receiver, completion)
+): Any? = startCoroutineUninterceptedOrReturnIntrinsic1(
+    this, receiver, if (this !is CoroutineImpl) createSimpleCoroutineFromSuspendFunction(completion) else completion
+)
 
 @Suppress("UNCHECKED_CAST")
 @kotlin.internal.InlineOnly
@@ -51,7 +55,9 @@ internal actual inline fun <R, P, T> (suspend R.(P) -> T).startCoroutineUninterc
     receiver: R,
     param: P,
     completion: Continuation<T>
-): Any? = startCoroutineUninterceptedOrReturnIntrinsic2(this, receiver, param, completion)
+): Any? = startCoroutineUninterceptedOrReturnIntrinsic2(
+    this, receiver, param, if (this !is CoroutineImpl) createSimpleCoroutineFromSuspendFunction(completion) else completion
+)
 
 /**
  * Creates unintercepted coroutine without receiver and with result type [T].
@@ -137,4 +143,14 @@ private inline fun <T> createCoroutineFromSuspendFunction(
             return block()
         }
     }
-}
+}
+
+@PublishedApi
+internal fun <T> createSimpleCoroutineFromSuspendFunction(
+    completion: Continuation<T>
+): CoroutineImpl = object : CoroutineImpl(completion as Continuation<Any?>) {
+    override fun doResume(): Any? {
+        if (exception != null) throw exception as Throwable
+        return result
+    }
+}