Refactor WorkflowNode into AbstractWorkflowNode <- StatefulWorkflowNode.

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

@@ -0,0 +1,143 @@
+package com.squareup.workflow1.internal
+
+import com.squareup.workflow1.ActionApplied
+import com.squareup.workflow1.ActionProcessingResult
+import com.squareup.workflow1.NoopWorkflowInterceptor
+import com.squareup.workflow1.RuntimeConfig
+import com.squareup.workflow1.RuntimeConfigOptions
+import com.squareup.workflow1.TreeSnapshot
+import com.squareup.workflow1.Workflow
+import com.squareup.workflow1.WorkflowIdentifier
+import com.squareup.workflow1.WorkflowInterceptor
+import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
+import com.squareup.workflow1.WorkflowTracer
+import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.selects.SelectBuilder
+import kotlin.coroutines.CoroutineContext
+
+internal fun <PropsT, OutputT, RenderingT> createWorkflowNode(
+  id: WorkflowNodeId,
+  workflow: Workflow<PropsT, OutputT, RenderingT>,
+  initialProps: PropsT,
+  snapshot: TreeSnapshot?,
+  baseContext: CoroutineContext,
+  // Providing default value so we don't need to specify in test.
+  runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG,
+  workflowTracer: WorkflowTracer? = null,
+  emitAppliedActionToParent: (ActionApplied<OutputT>) -> ActionProcessingResult = { it },
+  parent: WorkflowSession? = null,
+  interceptor: WorkflowInterceptor = NoopWorkflowInterceptor,
+  idCounter: IdCounter? = null
+): AbstractWorkflowNode<PropsT, OutputT, RenderingT> = when (workflow) {
+  // is ComposeWorkflow<*, *, *> -> ComposeWorkflowNode(
+  //   id = id,
+  //   workflow = workflow,
+  //   initialProps = initialProps,
+  //   snapshot = snapshot,
+  //   baseContext = baseContext,
+  //   runtimeConfig = runtimeConfig,
+  //   workflowTracer = workflowTracer,
+  //   emitAppliedActionToParent = emitAppliedActionToParent,
+  //   parent = parent,
+  //   interceptor = interceptor,
+  //   idCounter = idCounter,
+  // )
+
+  else -> StatefulWorkflowNode(
+    id = id,
+    workflow = workflow.asStatefulWorkflow(),
+    initialProps = initialProps,
+    snapshot = snapshot,
+    baseContext = baseContext,
+    runtimeConfig = runtimeConfig,
+    workflowTracer = workflowTracer,
+    emitAppliedActionToParent = emitAppliedActionToParent,
+    parent = parent,
+    interceptor = interceptor,
+    idCounter = idCounter,
+  )
+}
+
+internal abstract class AbstractWorkflowNode<PropsT, OutputT, RenderingT>(
+  val id: WorkflowNodeId,
+  final override val parent: WorkflowSession?,
+  final override val workflowTracer: WorkflowTracer?,
+  final override val runtimeConfig: RuntimeConfig,
+  protected val interceptor: WorkflowInterceptor,
+  protected val emitAppliedActionToParent: (ActionApplied<OutputT>) -> ActionProcessingResult,
+  baseContext: CoroutineContext,
+  idCounter: IdCounter?,
+) : WorkflowSession,
+  CoroutineScope {
+
+  /**
+   * Context that has a job that will live as long as this node.
+   * Also adds a debug name to this coroutine based on its ID.
+   */
+  final override val coroutineContext = baseContext +
+    Job(baseContext[Job]) +
+    CoroutineName(id.toString())
+
+  // WorkflowSession properties
+  final override val identifier: WorkflowIdentifier get() = id.identifier
+  final override val renderKey: String get() = id.name
+  final override val sessionId: Long = idCounter.createId()
+
+  final override fun toString(): String {
+    val parentDescription = parent?.let { "WorkflowInstance(…)" }
+    return "WorkflowInstance(" +
+      "identifier=$identifier, " +
+      "renderKey=$renderKey, " +
+      "instanceId=$sessionId, " +
+      "parent=$parentDescription" +
+      ")"
+  }
+
+  /**
+   * Walk the tree of workflows, rendering each one and using
+   * [RenderContext][com.squareup.workflow1.BaseRenderContext] to give its children a chance to
+   * render themselves and aggregate those child renderings.
+   *
+   * @param workflow The "template" workflow instance used in the current render pass. This isn't
+   * necessarily the same _instance_ every call, but will be the same _type_.
+   */
+  abstract fun render(
+    workflow: Workflow<PropsT, OutputT, RenderingT>,
+    input: PropsT
+  ): RenderingT
+
+  /**
+   * Walk the tree of state machines again, this time gathering snapshots and aggregating them
+   * automatically.
+   */
+  abstract fun snapshot(): TreeSnapshot
+
+  /**
+   * Gets the next [result][ActionProcessingResult] from the state machine. This will be an
+   * [OutputT] or null.
+   *
+   * Walk the tree of state machines, asking each one to wait for its next event. If something happen
+   * that results in an output, that output is returned. Null means something happened that requires
+   * a re-render, e.g. my state changed or a child state changed.
+   *
+   * It is an error to call this method after calling [cancel].
+   *
+   * @return [Boolean] whether or not the queues were empty for this node and its children at the
+   *    time of suspending.
+   */
+  abstract fun onNextAction(selector: SelectBuilder<ActionProcessingResult>): Boolean
+
+  /**
+   * Cancels this state machine host, and any coroutines started as children of it.
+   *
+   * This must be called when the caller will no longer call [onNextAction]. It is an error to call [onNextAction]
+   * after calling this method.
+   */
+  open fun cancel(cause: CancellationException? = null) {
+    coroutineContext.cancel(cause)
+  }
+}

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

@@ -0,0 +1,63 @@
+package com.squareup.workflow1.internal
+
+import com.squareup.workflow1.ActionApplied
+import com.squareup.workflow1.ActionProcessingResult
+import com.squareup.workflow1.NoopWorkflowInterceptor
+import com.squareup.workflow1.RuntimeConfig
+import com.squareup.workflow1.RuntimeConfigOptions
+import com.squareup.workflow1.TreeSnapshot
+import com.squareup.workflow1.Workflow
+import com.squareup.workflow1.WorkflowInterceptor
+import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
+import com.squareup.workflow1.WorkflowTracer
+import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.selects.SelectBuilder
+import kotlin.coroutines.CoroutineContext
+
+internal class ComposeWorkflowNode<PropsT, OutputT, RenderingT>(
+  id: WorkflowNodeId,
+  // workflow: ComposeWorkflow<PropsT, OutputT, RenderingT>,
+  initialProps: PropsT,
+  snapshot: TreeSnapshot?,
+  baseContext: CoroutineContext,
+  // Providing default value so we don't need to specify in test.
+  runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG,
+  workflowTracer: WorkflowTracer? = null,
+  emitAppliedActionToParent: (ActionApplied<OutputT>) -> ActionProcessingResult = { it },
+  parent: WorkflowSession? = null,
+  interceptor: WorkflowInterceptor = NoopWorkflowInterceptor,
+  idCounter: IdCounter? = null
+) : AbstractWorkflowNode<PropsT, OutputT, RenderingT>(
+  id = id,
+  runtimeConfig = runtimeConfig,
+  workflowTracer = workflowTracer,
+  parent = parent,
+  baseContext = baseContext,
+  idCounter = idCounter,
+  interceptor = interceptor,
+  emitAppliedActionToParent = emitAppliedActionToParent,
+) {
+
+  init {
+    interceptor.onSessionStarted(workflowScope = this, session = this)
+  }
+
+  override fun render(
+    workflow: Workflow<PropsT, OutputT, RenderingT>,
+    input: PropsT
+  ): RenderingT {
+    TODO("Not yet implemented")
+  }
+
+  override fun snapshot(): TreeSnapshot {
+    TODO("Not yet implemented")
+  }
+
+  override fun onNextAction(selector: SelectBuilder<ActionProcessingResult>): Boolean {
+    TODO("Not yet implemented")
+  }
+
+  override fun cancel(cause: CancellationException?) {
+    TODO("Not yet implemented")
+  }
+}

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/WorkflowNode.kt → workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/StatefulWorkflowNode.kt

@@ -14,7 +14,6 @@ import com.squareup.workflow1.Workflow
 import com.squareup.workflow1.WorkflowAction
 import com.squareup.workflow1.WorkflowExperimentalApi
 import com.squareup.workflow1.WorkflowExperimentalRuntime
-import com.squareup.workflow1.WorkflowIdentifier
 import com.squareup.workflow1.WorkflowInterceptor
 import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
 import com.squareup.workflow1.WorkflowTracer
@@ -50,32 +49,32 @@ import kotlin.reflect.KType
  * structured concurrency).
  */
 @OptIn(WorkflowExperimentalApi::class, WorkflowExperimentalRuntime::class)
-internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
-  val id: WorkflowNodeId,
+internal class StatefulWorkflowNode<PropsT, StateT, OutputT, RenderingT>(
+  id: WorkflowNodeId,
   workflow: StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>,
   initialProps: PropsT,
   snapshot: TreeSnapshot?,
   baseContext: CoroutineContext,
   // Providing default value so we don't need to specify in test.
-  override val runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG,
-  override val workflowTracer: WorkflowTracer? = null,
-  private val emitAppliedActionToParent: (ActionApplied<OutputT>) -> ActionProcessingResult =
-    { it },
-  override val parent: WorkflowSession? = null,
-  private val interceptor: WorkflowInterceptor = NoopWorkflowInterceptor,
+  runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG,
+  workflowTracer: WorkflowTracer? = null,
+  emitAppliedActionToParent: (ActionApplied<OutputT>) -> ActionProcessingResult = { it },
+  parent: WorkflowSession? = null,
+  interceptor: WorkflowInterceptor = NoopWorkflowInterceptor,
   idCounter: IdCounter? = null
-) : CoroutineScope, SideEffectRunner, RememberStore, WorkflowSession {
+) : AbstractWorkflowNode<PropsT, OutputT, RenderingT>(
+  id = id,
+  runtimeConfig = runtimeConfig,
+  workflowTracer = workflowTracer,
+  parent = parent,
+  baseContext = baseContext,
+  idCounter = idCounter,
+  interceptor = interceptor,
+  emitAppliedActionToParent = emitAppliedActionToParent,
+),
+  SideEffectRunner,
+  RememberStore {
 
-  /**
-   * Context that has a job that will live as long as this node.
-   * Also adds a debug name to this coroutine based on its ID.
-   */
-  override val coroutineContext = baseContext + Job(baseContext[Job]) + CoroutineName(id.toString())
-
-  // WorkflowInstance properties
-  override val identifier: WorkflowIdentifier get() = id.identifier
-  override val renderKey: String get() = id.name
-  override val sessionId: Long = idCounter.createId()
   private var cachedWorkflowInstance: StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>
   private var interceptedWorkflowInstance: StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>
 
@@ -116,45 +115,37 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     state = interceptedWorkflowInstance.initialState(initialProps, snapshot?.workflowSnapshot, this)
   }
 
-  override fun toString(): String {
-    val parentDescription = parent?.let { "WorkflowInstance(…)" }
-    return "WorkflowInstance(" +
-      "identifier=$identifier, " +
-      "renderKey=$renderKey, " +
-      "instanceId=$sessionId, " +
-      "parent=$parentDescription" +
-      ")"
-  }
-
   /**
    * Walk the tree of workflows, rendering each one and using
    * [RenderContext][com.squareup.workflow1.BaseRenderContext] to give its children a chance to
    * render themselves and aggregate those child renderings.
    */
   @Suppress("UNCHECKED_CAST")
-  fun render(
-    workflow: StatefulWorkflow<PropsT, *, OutputT, RenderingT>,
+  override fun render(
+    workflow: Workflow<PropsT, OutputT, RenderingT>,
     input: PropsT
-  ): RenderingT =
-    renderWithStateType(workflow as StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>, input)
+  ): RenderingT = renderWithStateType(
+    workflow = workflow.asStatefulWorkflow() as StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>,
+    props = input
+  )
 
   /**
    * Walk the tree of state machines again, this time gathering snapshots and aggregating them
    * automatically.
    */
-  fun snapshot(workflow: StatefulWorkflow<*, *, *, *>): TreeSnapshot {
-    @Suppress("UNCHECKED_CAST")
-    val typedWorkflow = workflow as StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>
-    maybeUpdateCachedWorkflowInstance(typedWorkflow)
-    return interceptor.onSnapshotStateWithChildren({
-      val childSnapshots = subtreeManager.createChildSnapshots()
-      val rootSnapshot = interceptedWorkflowInstance.snapshotState(state)
-      TreeSnapshot(
-        workflowSnapshot = rootSnapshot,
-        // Create the snapshots eagerly since subtreeManager is mutable.
-        childTreeSnapshots = { childSnapshots }
-      )
-    }, this)
+  override fun snapshot(): TreeSnapshot {
+    return interceptor.onSnapshotStateWithChildren(
+      proceed = {
+        val childSnapshots = subtreeManager.createChildSnapshots()
+        val rootSnapshot = interceptedWorkflowInstance.snapshotState(state)
+        TreeSnapshot(
+          workflowSnapshot = rootSnapshot,
+          // Create the snapshots eagerly since subtreeManager is mutable.
+          childTreeSnapshots = { childSnapshots }
+        )
+      },
+      session = this
+    )
   }
 
   override fun runningSideEffect(
@@ -212,7 +203,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
    *    time of suspending.
    */
   @OptIn(ExperimentalCoroutinesApi::class, DelicateCoroutinesApi::class)
-  fun onNextAction(selector: SelectBuilder<ActionProcessingResult>): Boolean {
+  override fun onNextAction(selector: SelectBuilder<ActionProcessingResult>): Boolean {
     // Listen for any child workflow updates.
     var empty = subtreeManager.onNextChildAction(selector)
 
@@ -230,11 +221,11 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
   /**
    * Cancels this state machine host, and any coroutines started as children of it.
    *
-   * This must be called when the caller will no longer call [onNextAction]. It is an error to call [onNextAction]
-   * after calling this method.
+   * This must be called when the caller will no longer call [onNextAction]. It is an error to call
+   * [onNextAction] after calling this method.
    */
-  fun cancel(cause: CancellationException? = null) {
-    coroutineContext.cancel(cause)
+  override fun cancel(cause: CancellationException?) {
+    super.cancel(cause)
     lastRendering = NullableInitBox()
   }
 
@@ -314,7 +305,6 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
    * Applies [action] to this workflow's [state] and then passes the resulting [ActionApplied]
    * via [emitAppliedActionToParent] to the parent, with additional information as to whether or
    * not this action has changed the current node's state.
-   *
    */
   private fun applyAction(
     action: WorkflowAction<PropsT, StateT, OutputT>,