Add WorkflowInterceptor API to intercept ComposeWorkflow render passes.

Author: zach-klippenstein

samples/hello-compose-workflow/src/main/java/com/squareup/sample/hellocomposeworkflow/HelloComposeWorkflow.kt

@@ -9,6 +9,7 @@ import com.squareup.sample.hellocomposeworkflow.HelloComposeWorkflow.State.Goodb
 import com.squareup.sample.hellocomposeworkflow.HelloComposeWorkflow.State.Hello
 import com.squareup.workflow1.WorkflowExperimentalApi
 import com.squareup.workflow1.compose.ComposeWorkflow
+import java.util.concurrent.atomic.AtomicInteger
 
 @OptIn(WorkflowExperimentalApi::class)
 object HelloComposeWorkflow : ComposeWorkflow<Unit, Nothing, HelloRendering>() {
@@ -23,7 +24,9 @@ object HelloComposeWorkflow : ComposeWorkflow<Unit, Nothing, HelloRendering>() {
     emitOutput: (Nothing) -> Unit
   ): HelloRendering {
     var state by remember { mutableStateOf(Hello) }
-    println("OMG recomposing state=$state")
+    val compositions = remember { AtomicInteger(0) }
+    println("OMG recomposing state=$state (count=${compositions.incrementAndGet()})")
+
     return HelloRendering(
       message = state.name,
       onClick = {

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

@@ -1,6 +1,7 @@
 package com.squareup.workflow1
 
 import com.squareup.workflow1.WorkflowInterceptor.RenderContextInterceptor
+import com.squareup.workflow1.WorkflowInterceptor.RuntimeLoopOutcome
 import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
@@ -110,6 +111,20 @@ public interface WorkflowInterceptor {
   ): RenderingAndSnapshot<R> =
     proceed(renderProps)
 
+  /**
+   * Called before a workflow is rendered, and before the interceptor's [onRender] method. Every
+   * call to this method is eventually followed by a call to [onAfterRender].
+   *
+   * This method should be preferred over [onRender] since it supports additional types of workflows
+   * that will not call [onRender].
+   */
+  public fun <P, S> onBeforeRender(
+    renderProps: P,
+    renderState: S,
+    session: WorkflowSession
+  ) {
+  }
+
   /**
    * Intercepts calls to [StatefulWorkflow.render].
    */
@@ -121,6 +136,15 @@ public interface WorkflowInterceptor {
     session: WorkflowSession
   ): R = proceed(renderProps, renderState, null)
 
+  /**
+   * Called after a workflow is rendered, and after the interceptor's [onRender] method. Every
+   * call to this method is preceded by a call to [onBeforeRender].
+   *
+   * This method should be preferred over [onRender] since it supports additional types of workflows
+   * that will not call [onRender].
+   */
+  public fun <R> onAfterRender(rendering: R) {}
+
   /**
    * Intercept calls to [StatefulWorkflow.snapshotState] including the children calls.
    * This is useful to intercept a rendering + snapshot traversal for tracing purposes.
@@ -375,29 +399,35 @@ internal fun <P, S, O, R> WorkflowInterceptor.intercept(
       renderProps: P,
       renderState: S,
       context: RenderContext<P, S, O>
-    ): R = onRender(
-      renderProps,
-      renderState,
-      context,
-      proceed = { props, state, interceptor ->
-        // 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,
-    )
+    ): R {
+      onBeforeRender(renderProps, renderState, workflowSession)
+      return onRender(
+        renderProps,
+        renderState,
+        context,
+        proceed = { props, state, interceptor ->
+          // 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,
+      ).also {
+        onAfterRender(it)
+      }
+    }
 
     override fun snapshotState(state: S) =
       onSnapshotState(state, workflow::snapshotState, workflowSession)

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

@@ -75,6 +75,16 @@ internal class ChainedWorkflowInterceptor(
     return chainedProceed(renderProps)
   }
 
+  override fun <P, S> onBeforeRender(
+    renderProps: P,
+    renderState: S,
+    session: WorkflowSession
+  ) {
+    for (i in interceptors.indices) {
+      interceptors[i].onBeforeRender(renderProps, renderState, session)
+    }
+  }
+
   override fun <P, S, O, R> onRender(
     renderProps: P,
     renderState: S,
@@ -99,6 +109,13 @@ internal class ChainedWorkflowInterceptor(
     return chainedProceed(renderProps, renderState, null)
   }
 
+  override fun <R> onAfterRender(rendering: R) {
+    // TODO does this still give optimized bytecode?
+    for (i in interceptors.indices.reversed()) {
+      interceptors[i].onAfterRender(rendering)
+    }
+  }
+
   override fun onSnapshotStateWithChildren(
     proceed: () -> TreeSnapshot,
     session: WorkflowSession

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

@@ -77,7 +77,8 @@ internal class ComposeWorkflowNode<PropsT, OutputT, RenderingT>(
     }
   }
 
-  private var cachedComposeWorkflow by mutableStateOf(workflow)
+  private var cachedComposeWorkflow: ComposeWorkflow<PropsT, OutputT, RenderingT>? by
+  mutableStateOf(null)
   private var lastProps by mutableStateOf(initialProps)
   private var lastRendering = NullableInitBox<RenderingT>()
   private val recomposer: Recomposer = Recomposer(coroutineContext)
@@ -196,6 +197,29 @@ internal class ComposeWorkflowNode<PropsT, OutputT, RenderingT>(
       session = this,
       proceed = { _, _, _ -> ComposeWorkflowState }
     )
+
+    // By not calling setContent directly every time, we ensure that if neither the workflow
+    // instance nor input changed, we don't recompose.
+    // setContent will synchronously perform the first recomposition before returning, which is why
+    // we leave cachedComposeWorkflow null for now: we don't want its produceRendering to be called
+    // until we're actually doing a render pass.
+    composition.setContent {
+      @Suppress("NAME_SHADOWING")
+      val workflow = cachedComposeWorkflow
+      if (workflow != null) {
+        val rendering = workflow.produceRendering(
+          props = lastProps,
+          emitOutput = sendOutputToChannel
+        )
+
+        // lastRendering isn't snapshot state, so wait until the composition is applied to update
+        // it.
+        SideEffect {
+          lastRendering = NullableInitBox(rendering)
+        }
+      }
+    }
+    cachedComposeWorkflow = workflow
   }
 
   override fun render(
@@ -221,6 +245,11 @@ internal class ComposeWorkflowNode<PropsT, OutputT, RenderingT>(
         }
       )
 
+      interceptor.onBeforeRender(
+        renderProps = input,
+        renderState = ComposeWorkflowState,
+        session = this
+      )
       interceptor.onRender(
         session = this,
         renderProps = input,
@@ -233,7 +262,9 @@ internal class ComposeWorkflowNode<PropsT, OutputT, RenderingT>(
           //  But we need to be able to support sending actions to support this anyway.
           doRender()
         }
-      )
+      ).also {
+        interceptor.onAfterRender(it)
+      }
     } else {
       this.lastProps = input
       doRender()
@@ -242,22 +273,6 @@ internal class ComposeWorkflowNode<PropsT, OutputT, RenderingT>(
 
   private fun doRender(): RenderingT {
     val frameRequest = if (!lastRendering.isInitialized) {
-      // By not calling setContent directly every time, we ensure that if neither the workflow
-      // instance nor input changed, we don't recompose.
-      // setContent will synchronously perform the first recomposition before returning!
-      composition.setContent {
-        val rendering = cachedComposeWorkflow.produceRendering(
-          props = lastProps,
-          emitOutput = sendOutputToChannel
-        )
-
-        // lastRendering isn't snapshot state, so wait until the composition is applied to update
-        // it.
-        SideEffect {
-          lastRendering = NullableInitBox(rendering)
-        }
-      }
-
       // Initial render kicks off the render loop. This should always synchronously request a frame.
       startComposition()
 
@@ -297,8 +312,16 @@ internal class ComposeWorkflowNode<PropsT, OutputT, RenderingT>(
   }
 
   override fun snapshot(): TreeSnapshot {
-    // TODO Support snapshots from rememberSaveable.
-    return TreeSnapshot(workflowSnapshot = null, childTreeSnapshots = ::emptyMap)
+    return interceptor.onSnapshotStateWithChildren(
+      session = this,
+      proceed = {
+        // Compose workflows do not support the onSnapshotState interceptor since they don't
+        // distinguish between snapshot state objects for themselves and their child
+        // ComposeWorkflows.
+        // TODO Support snapshots from rememberSaveable.
+        TreeSnapshot(workflowSnapshot = null, childTreeSnapshots = ::emptyMap)
+      }
+    )
   }
 
   @OptIn(ExperimentalCoroutinesApi::class, DelicateCoroutinesApi::class)