run is optimized with fast-path case and no longer has CoroutineScope in its block

Author: elizarov

kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/Builders.kt

@@ -18,6 +18,8 @@ package kotlinx.coroutines.experimental
 
 import java.util.concurrent.locks.LockSupport
 import kotlin.coroutines.experimental.*
+import kotlin.coroutines.experimental.intrinsics.startCoroutineUninterceptedOrReturn
+import kotlin.coroutines.experimental.intrinsics.suspendCoroutineOrReturn
 
 // --------------- basic coroutine builders ---------------
 
@@ -58,10 +60,28 @@ fun launch(context: CoroutineContext, start: Boolean = true, block: suspend Coro
  * different thread inside the block, and back when it completes.
  * The specified [context] is added onto the current coroutine context for the execution of the block.
  */
-public suspend fun <T> run(context: CoroutineContext, block: suspend CoroutineScope.() -> T): T =
-    suspendCoroutine { cont ->
-        // new don't invoke `newCoroutineContext`, but consider this being the same coroutine in the new context
-        InnerCoroutine(cont.context + context, cont).also { block.startCoroutine(it, it) }
+public suspend fun <T> run(context: CoroutineContext, block: suspend () -> T): T =
+    suspendCoroutineOrReturn sc@ { cont ->
+        val oldContext = cont.context
+        // fast path #1 if there is no change in the actual context:
+        if (context === oldContext || context is CoroutineContext.Element && oldContext[context.key] === context)
+            return@sc block.startCoroutineUninterceptedOrReturn(cont)
+        // compute new context
+        val newContext = oldContext + context
+        // fast path #2 if the result is actually the same
+        if (newContext === oldContext)
+            return@sc block.startCoroutineUninterceptedOrReturn(cont)
+        // fast path #3 if the new dispatcher is the same as the old one
+        if (newContext[ContinuationInterceptor] === oldContext[ContinuationInterceptor]) {
+            val newContinuation = RunContinuationDirect(newContext, cont)
+            return@sc block.startCoroutineUninterceptedOrReturn(newContinuation)
+        }
+        // slowest path otherwise -- use new interceptor, sync to its result via a
+        // full-blown instance of CancellableContinuation
+        val newContinuation = RunContinuationCoroutine(newContext, cont)
+        newContinuation.initCancellability()
+        block.startCoroutine(newContinuation)
+        newContinuation.getResult()
     }
 
 /**
@@ -111,12 +131,15 @@ private class LazyStandaloneCoroutine(
     }
 }
 
-private class InnerCoroutine<in T>(
+private class RunContinuationDirect<in T>(
     override val context: CoroutineContext,
     continuation: Continuation<T>
-) : Continuation<T> by continuation, CoroutineScope {
-    override val isActive: Boolean = context[Job]?.isActive ?: true
-}
+) : Continuation<T> by continuation
+
+private class RunContinuationCoroutine<in T>(
+    override val parentContext: CoroutineContext,
+    continuation: Continuation<T>
+) : CancellableContinuationImpl<T>(continuation, active = true)
 
 private class BlockingCoroutine<T>(
     override val parentContext: CoroutineContext,

kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/CancellableContinuation.kt

@@ -179,7 +179,7 @@ internal open class CancellableContinuationImpl<in T>(
     }
 
     override fun initCancellability() {
-        initParentJob(delegate.context[Job])
+        initParentJob(parentContext[Job])
     }
 
     @PublishedApi

kotlinx-coroutines-core/src/test/kotlin/kotlinx/coroutines/experimental/RunTest.kt

@@ -0,0 +1,127 @@
+/*
+ * Copyright 2016-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kotlinx.coroutines.experimental
+
+import org.hamcrest.MatcherAssert.assertThat
+import org.hamcrest.core.IsEqual
+import org.junit.Test
+
+class RunTest : TestBase() {
+    @Test
+    fun testSameContextNoSuspend() = runBlocking<Unit> {
+        expect(1)
+        launch(context) { // make sure there is not early dispatch here
+            finish(5)
+        }
+        expect(2)
+        val result = run(context) { // same context!
+            expect(3) // still here
+            "OK"
+        }
+        assertThat(result, IsEqual("OK"))
+        expect(4)
+    }
+
+    @Test
+    fun testSameContextWithSuspend() = runBlocking<Unit> {
+        expect(1)
+        launch(context) { // make sure there is not early dispatch here
+            expect(4)
+        }
+        expect(2)
+        val result = run(context) { // same context!
+            expect(3) // still here
+            yield() // now yields to launch!
+            expect(5)
+            "OK"
+        }
+        assertThat(result, IsEqual("OK"))
+        finish(6)
+    }
+
+    @Test
+    fun testCancelWithJobNoSuspend() = runBlocking<Unit> {
+        expect(1)
+        launch(context) { // make sure there is not early dispatch to here
+            finish(6)
+        }
+        expect(2)
+        val job = Job()
+        val result = run(context + job) { // same context + new job
+            expect(3) // still here
+            job.cancel() // cancel out job!
+            try {
+                yield() // shall throw CancellationException
+                expectUnreached()
+            } catch (e: CancellationException) {
+                expect(4)
+            }
+            "OK"
+        }
+        assertThat(result, IsEqual("OK"))
+        expect(5)
+    }
+
+    @Test
+    fun testCancelWithJobWithSuspend() = runBlocking<Unit> {
+        expect(1)
+        launch(context) { // make sure there is not early dispatch to here
+            expect(4)
+        }
+        expect(2)
+        val job = Job()
+        val result = run(context + job) { // same context + new job
+            expect(3) // still here
+            yield() // now yields to launch!
+            expect(5)
+            job.cancel() // cancel out job!
+            try {
+                yield() // shall throw CancellationExpcetion
+                expectUnreached()
+            } catch (e: CancellationException) {
+                expect(6)
+            }
+            "OK"
+        }
+        assertThat(result, IsEqual("OK"))
+        finish(7)
+    }
+
+    @Test
+    fun testCommonPoolNoSuspend() = runBlocking<Unit> {
+        expect(1)
+        val result = run(CommonPool) {
+            expect(2)
+            "OK"
+        }
+        assertThat(result, IsEqual("OK"))
+        finish(3)
+    }
+
+    @Test
+    fun testCommonPoolWithSuspend() = runBlocking<Unit> {
+        expect(1)
+        val result = run(CommonPool) {
+            expect(2)
+            delay(100)
+            expect(3)
+            "OK"
+        }
+        assertThat(result, IsEqual("OK"))
+        finish(4)
+    }
+}