988: Cache RenderContext per instance

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,12 +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)) }
-  )
-
   /**
    * Called at least once any time one of the following things happens:
    *  - This workflow's [renderProps] change (via the parent passing a different one in).
@@ -70,11 +63,43 @@ public abstract class StatelessWorkflow<in PropsT, out OutputT, out RenderingT>
    * Satisfies the [Workflow] interface by wrapping `this` in a [StatefulWorkflow] with `Unit`
    * state.
    *
-   * This method is called a few times per instance, but we don't need to allocate a new
-   * [StatefulWorkflow] every time, so we store it in a private property.
+   * This is only called when the instance of the Workflow is created, so returning a new instance
+   * of the [StatefulWorkflow] each time is fine.
    */
-  final override fun asStatefulWorkflow(): StatefulWorkflow<PropsT, *, OutputT, RenderingT> =
-    statefulWorkflow
+  final override fun asStatefulWorkflow(): StatefulWorkflow<PropsT, *, OutputT, RenderingT> {
+    return object : StatefulWorkflow<PropsT, Unit, OutputT, RenderingT>() {
+      // We want to cache the render context so that we don't have to recreate it each time
+      // render() is called.
+      private var cachedStatelessRenderContext:
+        StatelessWorkflow<PropsT, OutputT, RenderingT>.RenderContext? = null
+      private var cachedStatefulRenderContext:
+        StatefulWorkflow<PropsT, Unit, OutputT, RenderingT>.RenderContext? = null
+
+      override fun initialState(
+        props: PropsT,
+        snapshot: Snapshot?
+      ) = Unit
+
+      override fun render(
+        renderProps: PropsT,
+        renderState: Unit,
+        context: RenderContext
+      ): RenderingT {
+        // The `RenderContext` used *might* change - primarily in the case of our tests. E.g., The
+        // `RenderTester` uses a special NoOp context to render twice to test for idempotency.
+        // In order to support a changed render context but keep caching, we check to see if the
+        // instance passed in has changed.
+        if (cachedStatelessRenderContext == null || context != cachedStatefulRenderContext) {
+          // Recreate it if the underlying context changed.
+          cachedStatelessRenderContext = RenderContext(context, this@StatelessWorkflow)
+        }
+        cachedStatefulRenderContext = context
+        return render(renderProps, cachedStatelessRenderContext!!)
+      }
+
+      override fun snapshotState(state: Unit): Snapshot? = null
+    }
+  }
 }
 
 /**

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

@@ -268,6 +268,9 @@ public object NoopWorkflowInterceptor : WorkflowInterceptor
 /**
  * Returns a [StatefulWorkflow] that will intercept all calls to [workflow] via this
  * [WorkflowInterceptor].
+ *
+ * This is called once for each instance/session of a Workflow being intercepted. So we cache the
+ * render context for re-use within that [WorkflowSession].
  */
 @OptIn(WorkflowExperimentalApi::class)
 internal fun <P, S, O, R> WorkflowInterceptor.intercept(
@@ -277,6 +280,10 @@ internal fun <P, S, O, R> WorkflowInterceptor.intercept(
   workflow
 } else {
   object : SessionWorkflow<P, S, O, R>() {
+
+    private var cachedOriginalRenderContext: StatefulWorkflow<P, S, O, R>.RenderContext? = null
+    private var cachedInterceptedRenderContext: StatefulWorkflow<P, S, O, R>.RenderContext? = null
+
     override fun initialState(
       props: P,
       snapshot: Snapshot?,
@@ -298,9 +305,17 @@ 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))
+        // The `RenderContext` used *might* change - primarily in the case of our tests. E.g., The
+        // `RenderTester` uses a special NoOp context to render twice to test for idempotency.
+        // In order to support a changed render context but keep caching, we check to see if the
+        // instance passed in has changed.
+        if (cachedInterceptedRenderContext == null || cachedInterceptedRenderContext != context) {
+          val interceptedRenderContext = interceptor?.let { InterceptedRenderContext(context, it) }
+            ?: context
+          cachedInterceptedRenderContext = RenderContext(interceptedRenderContext, this)
+        }
+        cachedOriginalRenderContext = context
+        workflow.render(props, state, cachedInterceptedRenderContext!!)
       },
       session = workflowSession,
     )