Refactor RecomposerDriver -> SynchronizedMolecule.

SynchronizedMolecule fully encapsulates the Compose runtime into a clean, minimal interface with a single entry point, much like launchMolecule does. This should be much more testable, and we could potentially even pull it out into the Molecule library. It also cleans up the code inside ComposeWorkflowNodeAdapter even more.

Author: zach-klippenstein

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/compose/coroutines/Synchronization.kt renamed to workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/Channels.kt

@@ -1,11 +1,7 @@
-package com.squareup.workflow1.internal.compose.coroutines
+package com.squareup.workflow1.internal
 
 import kotlinx.coroutines.channels.SendChannel
 
-internal expect class Lock()
-
-internal expect inline fun <R> Lock.withLock(block: () -> R): R
-
 /**
  * Tries to send [element] to this channel and throws an [IllegalStateException] if the channel is
  * full or closed.

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

@@ -29,7 +29,7 @@ import com.squareup.workflow1.compose.WorkflowComposableRenderer
 import com.squareup.workflow1.identifier
 import com.squareup.workflow1.internal.IdCounter
 import com.squareup.workflow1.internal.WorkflowNodeId
-import com.squareup.workflow1.internal.compose.coroutines.requireSend
+import com.squareup.workflow1.internal.requireSend
 import com.squareup.workflow1.internal.createId
 import com.squareup.workflow1.workflowSessionToString
 import kotlinx.coroutines.CoroutineName

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

@@ -1,16 +1,9 @@
 package com.squareup.workflow1.internal.compose
 
-import androidx.compose.runtime.Composition
-import androidx.compose.runtime.Recomposer
-import androidx.compose.runtime.SideEffect
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableStateOf
-import androidx.compose.runtime.setValue
 import androidx.compose.runtime.snapshots.Snapshot
 import com.squareup.workflow1.ActionApplied
 import com.squareup.workflow1.ActionProcessingResult
 import com.squareup.workflow1.NoopWorkflowInterceptor
-import com.squareup.workflow1.NullableInitBox
 import com.squareup.workflow1.RuntimeConfig
 import com.squareup.workflow1.RuntimeConfigOptions
 import com.squareup.workflow1.TreeSnapshot
@@ -22,11 +15,10 @@ import com.squareup.workflow1.WorkflowTracer
 import com.squareup.workflow1.compose.ComposeWorkflow
 import com.squareup.workflow1.internal.AbstractWorkflowNode
 import com.squareup.workflow1.internal.IdCounter
-import com.squareup.workflow1.internal.UnitApplier
 import com.squareup.workflow1.internal.WorkflowNodeId
-import kotlinx.coroutines.CoroutineStart.ATOMIC
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.launch
+import com.squareup.workflow1.internal.compose.runtime.SynchronizedMolecule
+import com.squareup.workflow1.internal.compose.runtime.launchSynchronizedMolecule
+import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.selects.SelectBuilder
 import kotlin.coroutines.CoroutineContext
 
@@ -69,154 +61,104 @@ internal class ComposeWorkflowNodeAdapter<PropsT, OutputT, RenderingT>(
   emitAppliedActionToParent = emitAppliedActionToParent,
 ) {
 
-  private val recomposer: Recomposer = Recomposer(coroutineContext)
-
-  private val recomposerDriver = RecomposerDriver(recomposer)
-  private val composition: Composition = Composition(UnitApplier, recomposer)
-
-  private var cachedComposeWorkflow: ComposeWorkflow<PropsT, OutputT, RenderingT> by
-  mutableStateOf(workflow)
-  private var lastProps by mutableStateOf(initialProps)
-  private var lastRendering = NullableInitBox<RenderingT>()
+  private val recomposeChannel = Channel<Unit>(capacity = 1)
+  private val molecule: SynchronizedMolecule<RenderingT> = launchSynchronizedMolecule(
+    onNeedsRecomposition = { recomposeChannel.trySend(Unit) }
+  )
+
+  private val childNode = ComposeWorkflowChildNode<PropsT, OutputT, RenderingT>(
+    id = id,
+    initialProps = initialProps,
+    snapshot = snapshot,
+    baseContext = coroutineContext,
+    parent = parent,
+    workflowTracer = workflowTracer,
+    runtimeConfig = runtimeConfig,
+    interceptor = interceptor,
+    idCounter = idCounter,
+    emitAppliedActionToParent = { actionApplied ->
+      // Ensure any state updates performed by the output sender gets to invalidate any
+      // compositions that read them, so we can check needsRecompose below.
+      Snapshot.sendApplyNotifications()
+      log(
+        "adapter node sent apply notifications from action cascade (" +
+          "actionApplied=$actionApplied, needsRecompose=${molecule.needsRecomposition})"
+      )
+
+      // ComposeWorkflowChildNode can't tell if its own state changed since that information about
+      // specific composables/recompose scopes is only visible inside the compose runtime, so
+      // individual ComposeWorkflow nodes always report no state changes (unless they have a
+      // traditional child that reported a state change).
+      // However, we *can* check if any state changed that was read by anything in the
+      // composition, so when an action bubbles up to here, the top of the composition, we use
+      // that information to set the state changed flag if necessary.
+      val aggregateAction = if (molecule.needsRecomposition && !actionApplied.stateChanged) {
+        actionApplied.copy(stateChanged = true)
+      } else {
+        actionApplied
+      }
 
-  /**
-   * This is initialized to null so we don't render the workflow when initially calling
-   * [composition.setContent]. It is then set, and never nulled out again.
-   */
-  private var childNode: ComposeWorkflowChildNode<PropsT, OutputT, RenderingT>? = null
+      // Don't bubble up if no state changed and there was no output.
+      if (aggregateAction.stateChanged || aggregateAction.output != null) {
+        log("adapter node propagating action cascade up (aggregateAction=$aggregateAction)")
+        emitAppliedActionToParent(aggregateAction)
+      } else {
+        log("adapter node not propagating action cascade since nothing happened (aggregateAction=$aggregateAction)")
+        aggregateAction
+      }
+    }
+  )
 
   /**
-   * Function invoked when [onNextAction] receives a frame request from [withFrameNanos].
+   * Function invoked when [onNextAction] receives a recompose request.
    * This handles the case where some state read by the composition is changed but emitOutput is
    * not called.
    */
-  private val processFrameRequestFromChannel: () -> ActionProcessingResult = {
+  private val processRecompositionRequestFromChannel: (Unit) -> ActionProcessingResult = {
     // A pure frame request means compose state was updated that the composition read, but
     // emitOutput was not called, so we don't have any outputs to report.
     val applied = ActionApplied<OutputT>(
       output = null,
-      stateChanged = recomposerDriver.needsRecompose
+      // needsRecomposition should always be true now since the runtime explicitly requested
+      // recomposition, but check anyway.
+      stateChanged = molecule.needsRecomposition
     )
 
     // Propagate the action up the workflow tree.
     log("frame request received from channel, sending no output to parent: $applied")
     emitAppliedActionToParent(applied)
   }
 
-  init {
-    GlobalSnapshotManager.ensureStarted()
-
-    // By not calling setContent directly every time, we ensure that if neither the workflow
-    // instance nor input changed, we don't recompose.
-    // setContent will synchronously perform the first recomposition before returning, which is why
-    // we leave cachedComposeWorkflow null for now: we don't want its produceRendering to be called
-    // until we're actually doing a render pass.
-    // We also need to set the composition content before calling startComposition so it doesn't
-    // need to suspend to wait for it.
-    composition.setContent {
-      // childNode isn't snapshot state but that's fine, since when the recomposer is started it
-      // will always recompose, childNode will be non-null by then, and it will never change again.
-      val childNode = this.childNode
-      if (childNode != null) {
-        val rendering = childNode.produceRendering(
-          workflow = cachedComposeWorkflow,
-          props = lastProps
-        )
-
-        SideEffect {
-          this.lastRendering = NullableInitBox(rendering)
-        }
-      }
-    }
-
-    childNode = ComposeWorkflowChildNode(
-      id = id,
-      initialProps = initialProps,
-      snapshot = snapshot,
-      baseContext = coroutineContext,
-      parent = parent,
-      workflowTracer = workflowTracer,
-      runtimeConfig = runtimeConfig,
-      interceptor = interceptor,
-      idCounter = idCounter,
-      emitAppliedActionToParent = { actionApplied ->
-        // Ensure any state updates performed by the output sender gets to invalidate any
-        // compositions that read them, so we can check needsRecompose below.
-        Snapshot.sendApplyNotifications()
-        log(
-          "adapter node sent apply notifications from action cascade (" +
-            "actionApplied=$actionApplied, needsRecompose=${recomposerDriver.needsRecompose})"
-        )
-
-        // ComposeWorkflowChildNode can't tell if its own state changed since that information about
-        // specific composables/recompose scopes is only visible inside the compose runtime, so
-        // individual ComposeWorkflow nodes always report no state changes (unless they have a
-        // traditional child that reported a state change).
-        // However, we *can* check if any state changed that was read by anything in the
-        // composition, so when an action bubbles up to here, the top of the composition, we use
-        // that information to set the state changed flag if necessary.
-        val aggregateAction = if (recomposerDriver.needsRecompose && !actionApplied.stateChanged) {
-          actionApplied.copy(stateChanged = true)
-        } else {
-          actionApplied
-        }
-
-        // Don't bubble up if no state changed and there was no output.
-        if (aggregateAction.stateChanged || aggregateAction.output != null) {
-          log("adapter node propagating action cascade up (aggregateAction=$aggregateAction)")
-          emitAppliedActionToParent(aggregateAction)
-        } else {
-          log("adapter node not propagating action cascade since nothing happened (aggregateAction=$aggregateAction)")
-          aggregateAction
-        }
-      }
-    )
-  }
-
   override fun render(
     workflow: Workflow<PropsT, OutputT, RenderingT>,
     input: PropsT
   ): RenderingT {
-    this.cachedComposeWorkflow = workflow as ComposeWorkflow
-    this.lastProps = input
-
     // Ensure that recomposer has a chance to process any state changes from the action cascade that
     // triggered this render before we check for a frame.
-    log("render sending apply notifications again needsRecompose=${recomposerDriver.needsRecompose}")
+    log("render sending apply notifications again needsRecompose=${molecule.needsRecomposition}")
     // TODO Consider pulling this up into the workflow runtime loop, since we only need to run it
     //  once before the entire tree renders, not at every level. In fact, if this is only here to
     //  ensure cachedComposeWorkflow and lastProps are seen, that will only work if this
     //  ComposeWorkflow is not nested below another traditional and compose workflow, since anything
     //  rendering under the first CW will be in a snapshot.
     Snapshot.sendApplyNotifications()
-    log("sent apply notifications, needsRecompose=${recomposerDriver.needsRecompose}")
+    log("sent apply notifications, needsRecompose=${molecule.needsRecomposition}")
 
-    val initialRender = !lastRendering.isInitialized
-    if (initialRender) {
-      // Initial render kicks off the render loop. This should synchronously request a frame.
-      startComposition()
-    }
+    // If this re-render was not triggered by the channel handler, then clear it so we don't
+    // immediately trigger another redundant render pass after this.
+    recomposeChannel.tryReceive()
 
-    // Synchronously recompose any invalidated composables, if any, and update lastRendering.
-    // It is very likely that trySendFrame will fail: any time the workflow runtime is doing a
+    // It is very likely that this will be a noop: any time the workflow runtime is doing a
     // render pass and no state read by our composition changed, there shouldn't be a frame request.
-    // Hard-code unchanging frame time since there's no actual frame time and workflow code
-    // shouldn't rely on this value.
-    log("renderFrame")
-    val recomposed = recomposerDriver.tryPerformRecompose(frameTimeNanos = 0L)
-    if (recomposed) {
-      log("renderFrame finished executing frame")
-    } else {
-      log("no frame request at time of render!")
-      if (initialRender) {
-        error("Expected initial composition to synchronously request initial frame.")
-      }
+    return molecule.recomposeWithContent {
+      childNode.produceRendering(
+        workflow = workflow,
+        props = input
+      )
     }
-
-    return lastRendering.getOrThrow()
   }
 
-  override fun snapshot(): TreeSnapshot = childNode!!.snapshot()
+  override fun snapshot(): TreeSnapshot = childNode.snapshot()
 
   override fun onNextAction(selector: SelectBuilder<ActionProcessingResult>): Boolean {
     // We must register for child actions before frame requests, because selection is
@@ -225,26 +167,15 @@ internal class ComposeWorkflowNodeAdapter<PropsT, OutputT, RenderingT>(
     // the output handler will implicitly also handle frame requests. If a frame request happens at
     // the same time or the output handler enqueues a frame request, then the subsequent render pass
     // will dequeue the frame request itself before the next call to onNextAction.
-    var empty = childNode!!.onNextAction(selector)
+    var empty = childNode.onNextAction(selector)
 
     // If there's a frame request, then some state changed, which is equivalent to the traditional
     // case of a WorkflowAction being enqueued that just modifies state.
-    empty = empty && !recomposerDriver.needsRecompose
-    recomposerDriver.onAwaitFrameAvailable(selector, processFrameRequestFromChannel)
+    empty = empty && !molecule.needsRecomposition
+    with(selector) {
+      recomposeChannel.onReceive(processRecompositionRequestFromChannel)
+    }
 
     return empty
   }
-
-  @OptIn(ExperimentalCoroutinesApi::class)
-  private fun startComposition() {
-    // Launch as atomic to ensure the composition is always disposed, even if our job is cancelled
-    // before this coroutine has a chance to start running.
-    launch(start = ATOMIC) {
-      try {
-        recomposerDriver.runRecomposeAndApplyChanges()
-      } finally {
-        composition.dispose()
-      }
-    }
-  }
 }