Introduce renderComposable.

samples/compose-samples/src/main/java/com/squareup/sample/compose/nestedrenderings/NestedRenderingsActivity.kt

@@ -10,6 +10,7 @@ import androidx.compose.ui.graphics.Color
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
+import com.squareup.workflow1.SimpleLoggingWorkflowInterceptor
 import com.squareup.workflow1.WorkflowExperimentalRuntime
 import com.squareup.workflow1.config.AndroidRuntimeConfigTools
 import com.squareup.workflow1.mapRendering
@@ -47,7 +48,8 @@ class NestedRenderingsActivity : AppCompatActivity() {
         workflow = RecursiveWorkflow.mapRendering { it.withEnvironment(viewEnvironment) },
         scope = viewModelScope,
         savedStateHandle = savedState,
-        runtimeConfig = AndroidRuntimeConfigTools.getAppWorkflowRuntimeConfig()
+        runtimeConfig = AndroidRuntimeConfigTools.getAppWorkflowRuntimeConfig(),
+        interceptors = listOf(SimpleLoggingWorkflowInterceptor())
       )
     }
   }

settings.gradle.kts

@@ -7,6 +7,7 @@ pluginManagement {
     google()
     // For binary compatibility validator.
     maven { url = uri("https://kotlin.bintray.com/kotlinx") }
+    maven("https://maven.pkg.jetbrains.space/public/p/compose/dev")
   }
   includeBuild("build-logic")
 }

workflow-core/build.gradle.kts

@@ -3,6 +3,7 @@ import com.squareup.workflow1.buildsrc.iosWithSimulatorArm64
 plugins {
   id("kotlin-multiplatform")
   id("published")
+  id("org.jetbrains.compose") version "1.7.3"
 }
 
 kotlin {
@@ -23,6 +24,8 @@ dependencies {
   commonMainApi(libs.kotlinx.coroutines.core)
   // For Snapshot.
   commonMainApi(libs.squareup.okio)
+  commonMainApi("org.jetbrains.compose.runtime:runtime:1.7.3")
+  commonMainApi("org.jetbrains.compose.runtime:runtime-saveable:1.7.3")
 
   commonTestImplementation(libs.kotlinx.atomicfu)
   commonTestImplementation(libs.kotlinx.coroutines.test.common)

workflow-core/src/commonMain/kotlin/com/squareup/workflow1/BaseRenderContext.kt

@@ -8,7 +8,9 @@
 
 package com.squareup.workflow1
 
+import androidx.compose.runtime.Composable
 import com.squareup.workflow1.WorkflowAction.Companion.noAction
+import com.squareup.workflow1.compose.WorkflowComposable
 import kotlinx.coroutines.CoroutineScope
 import kotlin.jvm.JvmMultifileClass
 import kotlin.jvm.JvmName
@@ -86,6 +88,27 @@ public interface BaseRenderContext<PropsT, StateT, OutputT> {
     handler: (ChildOutputT) -> WorkflowAction<PropsT, StateT, OutputT>
   ): ChildRenderingT
 
+  // /**
+  //  * Synchronously composes a [content] function and returns its rendering. Whenever [content] is
+  //  * invalidated (i.e. a compose snapshot state object is changed that was previously read by
+  //  * [content] or any functions it calls), this workflow will be re-rendered and the relevant
+  //  * composables will be recomposed.
+  //  *
+  //  * The `emitOutput` function passed to [content] should be used to trigger [WorkflowAction]s in
+  //  * this workflow via [handler]. Every invocation of `emitOutput` will result [handler]s action
+  //  * being sent to this context's [actionSink]. However, it's important for the composable never to
+  //  * send to [actionSink] directly because we need to ensure that any state writes the composable
+  //  * does invalidate their composables before sending into the [actionSink].
+  //  */
+  // @WorkflowExperimentalApi
+  // public fun <ChildOutputT, ChildRenderingT> renderComposable(
+  //   key: String = "",
+  //   handler: (ChildOutputT) -> WorkflowAction<PropsT, StateT, OutputT>,
+  //   content: @WorkflowComposable @Composable (
+  //     emitOutput: (ChildOutputT) -> Unit
+  //   ) -> ChildRenderingT
+  // ): ChildRenderingT
+
   /**
    * Ensures [sideEffect] is running with the given [key].
    *
@@ -209,6 +232,20 @@ public fun <PropsT, StateT, OutputT, ChildRenderingT>
   key: String = ""
 ): ChildRenderingT = renderChild(child, Unit, key) { noAction() }
 
+/**
+ * TODO
+ */
+@WorkflowExperimentalApi
+public fun <PropsT, StateT, OutputT, ChildRenderingT>
+  BaseRenderContext<PropsT, StateT, OutputT>.renderComposable(
+  key: String = "",
+  content: @WorkflowComposable @Composable () -> ChildRenderingT
+): ChildRenderingT = renderComposable<Nothing, ChildRenderingT>(
+  key = key,
+  handler = { noAction() },
+  content = { content() }
+)
+
 /**
  * Ensures a [LifecycleWorker] is running. Since [worker] can't emit anything,
  * it can't trigger any [WorkflowAction]s.

workflow-core/src/commonMain/kotlin/com/squareup/workflow1/compose/WorkflowComposable.kt

@@ -0,0 +1,25 @@
+package com.squareup.workflow1.compose
+
+import androidx.compose.runtime.ComposableTargetMarker
+import com.squareup.workflow1.WorkflowExperimentalApi
+import kotlin.annotation.AnnotationRetention.BINARY
+import kotlin.annotation.AnnotationTarget.FILE
+import kotlin.annotation.AnnotationTarget.FUNCTION
+import kotlin.annotation.AnnotationTarget.PROPERTY_GETTER
+import kotlin.annotation.AnnotationTarget.TYPE
+import kotlin.annotation.AnnotationTarget.TYPE_PARAMETER
+
+/**
+ * An annotation that can be used to mark a composable function as being expected to be use in a
+ * composable function that is also marked or inferred to be marked as a [WorkflowComposable], i.e.
+ * that can be called from [BaseRenderContext.renderComposable].
+ *
+ * Using this annotation explicitly is rarely necessary as the Compose compiler plugin will infer
+ * the necessary equivalent annotations automatically. See
+ * [androidx.compose.runtime.ComposableTarget] for details.
+ */
+@WorkflowExperimentalApi
+@ComposableTargetMarker(description = "Workflow Composable")
+@Target(FILE, FUNCTION, PROPERTY_GETTER, TYPE, TYPE_PARAMETER)
+@Retention(BINARY)
+public annotation class WorkflowComposable

workflow-runtime/build.gradle.kts

@@ -3,6 +3,7 @@ import com.squareup.workflow1.buildsrc.iosWithSimulatorArm64
 plugins {
   id("kotlin-multiplatform")
   id("published")
+  id("org.jetbrains.compose") version "1.7.3"
 }
 
 kotlin {
@@ -16,6 +17,13 @@ kotlin {
   if (targets == "kmp" || targets == "js") {
     js(IR) { browser() }
   }
+  // sourceSets {
+  //   getByName("commonMain") {
+  //     dependencies {
+  //       implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.8.1")
+  //     }
+  //   }
+  // }
 }
 
 dependencies {

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

@@ -1,5 +1,6 @@
 package com.squareup.workflow1
 
+import androidx.compose.runtime.Composable
 import com.squareup.workflow1.WorkflowInterceptor.RenderContextInterceptor
 import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
 import kotlinx.coroutines.CoroutineScope
@@ -150,5 +151,15 @@ public open class SimpleLoggingWorkflowInterceptor : WorkflowInterceptor {
         }
       }
     }
+
+    override fun <CR> onRenderComposable(
+      key: String,
+      content: @Composable () -> CR,
+      proceed: (key: String, content: @Composable () -> CR) -> CR
+    ): CR = proceed(key) {
+      logMethod("onRenderComposable", session, "key" to key, "content" to content) {
+        content()
+      }
+    }
   }
 }

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

@@ -1,7 +1,9 @@
 package com.squareup.workflow1
 
+import androidx.compose.runtime.Composable
 import com.squareup.workflow1.WorkflowInterceptor.RenderContextInterceptor
 import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
+import com.squareup.workflow1.compose.WorkflowComposable
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
 import kotlin.coroutines.CoroutineContext
@@ -322,6 +324,15 @@ public interface WorkflowInterceptor {
         calculation: () -> CResult
       ) -> CResult
     ): CResult = proceed(key, resultType, inputs, calculation)
+
+    public fun <CO, CR> onRenderComposable(
+      key: String,
+      content: @Composable (CO) -> CR,
+      proceed: (
+        key: String,
+        content: @Composable (CO) -> CR
+      ) -> CR
+    ): CR = proceed(key, content)
   }
 }
 
@@ -459,6 +470,23 @@ private class InterceptedRenderContext<P, S, O>(
     }
   }
 
+  @OptIn(WorkflowExperimentalApi::class)
+  override fun <ChildOutputT, ChildRenderingT> renderComposable(
+    key: String,
+    handler: (ChildOutputT) -> WorkflowAction<P, S, O>,
+    content: @WorkflowComposable @Composable (emitOutput: (ChildOutputT) -> Unit) -> ChildRenderingT
+  ): ChildRenderingT = interceptor.onRenderComposable(
+    key = key,
+    content = content,
+    proceed = { iKey, iContent ->
+      baseRenderContext.renderComposable(
+        key = iKey,
+        handler = handler,
+        content = iContent
+      )
+    }
+  )
+
   /**
    * In a block with a CoroutineScope receiver, calls to `coroutineContext` bind
    * to `CoroutineScope.coroutineContext` instead of `suspend val coroutineContext`.

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

@@ -1,11 +1,16 @@
+@file:OptIn(WorkflowExperimentalApi::class)
+
 package com.squareup.workflow1.internal
 
+import androidx.compose.runtime.Composable
 import com.squareup.workflow1.BaseRenderContext
 import com.squareup.workflow1.RuntimeConfig
 import com.squareup.workflow1.Sink
 import com.squareup.workflow1.Workflow
 import com.squareup.workflow1.WorkflowAction
+import com.squareup.workflow1.WorkflowExperimentalApi
 import com.squareup.workflow1.WorkflowTracer
+import com.squareup.workflow1.compose.WorkflowComposable
 import com.squareup.workflow1.identifier
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.channels.SendChannel
@@ -27,6 +32,12 @@ internal class RealRenderContext<PropsT, StateT, OutputT>(
       key: String,
       handler: (ChildOutputT) -> WorkflowAction<PropsT, StateT, OutputT>
     ): ChildRenderingT
+
+    fun <ChildOutputT, ChildRenderingT> renderComposable(
+      key: String,
+      handler: (ChildOutputT) -> WorkflowAction<PropsT, StateT, OutputT>,
+      content: @Composable (emitOutput: (ChildOutputT) -> Unit) -> ChildRenderingT
+    ): ChildRenderingT
   }
 
   interface SideEffectRunner {
@@ -78,6 +89,15 @@ internal class RealRenderContext<PropsT, StateT, OutputT>(
     return renderer.render(child, props, key, handler)
   }
 
+  override fun <ChildOutputT, ChildRenderingT> renderComposable(
+    key: String,
+    handler: (ChildOutputT) -> WorkflowAction<PropsT, StateT, OutputT>,
+    content: @WorkflowComposable @Composable (emitOutput: (ChildOutputT) -> Unit) -> ChildRenderingT
+  ): ChildRenderingT {
+    checkNotFrozen()
+    return renderer.renderComposable(key, handler, content)
+  }
+
   override fun runningSideEffect(
     key: String,
     sideEffect: suspend CoroutineScope.() -> Unit

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

@@ -0,0 +1,13 @@
+package com.squareup.workflow1.internal
+
+import com.squareup.workflow1.WorkflowAction
+
+/**
+ * This action doesn't actually update state, but it's special-cased inside WorkflowNode to always
+ * act like it updated state, to force a re-render and thus a recomposition.
+ */
+internal class RecomposeAction<PropsT, StateT, OutputT> : WorkflowAction<PropsT, StateT, OutputT>() {
+  override fun Updater.apply() {
+    // Noop
+  }
+}

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

@@ -1,5 +1,6 @@
 package com.squareup.workflow1.internal
 
+import androidx.compose.runtime.Composable
 import com.squareup.workflow1.ActionApplied
 import com.squareup.workflow1.ActionProcessingResult
 import com.squareup.workflow1.NoopWorkflowInterceptor
@@ -90,15 +91,18 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
   private val contextForChildren: CoroutineContext,
   private val emitActionToParent: (
     action: WorkflowAction<PropsT, StateT, OutputT>,
-    childResult: ActionApplied<*>
+    childResult: ActionApplied<*>?
   ) -> ActionProcessingResult,
   private val runtimeConfig: RuntimeConfig,
   private val workflowTracer: WorkflowTracer?,
   private val workflowSession: WorkflowSession? = null,
   private val interceptor: WorkflowInterceptor = NoopWorkflowInterceptor,
-  private val idCounter: IdCounter? = null
+  private val idCounter: IdCounter? = null,
+  private val requestRerender: () -> Unit = {},
+  private val sendActionFromComposable: (WorkflowAction<PropsT,StateT,OutputT>) -> Unit
 ) : RealRenderContext.Renderer<PropsT, StateT, OutputT> {
   private var children = ActiveStagingList<WorkflowChildNode<*, *, *, *, *>>()
+  private var composables = ActiveStagingList<WorkflowComposableNode<*, *, *, *, *>>()
 
   /**
    * Moves all the nodes that have been accumulated in the staging list to the active list, making
@@ -112,6 +116,7 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
     children.commitStaging { child ->
       child.workflowNode.cancel()
     }
+    composables.commitStaging(onRemove = WorkflowComposableNode<*, *, *, *, *>::dispose)
     // Get rid of any snapshots that weren't applied on the first render pass.
     // They belong to children that were saved but not restarted.
     snapshotCache = null
@@ -145,6 +150,30 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
     return stagedChild.render(child.asStatefulWorkflow(), props)
   }
 
+  override fun <ChildOutputT, ChildRenderingT> renderComposable(
+    key: String,
+    handler: (ChildOutputT) -> WorkflowAction<PropsT, StateT, OutputT>,
+    content: @Composable (emitOutput: (ChildOutputT) -> Unit) -> ChildRenderingT
+  ): ChildRenderingT {
+    // Prevent duplicate workflows with the same key.
+    workflowTracer.trace("CheckingUniqueMatchesComposable") {
+      composables.forEachStaging {
+        require(key != it.workflowKey) {
+          "Expected keys to be unique for composable: key=\"$key\""
+        }
+      }
+    }
+
+    val stagedComposable = workflowTracer.trace("RetainingComposables") {
+      composables.retainOrCreate(
+        predicate = { it.workflowKey == key },
+        create = { createComposableNode<ChildOutputT, ChildRenderingT>(key, handler) }
+      )
+    }
+    stagedComposable.setHandler(handler)
+    return stagedComposable.render(content)
+  }
+
   /**
    * Uses [selector] to invoke [WorkflowNode.onNextAction] for every running child workflow this instance
    * is managing.
@@ -165,6 +194,7 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
     val snapshots = mutableMapOf<WorkflowNodeId, TreeSnapshot>()
     children.forEachActive { child ->
       val childWorkflow = child.workflow.asStatefulWorkflow()
+      // Skip children who aren't snapshottable.
       snapshots[child.id] = child.workflowNode.snapshot(childWorkflow)
     }
     return snapshots
@@ -206,4 +236,19 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
     return WorkflowChildNode(child, handler, workflowNode)
       .also { node = it }
   }
+
+  private fun <ChildOutputT, ChildRenderingT> createComposableNode(
+    key: String,
+    handler: (ChildOutputT) -> WorkflowAction<PropsT, StateT, OutputT>,
+  ): WorkflowComposableNode<ChildOutputT, ChildRenderingT, PropsT, StateT, OutputT> {
+    return WorkflowComposableNode<ChildOutputT, ChildRenderingT, PropsT, StateT, OutputT>(
+      workflowKey = key,
+      handler = handler,
+      coroutineContext = contextForChildren,
+      requestRerender = requestRerender,
+      sendAction = sendActionFromComposable,
+    ).also {
+      it.start()
+    }
+  }
 }