Save more code switching to resumeWith idiom

This change also makes TimeoutCoroutine extends ScopedCoroutines.
In addition to avoid code duplication it also enables to do emit in
flow from withing withTimeout { ... } sections, which would be
properly considered scope.

This change also save more JVM stack entries before resume calls.

Author: elizarov

kotlinx-coroutines-core/common/src/ResumeMode.kt

@@ -4,53 +4,10 @@
 
 package kotlinx.coroutines
 
-import kotlin.coroutines.*
-import kotlin.coroutines.intrinsics.*
-
 @PublishedApi internal const val MODE_ATOMIC_DEFAULT = 0 // schedule non-cancellable dispatch for suspendCoroutine
 @PublishedApi internal const val MODE_CANCELLABLE = 1    // schedule cancellable dispatch for suspendCancellableCoroutine
 @PublishedApi internal const val MODE_DIRECT = 2         // when the context is right just invoke the delegate continuation direct
 @PublishedApi internal const val MODE_UNDISPATCHED = 3   // when the thread is right, but need to mark it with current coroutine
 
 internal val Int.isCancellableMode get() = this == MODE_CANCELLABLE
 internal val Int.isDispatchedMode get() = this == MODE_ATOMIC_DEFAULT || this == MODE_CANCELLABLE
-
-internal fun <T> Continuation<T>.resumeMode(value: T, mode: Int) {
-    when (mode) {
-        MODE_ATOMIC_DEFAULT -> resume(value)
-        MODE_CANCELLABLE -> resumeCancellable(value)
-        MODE_DIRECT -> resumeDirect(value)
-        MODE_UNDISPATCHED -> (this as DispatchedContinuation).resumeUndispatched(value)
-        else -> error("Invalid mode $mode")
-    }
-}
-
-internal fun <T> Continuation<T>.resumeWithExceptionMode(exception: Throwable, mode: Int) {
-    when (mode) {
-        MODE_ATOMIC_DEFAULT -> resumeWithException(exception)
-        MODE_CANCELLABLE -> resumeCancellableWithException(exception)
-        MODE_DIRECT -> resumeDirectWithException(exception)
-        MODE_UNDISPATCHED -> (this as DispatchedContinuation).resumeUndispatchedWithException(exception)
-        else -> error("Invalid mode $mode")
-    }
-}
-
-internal fun <T> Continuation<T>.resumeUninterceptedMode(value: T, mode: Int) {
-    when (mode) {
-        MODE_ATOMIC_DEFAULT -> intercepted().resume(value)
-        MODE_CANCELLABLE -> intercepted().resumeCancellable(value)
-        MODE_DIRECT -> resume(value)
-        MODE_UNDISPATCHED -> withCoroutineContext(context, null) { resume(value) }
-        else -> error("Invalid mode $mode")
-    }
-}
-
-internal fun <T> Continuation<T>.resumeUninterceptedWithExceptionMode(exception: Throwable, mode: Int) {
-    when (mode) {
-        MODE_ATOMIC_DEFAULT -> intercepted().resumeWithException(exception)
-        MODE_CANCELLABLE -> intercepted().resumeCancellableWithException(exception)
-        MODE_DIRECT -> resumeWithException(exception)
-        MODE_UNDISPATCHED -> withCoroutineContext(context, null) { resumeWithException(exception) }
-        else -> error("Invalid mode $mode")
-    }
-}

kotlinx-coroutines-core/common/src/Timeout.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
 
 package kotlinx.coroutines
@@ -80,26 +80,12 @@ private fun <U, T: U> setupTimeout(
 
 private open class TimeoutCoroutine<U, in T: U>(
     @JvmField val time: Long,
-    @JvmField val uCont: Continuation<U> // unintercepted continuation
-) : AbstractCoroutine<T>(uCont.context, active = true), Runnable, Continuation<T>, CoroutineStackFrame {
-    override val defaultResumeMode: Int get() = MODE_DIRECT
-    override val callerFrame: CoroutineStackFrame? get() = (uCont as? CoroutineStackFrame)
-    override fun getStackTraceElement(): StackTraceElement? = null
-    override val isScopedCoroutine: Boolean get() = true
-
-    @Suppress("LeakingThis", "Deprecation")
+    uCont: Continuation<U> // unintercepted continuation
+) : ScopeCoroutine<T>(uCont.context, uCont), Runnable {
     override fun run() {
         cancelCoroutine(TimeoutCancellationException(time, this))
     }
 
-    @Suppress("UNCHECKED_CAST")
-    override fun afterCompletionInternal(state: Any?, mode: Int) {
-        if (state is CompletedExceptionally)
-            uCont.resumeUninterceptedWithExceptionMode(state.cause, mode)
-        else
-            uCont.resumeUninterceptedMode(state as T, mode)
-    }
-
     override fun nameString(): String =
         "${super.nameString()}(timeMillis=$time)"
 }

kotlinx-coroutines-core/common/src/internal/DispatchedContinuation.kt

@@ -175,32 +175,16 @@ internal class DispatchedContinuation<in T>(
     }
 
     @Suppress("NOTHING_TO_INLINE") // we need it inline to save us an entry on the stack
-    inline fun resumeCancellable(value: T) {
-        if (dispatcher.isDispatchNeeded(context)) {
-            _state = value
-            resumeMode = MODE_CANCELLABLE
-            dispatcher.dispatch(context, this)
-        } else {
-            executeUnconfined(value, MODE_CANCELLABLE) {
-                if (!resumeCancelled()) {
-                    resumeUndispatched(value)
-                }
-            }
-        }
-    }
-
-    @Suppress("NOTHING_TO_INLINE") // we need it inline to save us an entry on the stack
-    inline fun resumeCancellableWithException(exception: Throwable) {
-        val context = continuation.context
-        val state = CompletedExceptionally(exception)
+    inline fun resumeCancellableWith(result: Result<T>) {
+        val state = result.toState()
         if (dispatcher.isDispatchNeeded(context)) {
-            _state = CompletedExceptionally(exception)
+            _state = state
             resumeMode = MODE_CANCELLABLE
             dispatcher.dispatch(context, this)
         } else {
             executeUnconfined(state, MODE_CANCELLABLE) {
                 if (!resumeCancelled()) {
-                    resumeUndispatchedWithException(exception)
+                    resumeUndispatchedWith(result)
                 }
             }
         }
@@ -218,16 +202,9 @@ internal class DispatchedContinuation<in T>(
     }
 
     @Suppress("NOTHING_TO_INLINE") // we need it inline to save us an entry on the stack
-    inline fun resumeUndispatched(value: T) {
-        withCoroutineContext(context, countOrElement) {
-            continuation.resume(value)
-        }
-    }
-
-    @Suppress("NOTHING_TO_INLINE") // we need it inline to save us an entry on the stack
-    inline fun resumeUndispatchedWithException(exception: Throwable) {
+    inline fun resumeUndispatchedWith(result: Result<T>) {
         withCoroutineContext(context, countOrElement) {
-            continuation.resumeWithStackTrace(exception)
+            continuation.resumeWith(result)
         }
     }
 
@@ -243,6 +220,12 @@ internal class DispatchedContinuation<in T>(
         "DispatchedContinuation[$dispatcher, ${continuation.toDebugString()}]"
 }
 
+@Suppress("NOTHING_TO_INLINE") // we need it inline to save us an entry on the stack
+internal inline fun <T> Continuation<T>.resumeCancellableWith(result: Result<T>) = when (this) {
+    is DispatchedContinuation -> resumeCancellableWith(result)
+    else -> resumeWith(result)
+}
+
 internal fun DispatchedContinuation<Unit>.yieldUndispatched(): Boolean =
     executeUnconfined(Unit, MODE_CANCELLABLE, doYield = true) {
         run()

kotlinx-coroutines-core/common/src/internal/DispatchedTask.kt

@@ -105,21 +105,29 @@ internal fun <T> DispatchedTask<T>.dispatch(mode: Int = MODE_CANCELLABLE) {
     }
 }
 
+@Suppress("UNCHECKED_CAST")
 internal fun <T> DispatchedTask<T>.resume(delegate: Continuation<T>, useMode: Int) {
     // slow-path - use delegate
     val state = takeState()
-    val exception = getExceptionalResult(state)
-    if (exception != null) {
+    val exception = getExceptionalResult(state)?.let {
         /*
          * Recover stacktrace for non-dispatched tasks.
          * We usually do not recover stacktrace in a `resume` as all resumes go through `DispatchedTask.run`
          * and we recover stacktraces there, but this is not the case for a `suspend fun main()` that knows nothing about
          * kotlinx.coroutines and DispatchedTask
          */
-        val recovered = if (delegate is DispatchedTask<*>) exception else recoverStackTrace(exception, delegate)
-        delegate.resumeWithExceptionMode(recovered, useMode)
-    } else {
-        delegate.resumeMode(getSuccessfulResult(state), useMode)
+        if (delegate is DispatchedTask<*>) it else recoverStackTrace(it, delegate)
+    }
+    val result = if (exception != null)
+        Result.failure(exception)
+    else
+        Result.success(state as T)
+    when (useMode) {
+        MODE_ATOMIC_DEFAULT -> delegate.resumeWith(result)
+        MODE_CANCELLABLE -> delegate.resumeCancellableWith(result)
+        MODE_DIRECT -> ((delegate as? DispatchedContinuation)?.continuation ?: delegate).resumeWith(result)
+        MODE_UNDISPATCHED -> (delegate as DispatchedContinuation).resumeUndispatchedWith(result)
+        else -> error("Invalid mode $useMode")
     }
 }
 
@@ -158,27 +166,6 @@ internal inline fun DispatchedTask<*>.runUnconfinedEventLoop(
     }
 }
 
-
-internal fun <T> Continuation<T>.resumeCancellable(value: T) = when (this) {
-    is DispatchedContinuation -> resumeCancellable(value)
-    else -> resume(value)
-}
-
-internal fun <T> Continuation<T>.resumeCancellableWithException(exception: Throwable) = when (this) {
-    is DispatchedContinuation -> resumeCancellableWithException(exception)
-    else -> resumeWithStackTrace(exception)
-}
-
-internal fun <T> Continuation<T>.resumeDirect(value: T) = when (this) {
-    is DispatchedContinuation -> continuation.resume(value)
-    else -> resume(value)
-}
-
-internal fun <T> Continuation<T>.resumeDirectWithException(exception: Throwable) = when (this) {
-    is DispatchedContinuation -> continuation.resumeWithStackTrace(exception)
-    else -> resumeWithStackTrace(exception)
-}
-
 @Suppress("NOTHING_TO_INLINE")
 internal inline fun Continuation<*>.resumeWithStackTrace(exception: Throwable) {
     resumeWith(Result.failure(recoverStackTrace(exception, this)))

kotlinx-coroutines-core/common/src/internal/Scopes.kt

@@ -6,6 +6,7 @@ package kotlinx.coroutines.internal
 
 import kotlinx.coroutines.*
 import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
 import kotlin.jvm.*
 
 /**
@@ -25,10 +26,16 @@ internal open class ScopeCoroutine<in T>(
 
     @Suppress("UNCHECKED_CAST")
     override fun afterCompletionInternal(state: Any?, mode: Int) {
-        if (state is CompletedExceptionally) {
-            uCont.resumeUninterceptedWithExceptionMode(recoverStackTrace(state.cause, uCont), mode)
-        } else {
-            uCont.resumeUninterceptedMode(state as T, mode)
+        val result = if (state is CompletedExceptionally)
+            Result.failure(recoverStackTrace(state.cause, uCont))
+        else
+            Result.success(state as T)
+        when (mode) {
+            MODE_ATOMIC_DEFAULT -> uCont.intercepted().resumeWith(result)
+            MODE_CANCELLABLE -> uCont.intercepted().resumeCancellableWith(result)
+            MODE_DIRECT -> uCont.resumeWith(result)
+            MODE_UNDISPATCHED -> withCoroutineContext(uCont.context, null) { uCont.resumeWith(result) }
+            else -> error("Invalid mode $mode")
         }
     }
 }

kotlinx-coroutines-core/common/src/intrinsics/Cancellable.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
 
 package kotlinx.coroutines.intrinsics
@@ -14,7 +14,7 @@ import kotlin.coroutines.intrinsics.*
  */
 @InternalCoroutinesApi
 public fun <T> (suspend () -> T).startCoroutineCancellable(completion: Continuation<T>) = runSafely(completion) {
-    createCoroutineUnintercepted(completion).intercepted().resumeCancellable(Unit)
+    createCoroutineUnintercepted(completion).intercepted().resumeCancellableWith(Result.success(Unit))
 }
 
 /**
@@ -23,7 +23,7 @@ public fun <T> (suspend () -> T).startCoroutineCancellable(completion: Continuat
  */
 internal fun <R, T> (suspend (R) -> T).startCoroutineCancellable(receiver: R, completion: Continuation<T>) =
     runSafely(completion) {
-        createCoroutineUnintercepted(receiver, completion).intercepted().resumeCancellable(Unit)
+        createCoroutineUnintercepted(receiver, completion).intercepted().resumeCancellableWith(Result.success(Unit))
     }
 
 /**
@@ -32,7 +32,7 @@ internal fun <R, T> (suspend (R) -> T).startCoroutineCancellable(receiver: R, co
  */
 internal fun Continuation<Unit>.startCoroutineCancellable(fatalCompletion: Continuation<*>) =
     runSafely(fatalCompletion) {
-        intercepted().resumeCancellable(Unit)
+        intercepted().resumeCancellableWith(Result.success(Unit))
     }
 
 /**

kotlinx-coroutines-core/common/src/selects/Select.kt

@@ -285,7 +285,7 @@ internal class SelectBuilderImpl<in R>(
     // Resumes in MODE_CANCELLABLE
     override fun resumeSelectCancellableWithException(exception: Throwable) {
         doResume({ CompletedExceptionally(exception) }) {
-            uCont.intercepted().resumeCancellableWithException(exception)
+            uCont.intercepted().resumeCancellableWith(Result.failure(exception))
         }
     }