From 79b24ba06042601d4caaadac5a69c71fb3aa6d71 Mon Sep 17 00:00:00 2001 From: Zach Klippenstein Date: Thu, 11 Sep 2025 10:18:40 -0700 Subject: [PATCH 1/2] Prototype of a stepper composable helper for making backstacks. --- .../com/squareup/workflow1/compose/Stepper.kt | 218 ++++++++++++++++++ .../squareup/workflow1/compose/StepperDemo.kt | 98 ++++++++ 2 files changed, 316 insertions(+) create mode 100644 workflow-ui/compose/src/main/java/com/squareup/workflow1/compose/Stepper.kt create mode 100644 workflow-ui/compose/src/main/java/com/squareup/workflow1/compose/StepperDemo.kt diff --git a/workflow-ui/compose/src/main/java/com/squareup/workflow1/compose/Stepper.kt b/workflow-ui/compose/src/main/java/com/squareup/workflow1/compose/Stepper.kt new file mode 100644 index 0000000000..ec46f3e2cc --- /dev/null +++ b/workflow-ui/compose/src/main/java/com/squareup/workflow1/compose/Stepper.kt @@ -0,0 +1,218 @@ +package com.squareup.workflow1.compose + +import androidx.collection.ScatterMap +import androidx.collection.mutableScatterMapOf +import androidx.compose.runtime.Composable +import androidx.compose.runtime.getValue +import androidx.compose.runtime.mutableStateListOf +import androidx.compose.runtime.mutableStateOf +import androidx.compose.runtime.remember +import androidx.compose.runtime.setValue +import androidx.compose.runtime.snapshots.Snapshot +import androidx.compose.runtime.snapshots.StateObject +import androidx.compose.runtime.snapshots.StateRecord +import androidx.compose.runtime.snapshots.readable +import androidx.compose.runtime.snapshots.writable + +/** + * Composes [content] and returns its return value in a list. + * + * Every time [content] calls [Stepper.advance] its argument is passed to [advance] and [advance] is + * expected to update some states that are read inside [content]. The current values of all states + * changed by [advance] are saved into a frame and pushed onto the backstack along with the last + * value returned by [content]. + * + * When [Stepper.goBack] is called the last frame is popped and all the states that were written by + * [advance] are restored before recomposing [content]. + * + * @sample com.squareup.workflow1.compose.StepperDemo + */ +@Composable +public fun stepper( + advance: (T) -> Unit, + content: @Composable Stepper.() -> R +): List { + // TODO figure out how to support rememberSaveable + val stepperImpl = remember { StepperImpl(advance = advance) } + stepperImpl.advance = advance + stepperImpl.lastRendering = content(stepperImpl) + return stepperImpl.renderings +} + +/** + * Composes [content] and returns its return value in a list. Every time [content] calls + * [Stepper.advance] the current values of all states changed by the `toState` block are + * saved into a frame and pushed onto the backstack along with the last value returned by [content]. + * When [Stepper.goBack] is called the last frame is popped and all the states that were + * written by the `toState` block are restored before recomposing [content]. + * + * This is an overload of [stepper] that makes it easier to specify the state update function when + * calling [Stepper.advance] instead of defining it ahead of time. + * + * @sample com.squareup.workflow1.compose.StepperInlineDemo + */ +// Impl note: Inline since this is just syntactic sugar, no reason to generate bytecode/API for it. +@Suppress("NOTHING_TO_INLINE") +@Composable +public inline fun stepper( + noinline content: @Composable Stepper<() -> Unit, R>.() -> R +): List = stepper(advance = { it() }, content = content) + +public interface Stepper { + + /** The (possibly empty) stack of steps that came before the current one. */ + val previousSteps: List> + + /** + * Pushes a new frame onto the backstack with the current state and then runs [toState]. + */ + fun advance(toState: T) + + /** + * Pops the last frame off the backstack and restores its state. + * + * @return False if the stack was empty (i.e. this is a noop). + */ + fun goBack(): Boolean +} + +public interface Step { + /** The last rendering produced by this step. */ + val rendering: T + + /** + * Runs [block] inside a snapshot such that the step state is set to its saved values from this + * step. The snapshot is read-only, so writing to any snapshot state objects will throw. + */ + fun peekStateFromStep(block: () -> R): R +} + +private class StepperImpl( + advance: (T) -> Unit +) : Stepper { + var advance: (T) -> Unit by mutableStateOf(advance) + private val savePoints = mutableStateListOf() + var lastRendering by mutableStateOf(NO_RENDERING) + + val renderings: List + get() = buildList(capacity = savePoints.size + 1) { + savePoints.mapTo(this) { it.rendering } + @Suppress("UNCHECKED_CAST") + add(lastRendering as R) + } + + override val previousSteps: List> + get() = savePoints + + override fun advance(toState: T) { + check(lastRendering !== NO_RENDERING) { "advance called before first composition" } + + // Take an outer snapshot so all the state mutations in withState get applied atomically with + // our internal state update (to savePoints). + Snapshot.withMutableSnapshot { + val savedRecords = mutableScatterMapOf() + val snapshot = Snapshot.takeMutableSnapshot( + writeObserver = { + // Don't save the value of the object yet, we want the value _before_ the write, so we + // need to read it outside this inner snapshot. + savedRecords[it as StateObject] = null + } + ) + try { + // Record what state objects are written by the block. + snapshot.enter { this.advance.invoke(toState) } + + // Save the _current_ values of those state objects so we can restore them later. + // TODO Need to think more about which state objects need to be saved and restored for a + // particular frame. E.g. probably we should track all objects that were written for the + // current frame, and save those as well, even if they're not written by the _next_ frame. + savedRecords.forEachKey { stateObject -> + savedRecords[stateObject] = stateObject.copyCurrentRecord() + } + + // This should never fail since we're already in a snapshot and no other state has been + // written by this point, but check just in case. + val advanceApplyResult = snapshot.apply() + if (advanceApplyResult.succeeded) { + // This cast is fine, we know we've assigned a non-null value to all entries. + @Suppress("UNCHECKED_CAST") + savePoints += SavePoint( + savedRecords = savedRecords as ScatterMap, + rendering = lastRendering as R, + ) + } + // If !succeeded, throw the standard error. + advanceApplyResult.check() + } finally { + snapshot.dispose() + } + } + } + + override fun goBack(): Boolean { + Snapshot.withMutableSnapshot { + if (savePoints.isEmpty()) return false + val toRestore = savePoints.removeAt(savePoints.lastIndex) + + // Restore all state objects' saved values. + toRestore.restoreState() + + // Don't need to restore the last rendering, it will be computed fresh by the imminent + // recomposition. + } + return true + } + + /** + * Returns a copy of the current readable record of this state object. A copy is needed since + * active records can be mutated by other snapshots. + */ + private fun StateObject.copyCurrentRecord(): StateRecord { + val record = firstStateRecord.readable(this) + // Records can be mutated in other snapshots, so create a copy. + return record.create().apply { assign(record) } + } + + /** + * Sets the value of this state object to a [record] that was previously copied via + * [copyCurrentRecord]. + */ + private fun StateObject.restoreRecord(record: StateRecord) { + firstStateRecord.writable(this) { assign(record) } + } + + private inner class SavePoint( + val savedRecords: ScatterMap, + override val rendering: R, + ) : Step { + override fun peekStateFromStep(block: () -> R): R { + // Need a mutable snapshot to restore state. + val restoreSnapshot = Snapshot.takeMutableSnapshot() + try { + restoreSnapshot.enter { + restoreState() + + // Now take a read-only snapshot to enforce contract. + val readOnlySnapshot = Snapshot.takeSnapshot() + try { + return readOnlySnapshot.enter(block) + } finally { + readOnlySnapshot.dispose() + } + } + } finally { + restoreSnapshot.dispose() + } + } + + fun restoreState() { + savedRecords.forEach { stateObject, record -> + stateObject.restoreRecord(record) + } + } + } + + companion object { + val NO_RENDERING = Any() + } +} diff --git a/workflow-ui/compose/src/main/java/com/squareup/workflow1/compose/StepperDemo.kt b/workflow-ui/compose/src/main/java/com/squareup/workflow1/compose/StepperDemo.kt new file mode 100644 index 0000000000..fb69a9403a --- /dev/null +++ b/workflow-ui/compose/src/main/java/com/squareup/workflow1/compose/StepperDemo.kt @@ -0,0 +1,98 @@ +package com.squareup.workflow1.compose + +import androidx.compose.runtime.Composable +import androidx.compose.runtime.getValue +import androidx.compose.runtime.mutableStateOf +import androidx.compose.runtime.saveable.rememberSaveable +import androidx.compose.runtime.setValue +import androidx.compose.ui.util.fastJoinToString +import com.squareup.workflow1.compose.DemoStep.ONE +import com.squareup.workflow1.compose.DemoStep.THREE +import com.squareup.workflow1.compose.DemoStep.TWO +import com.squareup.workflow1.compose.Screen.ScreenOne +import com.squareup.workflow1.compose.Screen.ScreenThree +import com.squareup.workflow1.compose.Screen.ScreenTwo + +internal enum class DemoStep { + ONE, + TWO, + THREE, +} + +internal sealed interface Screen { + val message: String + + data class ScreenOne( + override val message: String, + val onNextClicked: () -> Unit, + ) : Screen + + data class ScreenTwo( + override val message: String, + val onNextClicked: () -> Unit, + val onBack: () -> Unit, + ) : Screen + + data class ScreenThree( + override val message: String, + val onBack: () -> Unit, + ) : Screen +} + +@Composable +internal fun StepperDemo() { + var step by rememberSaveable { mutableStateOf(ONE) } + println("step=$step") + + val stack: List = stepper(advance = { step = it }) { + val breadcrumbs = previousSteps.fastJoinToString(separator = " > ") { it.rendering.message } + when (step) { + ONE -> ScreenOne( + message = "Step one", + onNextClicked = { advance(TWO) }, + ) + + TWO -> ScreenTwo( + message = "Step two", + onNextClicked = { advance(THREE) }, + onBack = { goBack() }, + ) + + THREE -> ScreenThree( + message = "Step three", + onBack = { goBack() }, + ) + } + } + + println("stack = ${stack.fastJoinToString()}") +} + +@Composable +internal fun StepperInlineDemo() { + var step by rememberSaveable { mutableStateOf(ONE) } + println("step=$step") + + val stack: List = stepper { + val breadcrumbs = previousSteps.fastJoinToString(separator = " > ") { it.rendering.message } + when (step) { + ONE -> ScreenOne( + message = "Step one", + onNextClicked = { advance { step = TWO } }, + ) + + TWO -> ScreenTwo( + message = "Step two", + onNextClicked = { advance { step = THREE } }, + onBack = { goBack() }, + ) + + THREE -> ScreenThree( + message = "Step three", + onBack = { goBack() }, + ) + } + } + + println("stack = ${stack.fastJoinToString()}") +} From 06f9c341498c59d4fec54ab6d68d752f2d4c3682 Mon Sep 17 00:00:00 2001 From: zach-klippenstein <101754+zach-klippenstein@users.noreply.github.com> Date: Thu, 11 Sep 2025 20:04:04 +0000 Subject: [PATCH 2/2] Apply changes from apiDump Signed-off-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- workflow-ui/compose/api/compose.api | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/workflow-ui/compose/api/compose.api b/workflow-ui/compose/api/compose.api index 5c5a794399..823e90b68a 100644 --- a/workflow-ui/compose/api/compose.api +++ b/workflow-ui/compose/api/compose.api @@ -1,3 +1,19 @@ +public abstract interface class com/squareup/workflow1/compose/Step { + public abstract fun getRendering ()Ljava/lang/Object; + public abstract fun peekStateFromStep (Lkotlin/jvm/functions/Function0;)Ljava/lang/Object; +} + +public abstract interface class com/squareup/workflow1/compose/Stepper { + public abstract fun advance (Ljava/lang/Object;)V + public abstract fun getPreviousSteps ()Ljava/util/List; + public abstract fun goBack ()Z +} + +public final class com/squareup/workflow1/compose/StepperKt { + public static final fun stepper (Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function3;Landroidx/compose/runtime/Composer;I)Ljava/util/List; + public static final fun stepper (Lkotlin/jvm/functions/Function3;Landroidx/compose/runtime/Composer;I)Ljava/util/List; +} + public final class com/squareup/workflow1/ui/compose/ComposableSingletons$ScreenComposableFactoryFinderKt { public static final field INSTANCE Lcom/squareup/workflow1/ui/compose/ComposableSingletons$ScreenComposableFactoryFinderKt; public fun ()V