1247: Partial Tree Rerendering

Track whether or not the state (or state of a child) has changed in the WorkflowNode. Pass lastRendering if its not.

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

@@ -6,11 +6,13 @@ import com.squareup.workflow1.NoopWorkflowInterceptor
 import com.squareup.workflow1.RenderContext
 import com.squareup.workflow1.RuntimeConfig
 import com.squareup.workflow1.RuntimeConfigOptions
+import com.squareup.workflow1.RuntimeConfigOptions.RENDER_ONLY_WHEN_STATE_CHANGES
 import com.squareup.workflow1.StatefulWorkflow
 import com.squareup.workflow1.TreeSnapshot
 import com.squareup.workflow1.Workflow
 import com.squareup.workflow1.WorkflowAction
 import com.squareup.workflow1.WorkflowExperimentalApi
+import com.squareup.workflow1.WorkflowExperimentalRuntime
 import com.squareup.workflow1.WorkflowIdentifier
 import com.squareup.workflow1.WorkflowInterceptor
 import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
@@ -44,7 +46,7 @@ import kotlin.coroutines.CoroutineContext
  * hard-coded values added to worker contexts. It must not contain a [Job] element (it would violate
  * structured concurrency).
  */
-@OptIn(WorkflowExperimentalApi::class)
+@OptIn(WorkflowExperimentalApi::class, WorkflowExperimentalRuntime::class)
 internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
   val id: WorkflowNodeId,
   workflow: StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>,
@@ -86,9 +88,11 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
   )
   private val sideEffects = ActiveStagingList<SideEffectNode>()
   private var lastProps: PropsT = initialProps
+  private var lastRendering: RenderingT? = null
   private val eventActionsChannel =
     Channel<WorkflowAction<PropsT, StateT, OutputT>>(capacity = UNLIMITED)
   private var state: StateT
+  private var subtreeStateDidChange: Boolean = true
 
   private val baseRenderContext = RealRenderContext(
     renderer = subtreeManager,
@@ -205,16 +209,21 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     coroutineContext.cancel(cause)
   }
 
-  // Call this after we have been passed any workflow instance, in [render] or [snapshot]. It may
-  // have changed and we should check to see if we need to update our cached instances.
+  /** Call this after we have been passed any workflow instance, in [render] or [snapshot]. It may
+   *  have changed and we should check to see if we need to update our cached instances.
+   *
+   *  @return true if the instance has changed, otherwise false.
+   */
   private fun updateCachedWorkflowInstance(
     workflow: StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>
-  ) {
+  ): Boolean {
     if (workflow !== cachedWorkflowInstance) {
       // instance has changed.
       interceptedWorkflowInstance = interceptor.intercept(workflow, this)
       cachedWorkflowInstance = workflow
+      return true
     }
+    return false
   }
 
   /**
@@ -225,33 +234,49 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     workflow: StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>,
     props: PropsT
   ): RenderingT {
-    updateCachedWorkflowInstance(workflow)
-    updatePropsAndState(props)
+    val didUpdateCachedInstance = updatePropsAndState(props, workflow)
 
-    baseRenderContext.unfreeze()
-    val rendering = interceptedWorkflowInstance.render(props, state, context)
-    baseRenderContext.freeze()
+    if (!runtimeConfig.contains(RENDER_ONLY_WHEN_STATE_CHANGES) ||
+      lastRendering == null ||
+      subtreeStateDidChange
+    ) {
+      if (!didUpdateCachedInstance) {
+        // If we haven't already updated the cached instance, better do it now!
+        updateCachedWorkflowInstance(workflow)
+      }
+      baseRenderContext.unfreeze()
+      lastRendering = interceptedWorkflowInstance.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() }
+      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
+    return lastRendering!!
   }
 
+  /**
+   * @return true if the [interceptedWorkflowInstance] has been updated, false otherwise.
+   */
   private fun updatePropsAndState(
-    newProps: PropsT
-  ) {
+    newProps: PropsT,
+    workflow: StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>,
+  ): Boolean {
+    var didUpdateCachedInstance = false
     if (newProps != lastProps) {
+      didUpdateCachedInstance = updateCachedWorkflowInstance(workflow)
       val newState = interceptedWorkflowInstance.onPropsChanged(lastProps, newProps, state)
       state = newState
+      subtreeStateDidChange = true
     }
     lastProps = newProps
+    return didUpdateCachedInstance
   }
 
   /**
@@ -270,6 +295,8 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
       // Changing state is sticky, we pass it up if it ever changed.
       stateChanged = actionApplied.stateChanged || (childResult?.stateChanged ?: false)
     )
+    // Our state changed or one of our children's state changed.
+    subtreeStateDidChange = aggregateActionApplied.stateChanged
     return if (actionApplied.output != null) {
       emitAppliedActionToParent(aggregateActionApplied)
     } else {

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

@@ -207,23 +207,37 @@ class WorkflowOperatorsTest {
   private abstract class StateFlowWorkflow<T>(
     val name: String,
     val flow: StateFlow<T>
-  ) : StatelessWorkflow<Unit, Nothing, T>() {
+  ) : StatefulWorkflow<Unit, T, Nothing, T>() {
     var starts: Int = 0
       private set
 
+    override fun initialState(
+      props: Unit,
+      snapshot: Snapshot?
+    ): T {
+      return flow.value
+    }
+
     private val rerenderWorker = object : Worker<T> {
       override fun run(): Flow<T> = flow.onStart { starts++ }
     }
 
     override fun render(
       renderProps: Unit,
+      renderState: T,
       context: RenderContext
     ): T {
       // Listen to the flow to trigger a re-render when it updates.
-      context.runningWorker(rerenderWorker as Worker<Any?>) { WorkflowAction.noAction() }
-      return flow.value
+      context.runningWorker(rerenderWorker) { output: T ->
+        action("rerenderUpdate") {
+          state = output
+        }
+      }
+      return renderState
     }
 
+    override fun snapshotState(state: T): Snapshot? = null
+
     override fun toString(): String = "StateFlowWorkflow($name)"
   }
 }