Merge pull request #1342 from square/sedwards/stateless-cache

1337: Cache `StatelessWorkflow.RenderContext` in the `StatefulWorkflow.RenderContext`

Author: steve-the-edwards

workflow-core/api/workflow-core.api

@@ -257,6 +257,7 @@ public final class com/squareup/workflow1/StatefulWorkflow$Companion {
 }
 
 public final class com/squareup/workflow1/StatefulWorkflow$RenderContext : com/squareup/workflow1/BaseRenderContext {
+	public final fun asStatelessRenderContext ()Lcom/squareup/workflow1/StatelessWorkflow$RenderContext;
 	public final fun eventHandler (Ljava/lang/String;Ljava/lang/Boolean;Lkotlin/jvm/functions/Function1;)Lkotlin/jvm/functions/Function0;
 	public static synthetic fun eventHandler$default (Lcom/squareup/workflow1/StatefulWorkflow$RenderContext;Ljava/lang/String;Ljava/lang/Boolean;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lkotlin/jvm/functions/Function0;
 	public fun getActionSink ()Lcom/squareup/workflow1/Sink;
@@ -288,17 +289,6 @@ public final class com/squareup/workflow1/StatelessWorkflow$RenderContext : com/
 	public fun runningSideEffect (Ljava/lang/String;Lkotlin/jvm/functions/Function2;)V
 }
 
-public final class com/squareup/workflow1/StatelessWorkflow$StatelessAsStatefulWorkflow : com/squareup/workflow1/StatefulWorkflow {
-	public fun <init> (Lcom/squareup/workflow1/StatelessWorkflow;)V
-	public final fun clearCache ()V
-	public synthetic fun initialState (Ljava/lang/Object;Lcom/squareup/workflow1/Snapshot;)Ljava/lang/Object;
-	public fun initialState (Ljava/lang/Object;Lcom/squareup/workflow1/Snapshot;)V
-	public synthetic fun render (Ljava/lang/Object;Ljava/lang/Object;Lcom/squareup/workflow1/StatefulWorkflow$RenderContext;)Ljava/lang/Object;
-	public fun render (Ljava/lang/Object;Lkotlin/Unit;Lcom/squareup/workflow1/StatefulWorkflow$RenderContext;)Ljava/lang/Object;
-	public synthetic fun snapshotState (Ljava/lang/Object;)Lcom/squareup/workflow1/Snapshot;
-	public fun snapshotState (Lkotlin/Unit;)Lcom/squareup/workflow1/Snapshot;
-}
-
 public final class com/squareup/workflow1/TypedWorker : com/squareup/workflow1/Worker {
 	public fun <init> (Lkotlin/reflect/KType;Lkotlinx/coroutines/flow/Flow;)V
 	public fun doesSameWorkAs (Lcom/squareup/workflow1/Worker;)Z
@@ -424,7 +414,6 @@ public final class com/squareup/workflow1/WorkflowTracerKt {
 
 public final class com/squareup/workflow1/Workflows {
 	public static final fun RenderContext (Lcom/squareup/workflow1/BaseRenderContext;Lcom/squareup/workflow1/StatefulWorkflow;)Lcom/squareup/workflow1/StatefulWorkflow$RenderContext;
-	public static final fun RenderContext (Lcom/squareup/workflow1/BaseRenderContext;Lcom/squareup/workflow1/StatelessWorkflow;)Lcom/squareup/workflow1/StatelessWorkflow$RenderContext;
 	public static final fun action (Lcom/squareup/workflow1/StatefulWorkflow;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lcom/squareup/workflow1/WorkflowAction;
 	public static final fun action (Lcom/squareup/workflow1/StatefulWorkflow;Lkotlin/jvm/functions/Function0;Lkotlin/jvm/functions/Function1;)Lcom/squareup/workflow1/WorkflowAction;
 	public static final fun action (Lcom/squareup/workflow1/StatelessWorkflow;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lcom/squareup/workflow1/WorkflowAction;

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

@@ -706,6 +706,20 @@ public abstract class StatefulWorkflow<
           }
           ?: onFailedCast(name, CurrentStateT::class, state)
       }
+
+    private var statelessContext: StatelessWorkflow.RenderContext<PropsT, OutputT>? = null
+
+    /**
+     * Fetch the [StatelessWorkflow.RenderContext] that is equivalent to this [RenderContext]
+     * with the `StateT` type erased.
+     *
+     * We cache this value so that it does not need to be continually recreated and will be stable
+     * for compose. See [statelessContext] for the instance.
+     */
+    public fun asStatelessRenderContext(): StatelessWorkflow.RenderContext<PropsT, OutputT> =
+      statelessContext ?: StatelessWorkflow.RenderContext(this).also {
+        statelessContext = it
+      }
   }
 
   /**

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

@@ -227,23 +227,9 @@ public abstract class StatelessWorkflow<PropsT, OutputT, out RenderingT> :
    * Class type returned by [asStatefulWorkflow].
    * See [statefulWorkflow] for the instance.
    */
-  inner class StatelessAsStatefulWorkflow :
+  private inner class StatelessAsStatefulWorkflow :
     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.RenderContext<PropsT, OutputT>? = 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.RenderContext<PropsT, Unit, OutputT>? = null
-
     override fun initialState(
       props: PropsT,
       snapshot: Snapshot?
@@ -254,37 +240,12 @@ public abstract class StatelessWorkflow<PropsT, OutputT, out RenderingT> :
       renderState: Unit,
       context: RenderContext<PropsT, Unit, OutputT>
     ): 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!!)
+      // Pass `asStatelessRenderContext()`. Caching for that will be handled by the
+      // `StatefulWorkflow.RenderContext` itself.
+      return render(renderProps, context.asStatelessRenderContext())
     }
 
     override fun snapshotState(state: Unit): Snapshot? = null
-
-    /**
-     * When we are finished with at least one node that holds on to this workflow instance,
-     * then we clear the cache. The reason we do that every time is that it *might* be the last
-     * node that is caching this instance, and if so, we do not want to leak these cached
-     * render contexts.
-     *
-     * Yes, that means that it might have to be re-created again when this instance is used
-     * multiple times. The current design for how we get a [StatefulWorkflow] from the
-     * [StatelessWorkflow] is a failed compromise between performance (caching) and type-safe
-     * brevity (erasing the `StateT` type from the concerns of [StatelessWorkflow]). It needs
-     * to be fixed with a bigger re-write (https://github.com/square/workflow-kotlin/issues/1337).
-     */
-    fun clearCache() {
-      cachedStatelessRenderContext = null
-      canonicalStatefulRenderContext = null
-    }
   }
 
   private val statefulWorkflow: StatefulWorkflow<PropsT, Unit, OutputT, RenderingT> =
@@ -325,17 +286,6 @@ public abstract class StatelessWorkflow<PropsT, OutputT, out RenderingT> :
     statefulWorkflow
 }
 
-/**
- * Creates a `RenderContext` from a [BaseRenderContext] for the given [StatelessWorkflow].
- */
-@Suppress("UNCHECKED_CAST")
-public fun <PropsT, OutputT, RenderingT> RenderContext(
-  baseContext: BaseRenderContext<PropsT, *, OutputT>,
-  workflow: StatelessWorkflow<PropsT, OutputT, RenderingT>
-): StatelessWorkflow.RenderContext<PropsT, OutputT> =
-  (baseContext as? StatelessWorkflow.RenderContext<PropsT, OutputT>)
-    ?: StatelessWorkflow.RenderContext<PropsT, OutputT>(baseContext)
-
 /**
  * Returns a stateless [Workflow] via the given [render] function.
  *

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

@@ -9,7 +9,6 @@ import com.squareup.workflow1.RuntimeConfig
 import com.squareup.workflow1.RuntimeConfigOptions
 import com.squareup.workflow1.RuntimeConfigOptions.PARTIAL_TREE_RENDERING
 import com.squareup.workflow1.StatefulWorkflow
-import com.squareup.workflow1.StatelessWorkflow
 import com.squareup.workflow1.TreeSnapshot
 import com.squareup.workflow1.Workflow
 import com.squareup.workflow1.WorkflowAction
@@ -237,10 +236,6 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
   fun cancel(cause: CancellationException? = null) {
     coroutineContext.cancel(cause)
     lastRendering = NullableInitBox()
-    (
-      cachedWorkflowInstance as?
-        StatelessWorkflow<PropsT, OutputT, RenderingT>.StatelessAsStatefulWorkflow
-      )?.clearCache()
   }
 
   /**