DNM: WIP: Keep the same RenderContext through Interceptor

workflow-core/src/commonMain/kotlin/com/squareup/workflow1/StatelessWorkflow.kt

@@ -3,7 +3,6 @@
 
 package com.squareup.workflow1
 
-import kotlin.LazyThreadSafetyMode.NONE
 import kotlin.jvm.JvmMultifileClass
 import kotlin.jvm.JvmName
 
@@ -33,7 +32,6 @@ public abstract class StatelessWorkflow<in PropsT, out OutputT, out RenderingT>
   ) : BaseRenderContext<@UnsafeVariance PropsT, Nothing, @UnsafeVariance OutputT> by
   baseContext as BaseRenderContext<PropsT, Nothing, OutputT>
 
-  @Suppress("UNCHECKED_CAST")
   private val statefulWorkflow = Workflow.stateful<PropsT, Unit, OutputT, RenderingT>(
     initialState = { Unit },
     render = { props, _ -> render(props, RenderContext(this, this@StatelessWorkflow)) }
@@ -122,7 +120,7 @@ public fun <RenderingT> Workflow.Companion.rendering(
  * @param update Function that defines the workflow update.
  */
 public fun <PropsT, OutputT, RenderingT>
-StatelessWorkflow<PropsT, OutputT, RenderingT>.action(
+  StatelessWorkflow<PropsT, OutputT, RenderingT>.action(
   name: String = "",
   update: WorkflowAction<PropsT, *, OutputT>.Updater.() -> Unit
 ): WorkflowAction<PropsT, Nothing, OutputT> = action({ name }, update)
@@ -137,7 +135,7 @@ StatelessWorkflow<PropsT, OutputT, RenderingT>.action(
  * @param update Function that defines the workflow update.
  */
 public fun <PropsT, OutputT, RenderingT>
-StatelessWorkflow<PropsT, OutputT, RenderingT>.action(
+  StatelessWorkflow<PropsT, OutputT, RenderingT>.action(
   name: () -> String,
   update: WorkflowAction<PropsT, *, OutputT>.Updater.() -> Unit
 ): WorkflowAction<PropsT, Nothing, OutputT> = object : WorkflowAction<PropsT, Nothing, OutputT>() {

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

@@ -263,6 +263,10 @@ internal fun <P, S, O, R> WorkflowInterceptor.intercept(
   workflow
 } else {
   object : StatefulWorkflow<P, S, O, R>() {
+
+    private lateinit var interceptedContext: BaseRenderContext<P, S, O>
+    private lateinit var concreteRenderContext: StatefulWorkflow<P, S, O, R>.RenderContext
+
     override fun initialState(
       props: P,
       snapshot: Snapshot?
@@ -283,9 +287,12 @@ internal fun <P, S, O, R> WorkflowInterceptor.intercept(
       renderState,
       context,
       proceed = { props, state, interceptor ->
-        val interceptedContext = interceptor?.let { InterceptedRenderContext(context, it) }
-          ?: context
-        workflow.render(props, state, RenderContext(interceptedContext, this))
+        if (!this::interceptedContext.isInitialized) {
+          interceptedContext = interceptor?.let { InterceptedRenderContext(context, it) }
+            ?: context
+          concreteRenderContext = RenderContext(interceptedContext, this)
+        }
+        workflow.render(props, state, concreteRenderContext)
       },
       session = workflowSession,
     )

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

@@ -129,7 +129,7 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
       create = { createChildNode(child, props, key, handler) }
     )
     stagedChild.setHandler(handler)
-    return stagedChild.render(child.asStatefulWorkflow(), props)
+    return stagedChild.render(props)
   }
 
   /**
@@ -151,8 +151,7 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
   fun createChildSnapshots(): Map<WorkflowNodeId, TreeSnapshot> {
     val snapshots = mutableMapOf<WorkflowNodeId, TreeSnapshot>()
     children.forEachActive { child ->
-      val childWorkflow = child.workflow.asStatefulWorkflow()
-      snapshots[child.id] = child.workflowNode.snapshot(childWorkflow)
+      snapshots[child.id] = child.workflowNode.snapshot()
     }
     return snapshots
   }

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

@@ -1,6 +1,5 @@
 package com.squareup.workflow1.internal
 
-import com.squareup.workflow1.StatefulWorkflow
 import com.squareup.workflow1.Workflow
 import com.squareup.workflow1.WorkflowAction
 import com.squareup.workflow1.internal.InlineLinkedList.InlineListNode
@@ -48,12 +47,10 @@ internal class WorkflowChildNode<
    * Wrapper around [WorkflowNode.render] that allows calling it with erased types.
    */
   fun <R> render(
-    workflow: StatefulWorkflow<*, *, *, *>,
     props: Any?
   ): R {
     @Suppress("UNCHECKED_CAST")
     return workflowNode.render(
-      workflow as StatefulWorkflow<ChildPropsT, out Any?, ChildOutputT, Nothing>,
       props as ChildPropsT
     ) as R
   }

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

@@ -81,12 +81,12 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     eventActionsChannel = eventActionsChannel
   )
   private val context = RenderContext(baseRenderContext, workflow)
+  private val interceptedWorkflow = interceptor.intercept(workflow, this)
 
   init {
     interceptor.onSessionStarted(this, this)
 
-    state = interceptor.intercept(workflow, this)
-      .initialState(initialProps, snapshot?.workflowSnapshot)
+    state = interceptedWorkflow.initialState(initialProps, snapshot?.workflowSnapshot)
   }
 
   override fun toString(): String {
@@ -104,23 +104,34 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
    * [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>,
-    input: PropsT
-  ): RenderingT =
-    renderWithStateType(workflow as StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>, input)
+    props: PropsT
+  ): RenderingT {
+    updatePropsAndState(props)
+
+    baseRenderContext.unfreeze()
+    val rendering = interceptedWorkflow
+      .render(props, state, context)
+    baseRenderContext.freeze()
+
+    // Tear down workflows and workers that are obsolete.
+    subtreeManager.commitRenderedChildren()
+    // Side effect jobs are launched lazily, since they can send actions to the sink, and can only
+    // be started after context is frozen.
+    sideEffects.forEachStaging { it.job.start() }
+    sideEffects.commitStaging { it.job.cancel() }
+
+    return rendering
+  }
 
   /**
    * 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>
+  fun snapshot(): TreeSnapshot {
     return interceptor.onSnapshotStateWithChildren({
       val childSnapshots = subtreeManager.createChildSnapshots()
-      val rootSnapshot = interceptor.intercept(typedWorkflow, this)
+      val rootSnapshot = interceptedWorkflow
         .snapshotState(state)
       TreeSnapshot(
         workflowSnapshot = rootSnapshot,
@@ -188,37 +199,11 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     coroutineContext.cancel(cause)
   }
 
-  /**
-   * Contains the actual logic for [render], after we've casted the passed-in [Workflow]'s
-   * state type to our `StateT`.
-   */
-  private fun renderWithStateType(
-    workflow: StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>,
-    props: PropsT
-  ): RenderingT {
-    updatePropsAndState(workflow, props)
-
-    baseRenderContext.unfreeze()
-    val rendering = interceptor.intercept(workflow, this)
-      .render(props, state, context)
-    baseRenderContext.freeze()
-
-    // Tear down workflows and workers that are obsolete.
-    subtreeManager.commitRenderedChildren()
-    // Side effect jobs are launched lazily, since they can send actions to the sink, and can only
-    // be started after context is frozen.
-    sideEffects.forEachStaging { it.job.start() }
-    sideEffects.commitStaging { it.job.cancel() }
-
-    return rendering
-  }
-
   private fun updatePropsAndState(
-    workflow: StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>,
     newProps: PropsT
   ) {
     if (newProps != lastProps) {
-      val newState = interceptor.intercept(workflow, this)
+      val newState = interceptedWorkflow
         .onPropsChanged(lastProps, newProps, state)
       state = newState
     }

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

@@ -66,8 +66,8 @@ internal class WorkflowRunner<PropsT, OutputT, RenderingT>(
    */
   fun nextRendering(): RenderingAndSnapshot<RenderingT> {
     return interceptor.onRenderAndSnapshot(currentProps, { props ->
-      val rendering = rootNode.render(workflow, props)
-      val snapshot = rootNode.snapshot(workflow)
+      val rendering = rootNode.render(props)
+      val snapshot = rootNode.snapshot()
       RenderingAndSnapshot(rendering, snapshot)
     }, rootNode)
   }

workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/WorkflowOperatorsTest.kt

@@ -11,6 +11,7 @@ import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.plus
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import kotlin.test.Test
 import kotlin.test.assertEquals
@@ -157,19 +158,24 @@ class WorkflowOperatorsTest {
     val childWorkflow = object : StateFlowWorkflow<String>("child", trigger) {}
     val parentWorkflow = Workflow.stateless<Int, Nothing, String> { props ->
       when (props) {
+        // I don't understand this test. [mapRendering] defers to its wrapped Workflow for the identifier
+        // which is equivalent - but we are relying on the fact that we pass in the Workflow each time we render
+        // so we pull out the *same* child WorkflowNode (starts = 1) but give it a new Workflow to render.
+        // Why?
         0 -> renderChild(childWorkflow.mapRendering { "rendering1: $it" })
         1 -> renderChild(childWorkflow.mapRendering { "rendering2: $it" })
         else -> fail()
       }
     }
     val props = MutableStateFlow(0)
 
-    runTest(UnconfinedTestDispatcher()) {
+    runTest {
       val renderings = mutableListOf<String>()
       val workflowJob = Job(coroutineContext[Job])
       renderWorkflowIn(parentWorkflow, this + workflowJob, props) {}
         .onEach { renderings += it.rendering }
         .launchIn(this + workflowJob)
+      runCurrent()
       assertEquals(
         listOf(
           "rendering1: initial"
@@ -179,6 +185,7 @@ class WorkflowOperatorsTest {
       assertEquals(1, childWorkflow.starts)
 
       trigger.value = "foo"
+      runCurrent()
       assertEquals(1, childWorkflow.starts)
       assertEquals(
         listOf(
@@ -189,7 +196,9 @@ class WorkflowOperatorsTest {
       )
 
       props.value = 1
+      runCurrent()
       trigger.value = "bar"
+      runCurrent()
       assertEquals(1, childWorkflow.starts)
       assertEquals(
         listOf(