1247: Partial Tree Rerendering

If PARTIAL_TREE_RENDERING is in the RuntimeConfig then,
Track whether or not the state (or state of a descendant) changed in the WorkflowNode. Pass lastRendering if its not.

Also adds tests for this behavior and expands the test runtime matrix.
Also adds shards for the new runtime for instrumentation tests.

Author: steve-the-edwards

.github/workflows/kotlin.yml

@@ -222,6 +222,27 @@ jobs :
         with :
           report_paths : '**/build/test-results/test/TEST-*.xml'
 
+  jvm-partial-runtime-test:
+    name: Partial Tree Rendering Only Runtime JVM Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    steps:
+      - name: Checkout
+        uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
+
+      - name: Check with Gradle
+        uses: ./.github/actions/gradle-task
+        with:
+          task: jvmTest --continue -Pworkflow.runtime=baseline-partial
+          restore-cache-key: main-build-artifacts
+
+      # Report as GitHub Pull Request Check.
+      - name: Publish Test Report
+        uses: mikepenz/action-junit-report@5f47764eec0e1c1f19f40c8e60a5ba47e47015c5 # v4
+        if: always() # always run even if the previous step fails
+        with:
+          report_paths: '**/build/test-results/test/TEST-*.xml'
+
   jvm-conflate-stateChange-runtime-test :
     name : Render On State Change Only and Conflate Stale Runtime JVM Tests
     runs-on : ubuntu-latest
@@ -243,6 +264,27 @@ jobs :
         with :
           report_paths : '**/build/test-results/test/TEST-*.xml'
 
+  jvm-conflate-partial-runtime-test:
+    name: Render On State Change Only and Conflate Stale Runtime JVM Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    steps:
+      - name: Checkout
+        uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
+
+      - name: Check with Gradle
+        uses: ./.github/actions/gradle-task
+        with:
+          task: jvmTest --continue -Pworkflow.runtime=conflate-partial
+          restore-cache-key: main-build-artifacts
+
+      # Report as GitHub Pull Request Check.
+      - name: Publish Test Report
+        uses: mikepenz/action-junit-report@5f47764eec0e1c1f19f40c8e60a5ba47e47015c5 # v4
+        if: always() # always run even if the previous step fails
+        with:
+          report_paths: '**/build/test-results/test/TEST-*.xml'
+
   ios-tests :
     name : iOS Tests
     runs-on : macos-latest
@@ -349,7 +391,7 @@ jobs :
         ### <start-connected-check-shards>
         shardNum: [ 1, 2, 3 ]
         ### <end-connected-check-shards>
-        runtime : [ conflate, baseline-stateChange, conflate-stateChange ]
+        runtime : [ conflate, baseline-stateChange, conflate-stateChange, baseline-partial, conflate-partial ]
     steps :
       - name: Checkout
         uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
@@ -378,8 +420,10 @@ jobs :
       - ios-tests
       - js-tests
       - jvm-conflate-runtime-test
-      - jvm-conflate-stateChange-runtime-test
       - jvm-stateChange-runtime-test
+      - jvm-partial-runtime-test
+      - jvm-conflate-stateChange-runtime-test
+      - jvm-conflate-partial-runtime-test
       - ktlint
       - performance-tests
       - runtime-instrumentation-tests

workflow-config/config-android/src/main/java/com/squareup/workflow1/config/AndroidRuntimeConfigTools.kt

@@ -3,6 +3,7 @@ package com.squareup.workflow1.config
 import com.squareup.workflow1.RuntimeConfig
 import com.squareup.workflow1.RuntimeConfigOptions
 import com.squareup.workflow1.RuntimeConfigOptions.CONFLATE_STALE_RENDERINGS
+import com.squareup.workflow1.RuntimeConfigOptions.PARTIAL_TREE_RENDERING
 import com.squareup.workflow1.RuntimeConfigOptions.RENDER_ONLY_WHEN_STATE_CHANGES
 import com.squareup.workflow1.WorkflowExperimentalRuntime
 
@@ -27,6 +28,9 @@ public class AndroidRuntimeConfigTools {
      * Then, these can be combined (via '-') with:
      * "stateChange" : Only re-render when the state of some WorkflowNode has been changed by an
      *   action cascade.
+     * "partial" : Which includes "stateChange" as well as partial tree rendering, which only
+     *   re-renders each Workflow node if: 1) its state changed; or 2) one of its descendant's state
+     *   changed.
      *
      * E.g., "baseline-stateChange" to turn on the stateChange option with the baseline runtime.
      *
@@ -37,6 +41,12 @@ public class AndroidRuntimeConfigTools {
         "conflate" -> setOf(CONFLATE_STALE_RENDERINGS)
         "conflate-stateChange" -> setOf(CONFLATE_STALE_RENDERINGS, RENDER_ONLY_WHEN_STATE_CHANGES)
         "baseline-stateChange" -> setOf(RENDER_ONLY_WHEN_STATE_CHANGES)
+        "conflate-partial" -> setOf(
+          CONFLATE_STALE_RENDERINGS,
+          RENDER_ONLY_WHEN_STATE_CHANGES,
+          PARTIAL_TREE_RENDERING
+        )
+        "baseline-partial" -> setOf(RENDER_ONLY_WHEN_STATE_CHANGES, PARTIAL_TREE_RENDERING)
         "", "baseline" -> RuntimeConfigOptions.RENDER_PER_ACTION
         else ->
           throw IllegalArgumentException("Unrecognized config \"${BuildConfig.WORKFLOW_RUNTIME}\"")

workflow-config/config-jvm/src/main/java/com/squareup/workflow1/config/JvmTestRuntimeConfigTools.kt

@@ -3,6 +3,7 @@ package com.squareup.workflow1.config
 import com.squareup.workflow1.RuntimeConfig
 import com.squareup.workflow1.RuntimeConfigOptions
 import com.squareup.workflow1.RuntimeConfigOptions.CONFLATE_STALE_RENDERINGS
+import com.squareup.workflow1.RuntimeConfigOptions.PARTIAL_TREE_RENDERING
 import com.squareup.workflow1.RuntimeConfigOptions.RENDER_ONLY_WHEN_STATE_CHANGES
 import com.squareup.workflow1.WorkflowExperimentalRuntime
 
@@ -27,6 +28,9 @@ public class JvmTestRuntimeConfigTools {
      * Then, these can be combined (via '-') with:
      * "stateChange" : Only re-render when the state of some WorkflowNode has been changed by an
      *   action cascade.
+     * "partial" : Which includes "stateChange" as well as partial tree rendering, which only
+     *   re-renders each Workflow node if: 1) its state changed; or 2) one of its descendant's state
+     *   changed.
      *
      * E.g., "baseline-stateChange" to turn on the stateChange option with the baseline runtime.
      *
@@ -38,6 +42,12 @@ public class JvmTestRuntimeConfigTools {
         "conflate" -> setOf(CONFLATE_STALE_RENDERINGS)
         "conflate-stateChange" -> setOf(CONFLATE_STALE_RENDERINGS, RENDER_ONLY_WHEN_STATE_CHANGES)
         "baseline-stateChange" -> setOf(RENDER_ONLY_WHEN_STATE_CHANGES)
+        "conflate-partial" -> setOf(
+          CONFLATE_STALE_RENDERINGS,
+          RENDER_ONLY_WHEN_STATE_CHANGES,
+          PARTIAL_TREE_RENDERING
+        )
+        "baseline-partial" -> setOf(RENDER_ONLY_WHEN_STATE_CHANGES, PARTIAL_TREE_RENDERING)
         "", "baseline" -> RuntimeConfigOptions.RENDER_PER_ACTION
         else ->
           throw IllegalArgumentException("Unrecognized config \"$runtimeConfig\"")

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

@@ -38,6 +38,22 @@ public enum class RuntimeConfigOptions {
   @WorkflowExperimentalRuntime
   RENDER_ONLY_WHEN_STATE_CHANGES,
 
+  /**
+   * Only re-render each active Workflow node if:
+   * 1. It's own state changed, OR
+   * 2. One of it's descendant's state has changed.
+   *
+   * Otherwise return the cached rendering (as there is no way it could have changed).
+   *
+   * Note however that you must be careful using this because there may be external
+   * state that your Workflow's draw in and re-render and if that is not explicitly
+   * tracked within that Workflow's state, then it will not re-render. In this case,
+   * make sure that the state is tracked within the Workflow's state (even through
+   * an artificial token) in some way.
+   */
+  @WorkflowExperimentalRuntime
+  PARTIAL_TREE_RENDERING,
+
   /**
    * If we have more actions to process, do so before passing the rendering to the UI layer.
    */

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

@@ -6,11 +6,14 @@ 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.PARTIAL_TREE_RENDERING
+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
@@ -33,6 +36,7 @@ import kotlinx.coroutines.launch
 import kotlinx.coroutines.plus
 import kotlinx.coroutines.selects.SelectBuilder
 import kotlin.coroutines.CoroutineContext
+import kotlin.jvm.JvmInline
 
 /**
  * A node in a state machine tree. Manages the actual state for a given [Workflow].
@@ -44,7 +48,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 +90,11 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
   )
   private val sideEffects = ActiveStagingList<SideEffectNode>()
   private var lastProps: PropsT = initialProps
+  private var lastRendering: Box<RenderingT> = Box()
   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,
@@ -211,12 +217,14 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
    */
   private fun updateCachedWorkflowInstance(
     workflow: StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>
-  ) {
+  ): Boolean {
     if (workflow !== cachedWorkflowInstance) {
       // The instance has changed.
       cachedWorkflowInstance = workflow
       interceptedWorkflowInstance = interceptor.intercept(cachedWorkflowInstance, this)
+      return true
     }
+    return false
   }
 
   /**
@@ -227,39 +235,59 @@ 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(PARTIAL_TREE_RENDERING) ||
+      !lastRendering.isInitialized ||
+      subtreeStateDidChange
+    ) {
+      if (!didUpdateCachedInstance) {
+        // If we haven't already updated the cached instance, better do it now!
+        updateCachedWorkflowInstance(workflow)
+      }
+      baseRenderContext.unfreeze()
+      lastRendering = Box(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() }
+      }
+      // After we have rendered this subtree, we need another action in order for us to be
+      // considered dirty again.
+      subtreeStateDidChange = false
     }
 
-    return rendering
+    return lastRendering.getOrThrow()
   }
 
+  /**
+   * @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
   }
 
   /**
    * Applies [action] to this workflow's [state] and then passes the resulting [ActionApplied]
    * via [emitAppliedActionToParent] to the parent, with additional information as to whether or
    * not this action has changed the current node's state.
+   *
    */
   private fun applyAction(
     action: WorkflowAction<PropsT, StateT, OutputT>,
@@ -272,7 +300,13 @@ 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)
     )
-    return if (actionApplied.output != null) {
+    // Our state changed or one of our children's state changed.
+    subtreeStateDidChange = aggregateActionApplied.stateChanged
+    return if (actionApplied.output != null ||
+      runtimeConfig.contains(PARTIAL_TREE_RENDERING)
+    ) {
+      // If we are using the optimization, always return to the parent, so we carry a path that
+      // notes that the subtree did change all the way to the root.
       emitAppliedActionToParent(aggregateActionApplied)
     } else {
       aggregateActionApplied
@@ -289,4 +323,17 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
       SideEffectNode(key, job)
     }
   }
+
+  @JvmInline
+  internal value class Box<T>(private val _value: Any? = Uninitialized) {
+    val isInitialized: Boolean get() = _value !== Uninitialized
+
+    @Suppress("UNCHECKED_CAST")
+    fun getOrThrow(): T {
+      check(isInitialized)
+      return _value as T
+    }
+  }
+
+  internal object Uninitialized
 }