1249: Cache Interceptor Instance; Do not pass Workflow to render()

Author: steve-the-edwards

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

@@ -113,11 +113,32 @@ public interface Workflow<in PropsT, out OutputT, out RenderingT> {
 
 /**
  * Uses the given [function][transform] to transform a [Workflow] that
- * renders [FromRenderingT] to one renders [ToRenderingT],
+ * renders [FromRenderingT] to one renders [ToRenderingT].
+ *
+ * Note that since this uses the [identifier] of the Workflow being mapped, this can only be used
+ * *once* in a given Workflow unless a different key is passed for the subsequent [renderChild]
+ * call.
+ *
+ * In other words, if you would like to use this for multiple children workflows, you *must* render
+ * them each with separate keys. e.g.:
+ * ```
+ *     when (props) {
+ *       0 -> renderChild(childWorkflow.mapRendering { "rendering1: $it" })
+ *       1 -> renderChild(childWorkflow.mapRendering { "rendering2: $it" }, "rendering2")
+ *       else -> fail()
+ *     }
+ * ```
+ * If you do not, the first workflow instance returned by [mapRendering] will simply be used for
+ * all subsequent.
+ *
+ * If this is cumbersome or inconvenient, I would recommend simply calling [transform] directly
+ * on the rendering itself if it is cheap and idempotent. If it is not, well then, you should
+ * likely just create 2 different concrete workflow types with the same interface, base class,
+ * or some other form of re-use.
  */
 public fun <PropsT, OutputT, FromRenderingT, ToRenderingT>
   Workflow<PropsT, OutputT, FromRenderingT>.mapRendering(
-    transform: (FromRenderingT) -> ToRenderingT
+    transform: (FromRenderingT) -> ToRenderingT,
   ): Workflow<PropsT, OutputT, ToRenderingT> =
   object : StatelessWorkflow<PropsT, OutputT, ToRenderingT>(), ImpostorWorkflow {
     override val realIdentifier: WorkflowIdentifier get() = this@mapRendering.identifier

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

@@ -129,9 +129,9 @@ public class WorkflowIdentifier internal constructor(
   }
 
   /**
-   * Used to detect when this [WorkflowIdentifier] is a deeply stubbed mock. Stubs are provided
-   * until a primitive, in this case the typeName. This lets us determine if we are a mock
-   * object, or a real one. Mea culpa.
+   * Used to detect when this [WorkflowIdentifier] is a deeply stubbed mock. When mocking, object
+   * Stubs are provided until a primitive is hit, in this case the typeName (a [String]).
+   * This lets us determine if we are a mock object, or a real one. Mea culpa.
    */
   internal val deepNameCheck = type.typeName
 
@@ -177,11 +177,11 @@ public val Workflow<*, *, *>.identifier: WorkflowIdentifier
       is IdCacheable -> {
         // The following lines look more complex than they need to be. If we have not yet cached
         // the identifier, we do. But we return the [computedIdentifier] value directly as in the
-        // case of tests which use mocks we want to ensure that we return what is on line 180 -
+        // case of tests which use mocks we want to ensure that we return what is on line 203 -
         // "WorkflowIdentifier(Snapshottable(this::class))" as that depends solely on types.
-        // We do the 'senseless' comparison of .type.typeName here to detect the case where we
+        // We do the 'senseless' comparison of .deepNameCheck here to detect the case where we
         // we have mocks with deep stubs but the name itself (a String) is null.
-        // The reason this is so complicated is this caching has been added afterword via the
+        // The reason this is so complicated is this caching has been added afterwards via the
         // [IdCacheable] interface so that the [Workflow] interface itself remains unchanged.
         @Suppress("SENSELESS_COMPARISON")
         if (cachedIdentifier == null || cachedIdentifier!!.deepNameCheck == null) {

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

@@ -141,7 +141,7 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
         )
       }
     stagedChild.setHandler(handler)
-    return stagedChild.render(child.asStatefulWorkflow(), props)
+    return stagedChild.render(props)
   }
 
   /**
@@ -163,8 +163,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.WorkflowTracer
@@ -51,14 +50,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
+    return workflowNode.render(props as ChildPropsT) as R
   }
 
   /**

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

@@ -71,6 +71,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
   override val identifier: WorkflowIdentifier get() = id.identifier
   override val renderKey: String get() = id.name
   override val sessionId: Long = idCounter.createId()
+  private val interceptedWorkflow: StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>
 
   private val subtreeManager = SubtreeManager(
     snapshotCache = snapshot?.childTreeSnapshots,
@@ -99,8 +100,8 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
   init {
     interceptor.onSessionStarted(this, this)
 
-    state = interceptor.intercept(workflow, this)
-      .initialState(initialProps, snapshot?.workflowSnapshot, this)
+    interceptedWorkflow = interceptor.intercept(workflow, this)
+    state = interceptedWorkflow.initialState(initialProps, snapshot?.workflowSnapshot, this)
   }
 
   override fun toString(): String {
@@ -118,24 +119,35 @@ 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()
+
+    workflowTracer.trace("UpdateRuntimeTree") {
+      // 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)
-        .snapshotState(state)
+      val rootSnapshot = interceptedWorkflow.snapshotState(state)
       TreeSnapshot(
         workflowSnapshot = rootSnapshot,
         // Create the snapshots eagerly since subtreeManager is mutable.
@@ -202,40 +214,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()
-
-    workflowTracer.trace("UpdateRuntimeTree") {
-      // 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)
-        .onPropsChanged(lastProps, newProps, state)
+      val newState = interceptedWorkflow.onPropsChanged(lastProps, newProps, state)
       state = newState
     }
     lastProps = newProps

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

@@ -70,8 +70,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

@@ -151,7 +151,7 @@ class WorkflowOperatorsTest {
   }
 
   @Test
-  fun mapRendering_with_same_upstream_workflow_in_two_different_passes_does_not_restart() {
+  fun mapRendering_with_same_upstream_workflow_in_two_different_passes_uses_same_instance() {
     val trigger = MutableStateFlow("initial")
     val childWorkflow = object : StateFlowWorkflow<String>("child", trigger) {}
     val parentWorkflow = Workflow.stateless<Int, Nothing, String> { props ->
@@ -194,8 +194,76 @@ class WorkflowOperatorsTest {
         listOf(
           "rendering1: initial",
           "rendering1: foo",
+          "rendering1: foo",
+          "rendering1: bar"
+        ),
+        renderings
+      )
+
+      workflowJob.cancel()
+    }
+  }
+
+  @Test
+  fun mapRendering_with_same_upstream_workflow_and_diff_keys_uses_different_instance() {
+    val trigger = MutableStateFlow("initial")
+    val childWorkflow = object : StateFlowWorkflow<String>("child", trigger) {}
+    val parentWorkflow = Workflow.stateless<Int, Nothing, String> { props ->
+      when (props) {
+        0 -> renderChild(childWorkflow.mapRendering { "rendering1: $it" })
+        1 -> renderChild(childWorkflow.mapRendering { "rendering2: $it" }, "rendering2")
+        else -> fail()
+      }
+    }
+    val props = MutableStateFlow(0)
+
+    runTest(UnconfinedTestDispatcher()) {
+      val renderings = mutableListOf<String>()
+      val workflowJob = Job(coroutineContext[Job])
+      renderWorkflowIn(parentWorkflow, this + workflowJob, props) {}
+        .onEach { renderings += it.rendering }
+        .launchIn(this + workflowJob)
+      assertEquals(
+        listOf(
+          "rendering1: initial"
+        ),
+        renderings
+      )
+      assertEquals(1, childWorkflow.starts)
+
+      trigger.value = "foo"
+      assertEquals(1, childWorkflow.starts)
+      assertEquals(
+        listOf(
+          "rendering1: initial",
+          "rendering1: foo"
+        ),
+        renderings
+      )
+
+      props.value = 1
+      // Start another child workflow node.
+      assertEquals(2, childWorkflow.starts)
+      assertEquals(
+        listOf(
+          "rendering1: initial",
+          "rendering1: foo",
+          // 2 renderings, 1 for the props, the 2nd for the worker triggering on the new instance.
+          "rendering2: foo",
+          "rendering2: foo",
+        ),
+        renderings
+      )
+
+      trigger.value = "bar"
+      assertEquals(2, childWorkflow.starts)
+      assertEquals(
+        listOf(
+          "rendering1: initial",
+          "rendering1: foo",
+          "rendering2: foo",
           "rendering2: foo",
-          "rendering2: bar"
+          "rendering2: bar",
         ),
         renderings
       )