Refactor WorkflowNode into AbstractWorkflowNode <- StatefulWorkflowNode.

This is to prepare for go/compose-based-workflows.

Author: zach-klippenstein

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

@@ -1,6 +1,7 @@
 package com.squareup.workflow1
 
 import com.squareup.workflow1.WorkflowInterceptor.RenderContextInterceptor
+import com.squareup.workflow1.WorkflowInterceptor.RuntimeUpdate
 import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
@@ -194,6 +195,8 @@ public interface WorkflowInterceptor {
   /**
    * Information about the session of a workflow in the runtime that a [WorkflowInterceptor] method
    * is intercepting.
+   *
+   * Implementations should override [toString] to call [WorkflowSession.workflowSessionToString].
    */
   public interface WorkflowSession {
     /** The [WorkflowIdentifier] that represents the type of this workflow. */
@@ -419,6 +422,16 @@ internal fun <P, S, O, R> WorkflowInterceptor.intercept(
   }
 }
 
+internal fun WorkflowSession.workflowSessionToString(): String {
+  val parentDescription = parent?.let { "WorkflowInstance(…)" }
+  return "WorkflowInstance(" +
+    "identifier=$identifier, " +
+    "renderKey=$renderKey, " +
+    "instanceId=$sessionId, " +
+    "parent=$parentDescription" +
+    ")"
+}
+
 private class InterceptedRenderContext<P, S, O>(
   private val baseRenderContext: BaseRenderContext<P, S, O>,
   private val interceptor: RenderContextInterceptor<P, S, O>

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

@@ -0,0 +1,130 @@
+package com.squareup.workflow1.internal
+
+import com.squareup.workflow1.ActionApplied
+import com.squareup.workflow1.ActionProcessingResult
+import com.squareup.workflow1.ActionsExhausted
+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.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> = 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,
+  protected val interceptor: WorkflowInterceptor,
+  protected val emitAppliedActionToParent: (ActionApplied<OutputT>) -> ActionProcessingResult,
+  baseContext: CoroutineContext,
+) {
+
+  /**
+   * Scope that has a job that will live as long as this node and be cancelled when [cancel] is
+   * called.
+   * Also adds a debug name to this coroutine based on its ID.
+   */
+  val scope: CoroutineScope = CoroutineScope(
+    baseContext +
+      Job(parent = baseContext[Job]) +
+      CoroutineName(id.toString())
+  )
+
+  /**
+   * The [WorkflowSession] that represents this node to [WorkflowInterceptor]s.
+   */
+  abstract val session: WorkflowSession
+
+  /**
+   * 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
+
+  /**
+   * Register select clauses for the next [result][ActionProcessingResult] from the state machine.
+   *
+   * Walk the tree of state machines, asking each one to wait for its next event. If something
+   * happens 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].
+   *
+   * Contrast this to [applyNextAvailableTreeAction], which is used to check for an action
+   * that is already available without waiting, and then _immediately_ apply it.
+   */
+  abstract fun registerTreeActionSelectors(selector: SelectBuilder<ActionProcessingResult>)
+
+  /**
+   * Will try to apply any immediately available actions in this action queue or any of our
+   * children's.
+   *
+   * Contrast this to [registerTreeActionSelectors] which will add select clauses that will await
+   * the next action.
+   *
+   * @param skipDirtyNodes Whether or not this should skip over any workflow nodes that are already
+   * 'dirty' - that is, they had their own state changed as the result of a previous action before
+   * the next render pass.
+   *
+   * @return [ActionProcessingResult] of the action processed, or [ActionsExhausted] if there were
+   * none immediately available.
+   */
+  abstract fun applyNextAvailableTreeAction(skipDirtyNodes: Boolean = false): ActionProcessingResult
+
+  /**
+   * Cancels this state machine host, and any coroutines started as children of it.
+   *
+   * This must be called when the caller will no longer call [registerTreeActionSelectors]. It is an
+   * error to call [registerTreeActionSelectors] after calling this method.
+   */
+  open fun cancel(cause: CancellationException? = null) {
+    scope.cancel(cause)
+  }
+}

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

@@ -24,12 +24,12 @@ import com.squareup.workflow1.intercept
 import com.squareup.workflow1.internal.RealRenderContext.RememberStore
 import com.squareup.workflow1.internal.RealRenderContext.SideEffectRunner
 import com.squareup.workflow1.trace
+import com.squareup.workflow1.workflowSessionToString
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.CoroutineStart.LAZY
 import kotlinx.coroutines.Job
-import kotlinx.coroutines.cancel
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.channels.Channel.Factory.UNLIMITED
 import kotlinx.coroutines.launch
@@ -49,38 +49,42 @@ 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 },
+  emitAppliedActionToParent: (ActionApplied<OutputT>) -> ActionProcessingResult = { it },
   override val parent: WorkflowSession? = null,
-  private val interceptor: WorkflowInterceptor = NoopWorkflowInterceptor,
+  interceptor: WorkflowInterceptor = NoopWorkflowInterceptor,
   idCounter: IdCounter? = null
-) : CoroutineScope, SideEffectRunner, RememberStore, WorkflowSession {
-
-  /**
-   * 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
+) : AbstractWorkflowNode<PropsT, OutputT, RenderingT>(
+  id = id,
+  baseContext = baseContext,
+  interceptor = interceptor,
+  emitAppliedActionToParent = emitAppliedActionToParent,
+),
+  SideEffectRunner,
+  RememberStore,
+  WorkflowSession {
+
+  // WorkflowSession 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>
 
+  override val session: WorkflowSession
+    get() = this
+
   private val subtreeManager = SubtreeManager(
     snapshotCache = snapshot?.childTreeSnapshots,
-    contextForChildren = coroutineContext,
+    contextForChildren = scope.coroutineContext,
     emitActionToParent = ::applyAction,
     runtimeConfig = runtimeConfig,
     workflowTracer = workflowTracer,
@@ -117,52 +121,54 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
   private val context = RenderContext(baseRenderContext, workflow)
 
   init {
-    interceptor.onSessionStarted(this, this)
+    interceptor.onSessionStarted(workflowScope = scope, session = this)
 
     cachedWorkflowInstance = workflow
-    interceptedWorkflowInstance = interceptor.intercept(cachedWorkflowInstance, this)
-    state = interceptedWorkflowInstance.initialState(initialProps, snapshot?.workflowSnapshot, this)
+    interceptedWorkflowInstance = interceptor.intercept(
+      workflow = cachedWorkflowInstance,
+      workflowSession = this
+    )
+    state = interceptedWorkflowInstance.initialState(
+      props = initialProps,
+      snapshot = snapshot?.workflowSnapshot,
+      workflowScope = scope
+    )
   }
 
-  override fun toString(): String {
-    val parentDescription = parent?.let { "WorkflowInstance(…)" }
-    return "WorkflowInstance(" +
-      "identifier=$identifier, " +
-      "renderKey=$renderKey, " +
-      "instanceId=$sessionId, " +
-      "parent=$parentDescription" +
-      ")"
-  }
+  override fun toString(): String = workflowSessionToString()
 
   /**
    * 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(
@@ -206,20 +212,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     return result.lastCalculated as ResultT
   }
 
-  /**
-   * Register select clauses for the next [result][ActionProcessingResult] from the state machine.
-   *
-   * 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].
-   *
-   * Contrast this to [applyNextAvailableTreeAction], which is used to check for an action
-   * that is already available without waiting, and then _immediately_ apply it.
-   * The select clauses added here also call [applyAction] when one of them is selected.
-   */
-  fun registerTreeActionSelectors(selector: SelectBuilder<ActionProcessingResult>) {
+  override fun registerTreeActionSelectors(selector: SelectBuilder<ActionProcessingResult>) {
     // Listen for any child workflow updates.
     subtreeManager.registerChildActionSelectors(selector)
 
@@ -231,22 +224,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     }
   }
 
-  /**
-   * Will try to apply any immediately available actions in this action queue or any of our
-   * children's.
-   *
-   * Contrast this to [registerTreeActionSelectors] which will add select clauses that will await
-   * the next action. That will also end up with [applyAction] being called when the clauses is
-   * selected.
-   *
-   * @param skipDirtyNodes Whether or not this should skip over any workflow nodes that are already
-   * 'dirty' - that is, they had their own state changed as the result of a previous action before
-   * the next render pass.
-   *
-   * @return [ActionProcessingResult] of the action processed, or [ActionsExhausted] if there were
-   * none immediately available.
-   */
-  fun applyNextAvailableTreeAction(skipDirtyNodes: Boolean = false): ActionProcessingResult {
+  override fun applyNextAvailableTreeAction(skipDirtyNodes: Boolean): ActionProcessingResult {
     if (skipDirtyNodes && selfStateDirty) return ActionsExhausted
 
     val result = subtreeManager.applyNextAvailableChildAction(skipDirtyNodes)
@@ -262,11 +240,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 [registerTreeActionSelectors]. It is an error to call [registerTreeActionSelectors]
-   * after calling this method.
+   * This must be called when the caller will no longer call [registerTreeActionSelectors]. It is an
+   * error to call [registerTreeActionSelectors] after calling this method.
    */
-  fun cancel(cause: CancellationException? = null) {
-    coroutineContext.cancel(cause)
+  override fun cancel(cause: CancellationException?) {
+    super.cancel(cause)
     lastRendering = NullableInitBox()
   }
 
@@ -348,7 +326,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>,
@@ -389,7 +366,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     sideEffect: suspend CoroutineScope.() -> Unit
   ): SideEffectNode {
     return workflowTracer.trace("CreateSideEffectNode") {
-      val scope = this + CoroutineName("sideEffect[$key] for $id")
+      val scope = scope + CoroutineName("sideEffect[$key] for $id")
       val job = scope.launch(start = LAZY, block = sideEffect)
       SideEffectNode(key, job)
     }