988: Cache RenderContext per instance

For the InterceptedRenderContext and for the StatelessWorkflow.RenderContext.

Author: steve-the-edwards

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,11 +32,43 @@ 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)) }
-  )
+  private val statefulWorkflow: StatefulWorkflow<PropsT, Unit, OutputT, RenderingT> =
+    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
+
+      // We must know if the RenderContext we are passed (which is a StatefulWorkflow.RenderContext)
+      // has changed, so keep track of it.
+      private var canonicalStatefulRenderContext:
+        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 !== canonicalStatefulRenderContext) {
+          // Recreate it if the StatefulWorkflow.RenderContext we are passed has changed.
+          cachedStatelessRenderContext = RenderContext(context, this@StatelessWorkflow)
+        }
+        canonicalStatefulRenderContext = context
+        // Pass the StatelessWorkflow.RenderContext to our StatelessWorkflow.
+        return render(renderProps, cachedStatelessRenderContext!!)
+      }
+
+      override fun snapshotState(state: Unit): Snapshot? = null
+    }
 
   /**
    * Called at least once any time one of the following things happens:
@@ -69,9 +100,6 @@ 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.
    */
   final override fun asStatefulWorkflow(): StatefulWorkflow<PropsT, *, OutputT, RenderingT> =
     statefulWorkflow
@@ -123,9 +151,9 @@ public fun <RenderingT> Workflow.Companion.rendering(
  */
 public fun <PropsT, OutputT, RenderingT>
   StatelessWorkflow<PropsT, OutputT, RenderingT>.action(
-    name: String,
-    update: WorkflowAction<PropsT, *, OutputT>.Updater.() -> Unit
-  ): WorkflowAction<PropsT, Nothing, OutputT> = action({ name }, update)
+  name: String,
+  update: WorkflowAction<PropsT, *, OutputT>.Updater.() -> Unit
+): WorkflowAction<PropsT, Nothing, OutputT> = action({ name }, update)
 
 /**
  * Convenience to create a [WorkflowAction] with parameter types matching those
@@ -139,9 +167,9 @@ public fun <PropsT, OutputT, RenderingT>
  */
 public fun <PropsT, OutputT, RenderingT>
   StatelessWorkflow<PropsT, OutputT, RenderingT>.action(
-    name: () -> String,
-    update: WorkflowAction<PropsT, *, OutputT>.Updater.() -> Unit
-  ): WorkflowAction<PropsT, Nothing, OutputT> = object : WorkflowAction<PropsT, Nothing, OutputT>() {
+  name: () -> String,
+  update: WorkflowAction<PropsT, *, OutputT>.Updater.() -> Unit
+): WorkflowAction<PropsT, Nothing, OutputT> = object : WorkflowAction<PropsT, Nothing, OutputT>() {
   override val debuggingName: String
     get() = name()
 

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,16 @@ internal fun <P, S, O, R> WorkflowInterceptor.intercept(
   workflow
 } else {
   object : SessionWorkflow<P, S, O, R>() {
+
+    // Render context that we are passed.
+    private var canonicalRenderContext: StatefulWorkflow<P, S, O, R>.RenderContext? = null
+
+    // Render context interceptor that we are passed.
+    private var canonicalRenderContextInterceptor: RenderContextInterceptor<P, S, O>? = null
+
+    // Cache of the intercepted render context.
+    private var cachedInterceptedRenderContext: StatefulWorkflow<P, S, O, R>.RenderContext? = null
+
     override fun initialState(
       props: P,
       snapshot: Snapshot?,
@@ -298,9 +311,20 @@ 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 || canonicalRenderContext !== context ||
+          canonicalRenderContextInterceptor != interceptor) {
+          val interceptedRenderContext = interceptor?.let { InterceptedRenderContext(context, it) }
+            ?: context
+          cachedInterceptedRenderContext = RenderContext(interceptedRenderContext, this)
+        }
+        canonicalRenderContext = context
+        canonicalRenderContextInterceptor = interceptor
+        // Use the intercepted RenderContext for rendering.
+        workflow.render(props, state, cachedInterceptedRenderContext!!)
       },
       session = workflowSession,
     )