~merge SafeCollector and ContinuationImpl, properly release intercepted continuation

Author: qwwdfsad

kotlinx-coroutines-core/common/src/flow/Builders.kt

@@ -55,7 +55,7 @@ private class SafeFlow<T>(private val block: suspend FlowCollector<T>.() -> Unit
         try {
             safeCollector.block()
         } finally {
-            safeCollector.release()
+            safeCollector.releaseIntercepted()
         }
     }
 }

kotlinx-coroutines-core/common/src/flow/Flow.kt

@@ -5,7 +5,7 @@
 package kotlinx.coroutines.flow
 
 import kotlinx.coroutines.*
-import kotlinx.coroutines.flow.internal.SafeCollector
+import kotlinx.coroutines.flow.internal.*
 import kotlin.coroutines.*
 
 /**
@@ -149,8 +149,8 @@ import kotlin.coroutines.*
  * it hard to reason about the code because an exception in the `collect { ... }` could be somehow "caught"
  * by an upstream flow, limiting the ability of local reasoning about the code.
  *
- * Flow machinery enforces exception transparency at runtime an throws [IllegalStateException] on any attempt to emit value,
- * if an exception has been thrown on previous attemp.
+ * Flow machinery enforces exception transparency at runtime and throws [IllegalStateException] on any attempt to emit a value,
+ * if an exception has been thrown on previous attempt.
  *
  * ### Reactive streams
  *
@@ -199,7 +199,12 @@ public abstract class AbstractFlow<T> : Flow<T> {
 
     @InternalCoroutinesApi
     public final override suspend fun collect(collector: FlowCollector<T>) {
-        collectSafely(SafeCollector(collector, coroutineContext))
+        val safeCollector = SafeCollector(collector, coroutineContext)
+        try {
+            collectSafely(safeCollector)
+        } finally {
+            safeCollector.releaseIntercepted()
+        }
     }
 
     /**

kotlinx-coroutines-core/common/src/flow/internal/SafeCollector.common.kt

@@ -17,7 +17,7 @@ internal expect class SafeCollector<T>(
     internal val collector: FlowCollector<T>
     internal val collectContext: CoroutineContext
     internal val collectContextSize: Int
-    public fun release()
+    public fun releaseIntercepted()
 }
 
 @JvmName("checkContext") // For prettier stack traces

kotlinx-coroutines-core/common/src/flow/operators/Emitters.kt

@@ -71,7 +71,12 @@ internal inline fun <T, R> Flow<T>.unsafeTransform(
 public fun <T> Flow<T>.onStart(
     action: suspend FlowCollector<T>.() -> Unit
 ): Flow<T> = unsafeFlow { // Note: unsafe flow is used here, but safe collector is used to invoke start action
-    SafeCollector<T>(this, coroutineContext).action()
+    val safeCollector = SafeCollector<T>(this, coroutineContext)
+    try {
+        safeCollector.action()
+    } finally {
+        safeCollector.releaseIntercepted()
+    }
     collect(this) // directly delegate
 }
 
@@ -141,7 +146,12 @@ public fun <T> Flow<T>.onCompletion(
         throw e
     }
     // Exception from the upstream or normal completion
-    SafeCollector(this, coroutineContext).invokeSafely(action, exception)
+    val safeCollector = SafeCollector(this, coroutineContext)
+    try {
+        safeCollector.invokeSafely(action, exception)
+    } finally {
+        safeCollector.releaseIntercepted()
+    }
     exception?.let { throw it }
 }
 

kotlinx-coroutines-core/js/src/flow/internal/SafeCollector.kt

@@ -25,6 +25,6 @@ internal actual class SafeCollector<T> actual constructor(
         collector.emit(value)
     }
 
-    public actual fun release() {
+    public actual fun releaseIntercepted() {
     }
 }

kotlinx-coroutines-core/jvm/src/flow/internal/SafeCollector.kt

@@ -10,41 +10,36 @@ import kotlin.coroutines.intrinsics.*
 import kotlin.coroutines.jvm.internal.*
 
 @Suppress("UNCHECKED_CAST")
+private val emitFun =
+    FlowCollector<Any?>::emit as Function3<FlowCollector<Any?>, Any?, Continuation<Unit>, Any?>
+/*
+ * Implementor of ContinuationImpl (that will be preserved as ABI nearly forever)
+ * in order to properly control 'intercepted()' lifecycle.
+ */
+@Suppress("CANNOT_OVERRIDE_INVISIBLE_MEMBER", "INVISIBLE_MEMBER", "INVISIBLE_REFERENCE", "UNCHECKED_CAST")
 internal actual class SafeCollector<T> actual constructor(
     @JvmField internal actual val collector: FlowCollector<T>,
     @JvmField internal actual val collectContext: CoroutineContext
-) : FlowCollector<T> {
+) : FlowCollector<T>, ContinuationImpl(NoOpContinuation, EmptyCoroutineContext) {
 
     @JvmField // Note, it is non-capturing lambda, so no extra allocation during init of SafeCollector
     internal actual val collectContextSize = collectContext.fold(0) { count, _ -> count + 1 }
     private var lastEmissionContext: CoroutineContext? = null
-    private var intercepted: InterceptedContinuationOwner = InterceptedContinuationOwner()
-    private val emitFun = run {
-        collector::emit as Function2<T, Continuation<Unit>, Any?>
-    }
+    private var completion: Continuation<Unit>? = null
 
-    /*
-     * Implementor of ContinuationImpl (that will be preserved by ABI implementation nearly forever)
-     * in order to properly control 'intercepted()' lifecycle.
-     */
-    @Suppress("CANNOT_OVERRIDE_INVISIBLE_MEMBER", "INVISIBLE_MEMBER", "INVISIBLE_REFERENCE")
-    private inner class InterceptedContinuationOwner : ContinuationImpl(NoOpContinuation, EmptyCoroutineContext) {
-        @JvmField
-        public var realCompletion: Continuation<Unit>? = null
-
-        override val context: CoroutineContext
-            get() = realCompletion?.context ?: EmptyCoroutineContext
-
-        override fun invokeSuspend(result: Result<Any?>): Any? {
-            result.onFailure { lastEmissionContext = DownstreamExceptionElement(it) }
-            realCompletion?.resumeWith(result as Result<Unit>)
-            return COROUTINE_SUSPENDED
-        }
+    // ContinuationImpl
+    override val context: CoroutineContext
+        get() = completion?.context ?: EmptyCoroutineContext
 
-        // Escalate visibility to manually release intercepted continuation
-        public override fun releaseIntercepted() {
-            super.releaseIntercepted()
-        }
+    override fun invokeSuspend(result: Result<Any?>): Any? {
+        result.onFailure { lastEmissionContext = DownstreamExceptionElement(it) }
+        completion?.resumeWith(result as Result<Unit>)
+        return COROUTINE_SUSPENDED
+    }
+
+    // Escalate visibility to manually release intercepted continuation
+    public actual override fun releaseIntercepted() {
+        super.releaseIntercepted()
     }
 
     /**
@@ -72,8 +67,8 @@ internal actual class SafeCollector<T> actual constructor(
         if (previousContext !== currentContext) {
             checkContext(currentContext, previousContext, value)
         }
-        intercepted.realCompletion = uCont
-        return emitFun(value, intercepted as Continuation<Unit>)
+        completion = uCont
+        return emitFun(collector as FlowCollector<Any?>, value, this as Continuation<Unit>)
     }
 
     private fun checkContext(
@@ -88,10 +83,6 @@ internal actual class SafeCollector<T> actual constructor(
         lastEmissionContext = currentContext
     }
 
-    public actual fun release() {
-        intercepted.releaseIntercepted()
-    }
-
     private fun exceptionTransparencyViolated(exception: DownstreamExceptionElement, value: Any?) {
         /*
          * Exception transparency ensures that if a `collect` block or any intermediate operator

kotlinx-coroutines-core/native/src/flow/internal/SafeCollector.kt

@@ -25,6 +25,6 @@ internal actual class SafeCollector<T> actual constructor(
         collector.emit(value)
     }
 
-    public actual fun release() {
+    public actual fun releaseIntercepted() {
     }
 }