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

This doesn't change any behavior, it is just more efficient: there's no need to
define the function at the tick pass, so this is not only more efficient but also
just easier to read, since there are fewer moving parts.

This is also required for #910.

Author: zach-klippenstein

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)
 }