Refactor RecomposerDriver a bit.

Author: zach-klippenstein

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/compose/ComposeWorkflowNodeAdapter.kt

@@ -72,10 +72,7 @@ internal class ComposeWorkflowNodeAdapter<PropsT, OutputT, RenderingT>(
 
   private val recomposer: Recomposer = Recomposer(coroutineContext)
 
-  // TODO we could create a PreemptingDispatcher and stash it in the coroutine context (not as a
-  //  dispatcher, but inside a private holder) so we can reuse it for all compose runtimes. But it
-  //  wouldn't really save that much and might not be worth it.
-  private val recomposerDriver = RecomposerDriver(recomposer = recomposer, scope = this)
+  private val recomposerDriver = RecomposerDriver(recomposer)
   private val composition: Composition = Composition(UnitApplier, recomposer)
 
   private var frameTimeCounter = 0L

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/compose/RecomposerDriver.kt

@@ -6,7 +6,6 @@ import com.squareup.workflow1.internal.compose.coroutines.PreemptingDispatcher
 import com.squareup.workflow1.internal.compose.coroutines.requireSend
 import kotlinx.coroutines.CancellableContinuation
 import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.channels.Channel
@@ -39,43 +38,48 @@ import kotlin.coroutines.Continuation
  * recompose loop, and by advancing it any time the recomposer reports pending work but hasn't
  * requested a frame yet.
  */
-internal class RecomposerDriver(
-  private val recomposer: Recomposer,
-  private val dispatcher: PreemptingDispatcher,
-) {
-  @OptIn(ExperimentalStdlibApi::class)
-  constructor(
-    recomposer: Recomposer,
-    scope: CoroutineScope
-  ) : this(
-    recomposer = recomposer,
-    dispatcher = PreemptingDispatcher(
-      scope.coroutineContext[CoroutineDispatcher] ?: Dispatchers.Default
+interface RecomposerDriver {
+  /**
+   * Returns true if the recomposer is ready to recompose. When true, the next call to
+   * [tryPerformRecompose] will succeed.
+   *
+   * Use [onAwaitFrameAvailable] to wait for this to be true.
+   */
+  val needsRecompose: Boolean
+
+  suspend fun runRecomposeAndApplyChanges()
+
+  /**
+   * If the [Recomposer] is ready to recompose ([needsRecompose] is true), performs the
+   * recomposition with the given frame time and returns true. Returns false if there is no work to
+   * do.
+   */
+  fun tryPerformRecompose(frameTimeNanos: Long): Boolean
+
+  /**
+   * Registers with selector to resume when [needsRecompose] becomes true.
+   */
+  fun <R> onAwaitFrameAvailable(
+    selector: SelectBuilder<R>,
+    block: suspend () -> R
+  )
+}
+
+@OptIn(ExperimentalStdlibApi::class)
+internal fun RecomposerDriver(recomposer: Recomposer): RecomposerDriver =
+  RealRecomposerDriver(
+    recomposer,
+    PreemptingDispatcher(
+      recomposer.effectCoroutineContext[CoroutineDispatcher] ?: Dispatchers.Default
     )
   )
 
-  private val frameRequestChannel = Channel<FrameRequest<*>>(capacity = 1)
+private class RealRecomposerDriver(
+  private val recomposer: Recomposer,
+  private val dispatcher: PreemptingDispatcher,
+) : RecomposerDriver, MonotonicFrameClock {
 
-  private val frameClock = object : MonotonicFrameClock {
-    @OptIn(ExperimentalStdlibApi::class)
-    override suspend fun <R> withFrameNanos(onFrame: (frameTimeNanos: Long) -> R): R {
-      log("compose workflow withFrameNanos (dispatcher=${currentCoroutineContext()[CoroutineDispatcher]})")
-      log(RuntimeException().stackTraceToString())
-
-      return suspendCancellableCoroutine { continuation ->
-        val frameRequest = FrameRequest(
-          onFrame = onFrame,
-          continuation = continuation
-        )
-
-        // This will throw if a frame request is already enqueued. If currently in an action cascade
-        // (i.e. handling a received output), then it will be picked up in the imminent re-render.
-        // Otherwise, onNextAction will have registered a receiver for it that will trigger a render
-        // pass.
-        frameRequestChannel.requireSend(frameRequest)
-      }
-    }
-  }
+  private val frameRequestChannel = Channel<FrameRequest<*>>(capacity = 1)
 
   /**
    * Returns true if the recomposer is ready to recompose. When true, the next call to
@@ -84,7 +88,7 @@ internal class RecomposerDriver(
    * Use [onAwaitFrameAvailable] to wait for this to be true.
    */
   @OptIn(ExperimentalCoroutinesApi::class)
-  val needsRecompose: Boolean
+  override val needsRecompose: Boolean
     get() {
       val wasEmpty = frameRequestChannel.isEmpty
       if (wasEmpty && recomposer.hasPendingWork) {
@@ -99,11 +103,12 @@ internal class RecomposerDriver(
       }
     }
 
-  suspend fun runRecomposeAndApplyChanges() {
+  override suspend fun runRecomposeAndApplyChanges() {
     // Note: This context is _only_ used for the actual recompose loop. Everything inside the
     // composition (rememberCoroutineScope, LaunchedEffects, etc) will NOT see these, and will see
     // only whatever context was passed into the Recomposer's constructor (plus the stuff it adds
     // to that context itself, like the BroadcastFrameClock).
+    val frameClock = this
     withContext(dispatcher + frameClock) {
       recomposer.runRecomposeAndApplyChanges()
     }
@@ -114,7 +119,7 @@ internal class RecomposerDriver(
    * recomposition with the given frame time and returns true. Returns false if there is no work to
    * do.
    */
-  fun tryPerformRecompose(frameTimeNanos: Long): Boolean {
+  override fun tryPerformRecompose(frameTimeNanos: Long): Boolean {
     tryGetFrameRequest()?.let { frameRequest ->
       frameRequest.execute(frameTimeNanos)
       return true
@@ -125,7 +130,7 @@ internal class RecomposerDriver(
   /**
    * Registers with selector to resume when [needsRecompose] becomes true.
    */
-  fun <R> onAwaitFrameAvailable(
+  override fun <R> onAwaitFrameAvailable(
     selector: SelectBuilder<R>,
     block: suspend () -> R
   ) {
@@ -140,6 +145,25 @@ internal class RecomposerDriver(
     }
   }
 
+  @OptIn(ExperimentalStdlibApi::class)
+  override suspend fun <R> withFrameNanos(onFrame: (frameTimeNanos: Long) -> R): R {
+    log("compose workflow withFrameNanos (dispatcher=${currentCoroutineContext()[CoroutineDispatcher]})")
+    log(RuntimeException().stackTraceToString())
+
+    return suspendCancellableCoroutine { continuation ->
+      val frameRequest = FrameRequest(
+        onFrame = onFrame,
+        continuation = continuation
+      )
+
+      // This will throw if a frame request is already enqueued. If currently in an action cascade
+      // (i.e. handling a received output), then it will be picked up in the imminent re-render.
+      // Otherwise, onNextAction will have registered a receiver for it that will trigger a render
+      // pass.
+      frameRequestChannel.requireSend(frameRequest)
+    }
+  }
+
   /**
    * Returns a [Continuation] representing the next frame request if the recomposer has work to do,
    * otherwise returns null. This method is best-effort: It tries to poke the recomposer to request