Fixed createCoroutineFromSuspendFunction

Passes an object that extends BaseContinuationImpl to the
suspend function invocation, so that it can use .intercepted on it.
This requires implementation to track its state via label variable.

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

@@ -82,7 +82,7 @@ public actual fun <T> (suspend () -> T).createCoroutineUnintercepted(
         create(completion)
     else
         createCoroutineFromSuspendFunction(completion) {
-            (this as Function1<Continuation<T>, Any?>).invoke(completion)
+            (this as Function1<Continuation<T>, Any?>).invoke(it)
         }
 
 /**
@@ -115,7 +115,7 @@ public actual fun <R, T> (suspend R.() -> T).createCoroutineUnintercepted(
         create(receiver, completion)
     else {
         createCoroutineFromSuspendFunction(completion) {
-            (this as Function2<R, Continuation<T>, Any?>).invoke(receiver, completion)
+            (this as Function2<R, Continuation<T>, Any?>).invoke(receiver, it)
         }
     }
 
@@ -144,25 +144,50 @@ public actual fun <T> Continuation<T>.intercepted(): Continuation<T> =
  * We must wrap it into an instance that extends [BaseContinuationImpl], because that is an expectation of all coroutines machinery.
  * As an optimization we use lighter-weight [RestrictedContinuationImpl] base class (it has less fields) if the context is
  * [EmptyCoroutineContext], and a full-blown [ContinuationImpl] class otherwise.
+ *
+ * The instance of [BaseContinuationImpl] is passed to the [block] so that it can be passed to the corresponding invocation.
  */
 @SinceKotlin("1.3")
 private inline fun <T> createCoroutineFromSuspendFunction(
     completion: Continuation<T>,
-    crossinline block: () -> Any?
+    crossinline block: (Continuation<T>) -> Any?
 ): Continuation<Unit> {
     val context = completion.context
+    // label == 0 when coroutine is not started yet (initially) or label == 1 when it was
     return if (context === EmptyCoroutineContext)
         object : RestrictedContinuationImpl(completion as Continuation<Any?>) {
-            override fun invokeSuspend(result: SuccessOrFailure<Any?>): Any? {
-                result.getOrThrow() // Rethrow exception if trying to start with exception (will be caught by BaseContinuationImpl.resumeWith
-                return block() // run the block
-            }
+            private var label = 0
+
+            override fun invokeSuspend(result: SuccessOrFailure<Any?>): Any? =
+                when (label) {
+                    0 -> {
+                        label = 1
+                        result.getOrThrow() // Rethrow exception if trying to start with exception (will be caught by BaseContinuationImpl.resumeWith
+                        block(this) // run the block, may return or suspend
+                    }
+                    1 -> {
+                        label = 2
+                        result.getOrThrow() // this is the result if the block had suspended
+                    }
+                    else -> error("This coroutine had already completed")
+                }
         }
     else
         object : ContinuationImpl(completion as Continuation<Any?>, context) {
-            override fun invokeSuspend(result: SuccessOrFailure<Any?>): Any? {
-                result.getOrThrow() // Rethrow exception if trying to start with exception (will be caught by BaseContinuationImpl.resumeWith
-                return block() // run the block
-            }
+            private var label = 0
+            
+            override fun invokeSuspend(result: SuccessOrFailure<Any?>): Any? =
+                when (label) {
+                    0 -> {
+                        label = 1
+                        result.getOrThrow() // Rethrow exception if trying to start with exception (will be caught by BaseContinuationImpl.resumeWith
+                        block(this) // run the block, may return or suspend
+                    }
+                    1 -> {
+                        label = 2
+                        result.getOrThrow() // this is the result if the block had suspended
+                    }
+                    else -> error("This coroutine had already completed")
+                }
         }
 }

libraries/stdlib/coroutines/jvm/test/CoroutinesTest.kt

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package test.kotlin.coroutines
+
+import java.io.Closeable
+import java.util.concurrent.ExecutorService
+import java.util.concurrent.Executors
+import java.util.concurrent.Semaphore
+import kotlin.coroutines.*
+import kotlin.coroutines.experimental.startCoroutine
+import kotlin.test.*
+
+/**
+ * Tests on coroutines standard library functions.
+ */
+class CoroutinesTest {
+    /**
+     * Makes sure that using [startCoroutine] with suspending references properly establishes intercepted context.
+     */
+    @Test
+    fun testStartInterceptedSuspendReference() {
+        val done = Semaphore(0)
+        TestDispatcher("Result").use { resumeDispatcher ->
+            TestDispatcher("Context").use { contextDispatcher ->
+                val switcher = DispatcherSwitcher(contextDispatcher, resumeDispatcher)
+                val ref = switcher::run // callable reference
+                ref.startCoroutine(Continuation(contextDispatcher) { result ->
+                    contextDispatcher.assertThread()
+                    // todo: below does not work due to a bug in inline classes
+//                    assertEquals(42, result.getOrThrow())
+                    done.release()
+                })
+                done.acquire()
+            }
+        }
+    }
+
+}
+
+class DispatcherSwitcher(
+    private val contextDispatcher: TestDispatcher,
+    private val resumeDispatcher: TestDispatcher
+) {
+    suspend fun run(): Int = suspendCoroutine { cont ->
+        contextDispatcher.assertThread()
+        resumeDispatcher.executor.execute {
+            resumeDispatcher.assertThread()
+            cont.resume(42)
+        }
+    }
+}
+
+class TestDispatcher(
+    private val name: String
+) : AbstractCoroutineContextElement(ContinuationInterceptor), ContinuationInterceptor, Closeable {
+    private lateinit var thread: Thread
+
+    val executor: ExecutorService = Executors.newSingleThreadExecutor { runnable ->
+        Thread(runnable, name).also { thread = it }
+    }
+
+    fun assertThread() {
+        assertEquals(thread, Thread.currentThread())
+    }
+
+    override fun <T> interceptContinuation(continuation: Continuation<T>): Continuation<T> =
+        DispatchedContinuation(continuation)
+
+    override fun close() {
+        executor.shutdown()
+    }
+
+    inner class DispatchedContinuation<T>(val delegate: Continuation<T>) : Continuation<T> {
+        override val context: CoroutineContext = delegate.context
+
+        override fun resumeWith(result: SuccessOrFailure<T>) {
+            executor.execute {
+                delegate.resumeWith(result)
+            }
+        }
+    }
+}
+