Pass acceptOutput function to WorkflowNode constructor instead of every tick pass.

kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/SubtreeManager.kt

@@ -92,6 +92,7 @@ import kotlin.coroutines.CoroutineContext
  */
 internal class SubtreeManager<StateT, OutputT : Any>(
   private val contextForChildren: CoroutineContext,
+  private val emitActionToParent: (WorkflowAction<StateT, OutputT>) -> Any?,
   private val parentDiagnosticId: Long,
   private val diagnosticListener: WorkflowDiagnosticListener? = null,
   private val idCounter: IdCounter? = null
@@ -150,16 +151,9 @@ internal class SubtreeManager<StateT, OutputT : Any>(
    * Uses [selector] to invoke [WorkflowNode.tick] for every running child workflow this instance
    * is managing.
    */
-  fun <T : Any> tickChildren(
-    selector: SelectBuilder<T?>,
-    handler: (WorkflowAction<StateT, OutputT>) -> T?
-  ) {
+  fun <T : Any> tickChildren(selector: SelectBuilder<T?>) {
     children.forEachActive { child ->
-      child.workflowNode.tick(selector) { output ->
-        val componentUpdate = child.acceptChildOutput(output)
-        @Suppress("UNCHECKED_CAST")
-        return@tick handler(componentUpdate as WorkflowAction<StateT, OutputT>)
-      }
+      child.workflowNode.tick(selector)
     }
   }
 
@@ -187,16 +181,25 @@ internal class SubtreeManager<StateT, OutputT : Any>(
     handler: (ChildOutputT) -> WorkflowAction<StateT, OutputT>
   ): WorkflowChildNode<ChildPropsT, ChildOutputT, StateT, OutputT> {
     val id = child.id(key)
+    lateinit var node: WorkflowChildNode<ChildPropsT, ChildOutputT, StateT, OutputT>
+
+    fun acceptChildOutput(output: ChildOutputT): Any? {
+      val action = node.acceptChildOutput(output)
+      return emitActionToParent(action)
+    }
+
     val workflowNode = WorkflowNode(
         id,
         child.asStatefulWorkflow(),
         initialProps,
         snapshotCache[id],
         contextForChildren,
+        ::acceptChildOutput,
         parentDiagnosticId,
         diagnosticListener,
         idCounter
     )
     return WorkflowChildNode(child, handler, workflowNode)
+        .also { node = it }
   }
 }

kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/WorkflowLoop.kt

@@ -121,7 +121,7 @@ internal open class RealWorkflowLoop : WorkflowLoop {
             }
 
             // Tick the workflow tree.
-            rootNode.tick(this) { it }
+            rootNode.tick(this)
           }
         }
         // Compiler gets confused, and thinks both that this throw is unreachable, and without the

kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/WorkflowNode.kt

@@ -28,7 +28,6 @@ import com.squareup.workflow.diagnostic.createId
 import com.squareup.workflow.internal.RealRenderContext.WorkerRunner
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.InternalCoroutinesApi
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.channels.Channel
@@ -40,6 +39,8 @@ import kotlin.coroutines.CoroutineContext
 /**
  * A node in a state machine tree. Manages the actual state for a given [Workflow].
  *
+ * @param emitOutputToParent A function that this node will call when it needs to emit an output
+ * value to its parent. Returns either the output to be emitted from the root workflow, or null.
  * @param initialState Allows unit tests to start the node from a given state, instead of calling
  * [StatefulWorkflow.initialState].
  */
@@ -50,6 +51,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT : Any, RenderingT>(
   initialProps: PropsT,
   snapshot: ByteString?,
   baseContext: CoroutineContext,
+  private val emitOutputToParent: (OutputT) -> Any? = { it },
   parentDiagnosticId: Long? = null,
   private val diagnosticListener: WorkflowDiagnosticListener? = null,
   private val idCounter: IdCounter? = null,
@@ -69,8 +71,9 @@ internal class WorkflowNode<PropsT, StateT, OutputT : Any, RenderingT>(
    */
   internal val diagnosticId = idCounter.createId()
 
-  private val subtreeManager =
-    SubtreeManager<StateT, OutputT>(coroutineContext, diagnosticId, diagnosticListener, idCounter)
+  private val subtreeManager = SubtreeManager<StateT, OutputT>(
+      coroutineContext, ::applyAction, diagnosticId, diagnosticListener, idCounter
+  )
 
   private val workers = ActiveStagingList<WorkerChildNode<*, *, *>>()
 
@@ -158,20 +161,9 @@ internal class WorkflowNode<PropsT, StateT, OutputT : Any, RenderingT>(
    *
    * It is an error to call this method after calling [cancel].
    */
-  @UseExperimental(InternalCoroutinesApi::class)
-  fun <T : Any> tick(
-    selector: SelectBuilder<T?>,
-    handler: (OutputT) -> T?
-  ) {
-    fun acceptUpdate(action: WorkflowAction<StateT, OutputT>): T? {
-      val (newState, output) = action.applyTo(state)
-      diagnosticListener?.onWorkflowAction(diagnosticId, action, state, newState, output)
-      state = newState
-      return output?.let(handler)
-    }
-
+  fun <T : Any> tick(selector: SelectBuilder<T?>) {
     // Listen for any child workflow updates.
-    subtreeManager.tickChildren(selector, ::acceptUpdate)
+    subtreeManager.tickChildren(selector)
 
     // Listen for any subscription updates.
     workers.forEachActive { child ->
@@ -188,7 +180,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT : Any, RenderingT>(
           } else {
             val update = child.acceptUpdate(valueOrDone.value)
             @Suppress("UNCHECKED_CAST")
-            acceptUpdate(update as WorkflowAction<StateT, OutputT>)
+            return@onReceive applyAction(update as WorkflowAction<StateT, OutputT>)
           }
         }
       }
@@ -198,7 +190,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT : Any, RenderingT>(
     with(selector) {
       eventActionsChannel.onReceive { action ->
         diagnosticListener?.onSinkReceived(diagnosticId, action)
-        acceptUpdate(action)
+        return@onReceive applyAction(action)
       }
     }
   }
@@ -256,6 +248,18 @@ internal class WorkflowNode<PropsT, StateT, OutputT : Any, RenderingT>(
     lastProps = newProps
   }
 
+  /**
+   * Applies [action] to this workflow's [state] and
+   * [emits an output to its parent][emitOutputToParent] if necessary.
+   */
+  private fun <T : Any> applyAction(action: WorkflowAction<StateT, OutputT>): T? {
+    val (newState, output) = action.applyTo(state)
+    diagnosticListener?.onWorkflowAction(diagnosticId, action, state, newState, output)
+    state = newState
+    @Suppress("UNCHECKED_CAST")
+    return output?.let(emitOutputToParent) as T?
+  }
+
   private fun <T> createWorkerNode(
     worker: Worker<T>,
     key: String,

kotlin/workflow-runtime/src/test/java/com/squareup/workflow/internal/SubtreeManagerTest.kt

@@ -99,17 +99,15 @@ class SubtreeManagerTest {
   private val context = Unconfined
 
   @Test fun `render starts new child`() {
-    val manager =
-      SubtreeManager<String, String>(context, parentDiagnosticId = 0, diagnosticListener = null)
+    val manager = subtreeManagerForTest<String, String>()
     val workflow = TestWorkflow()
 
     manager.render(workflow, "props", key = "", handler = { fail() })
     assertEquals(1, workflow.started)
   }
 
   @Test fun `render doesn't start existing child`() {
-    val manager =
-      SubtreeManager<String, String>(context, parentDiagnosticId = 0, diagnosticListener = null)
+    val manager = subtreeManagerForTest<String, String>()
     val workflow = TestWorkflow()
     fun render() = manager.render(workflow, "props", key = "", handler = { fail() })
         .also { manager.commitRenderedChildren() }
@@ -121,8 +119,7 @@ class SubtreeManagerTest {
   }
 
   @Test fun `render restarts child after tearing down`() {
-    val manager =
-      SubtreeManager<String, String>(context, parentDiagnosticId = 0, diagnosticListener = null)
+    val manager = subtreeManagerForTest<String, String>()
     val workflow = TestWorkflow()
     fun render() = manager.render(workflow, "props", key = "", handler = { fail() })
         .also { manager.commitRenderedChildren() }
@@ -138,8 +135,7 @@ class SubtreeManagerTest {
   }
 
   @Test fun `render throws on duplicate key`() {
-    val manager =
-      SubtreeManager<String, String>(context, parentDiagnosticId = 0, diagnosticListener = null)
+    val manager = subtreeManagerForTest<String, String>()
     val workflow = TestWorkflow()
     manager.render(workflow, "props", "foo", handler = { fail() })
 
@@ -153,8 +149,7 @@ class SubtreeManagerTest {
   }
 
   @Test fun `render returns child rendering`() {
-    val manager =
-      SubtreeManager<String, String>(context, parentDiagnosticId = 0, diagnosticListener = null)
+    val manager = subtreeManagerForTest<String, String>()
     val workflow = TestWorkflow()
 
     val (composeProps, composeState) = manager.render(
@@ -165,8 +160,7 @@ class SubtreeManagerTest {
   }
 
   @Test fun `tick children handles child output`() {
-    val manager =
-      SubtreeManager<String, String>(context, parentDiagnosticId = 0, diagnosticListener = null)
+    val manager = subtreeManagerForTest<String, String>()
     val workflow = TestWorkflow()
     val handler: StringHandler = { output ->
       action { setOutput("case output:$output") }
@@ -189,8 +183,7 @@ class SubtreeManagerTest {
   }
 
   @Test fun `render updates child's output handler`() {
-    val manager =
-      SubtreeManager<String, String>(context, parentDiagnosticId = 0, diagnosticListener = null)
+    val manager = subtreeManagerForTest<String, String>()
     val workflow = TestWorkflow()
     fun render(handler: StringHandler) =
       manager.render(workflow, "props", key = "", handler = handler)
@@ -219,7 +212,7 @@ class SubtreeManagerTest {
 
   // See https://github.com/square/workflow/issues/404
   @Test fun `createChildSnapshot snapshots eagerly`() {
-    val manager = SubtreeManager<Unit, Nothing>(Unconfined, parentDiagnosticId = 0)
+    val manager = subtreeManagerForTest<Unit, Nothing>()
     val workflow = SnapshotTestWorkflow()
     assertEquals(0, workflow.snapshots)
 
@@ -232,7 +225,7 @@ class SubtreeManagerTest {
 
   // See https://github.com/square/workflow/issues/404
   @Test fun `createChildSnapshot serializes lazily`() {
-    val manager = SubtreeManager<Unit, Nothing>(Unconfined, parentDiagnosticId = 0)
+    val manager = subtreeManagerForTest<Unit, Nothing>()
     val workflow = SnapshotTestWorkflow()
     assertEquals(0, workflow.serializes)
 
@@ -246,11 +239,9 @@ class SubtreeManagerTest {
     assertEquals(1, workflow.serializes)
   }
 
-  private suspend fun <S, O : Any> SubtreeManager<S, O>.tickAction(): WorkflowAction<S, O>? {
-    return select {
-      tickChildren(this) { update ->
-        return@tickChildren update
-      }
-    }
-  }
+  private suspend fun <S, O : Any> SubtreeManager<S, O>.tickAction(): WorkflowAction<S, O>? =
+    select { tickChildren(this) }
+
+  private fun <S, O : Any> subtreeManagerForTest() =
+    SubtreeManager<S, O>(context, emitActionToParent = { it }, parentDiagnosticId = 0)
 }