Refactor LifecycleOwner to ComposeLifecycleOwner for better lifecycle management

- Extracted anonymous LifecycleOwner and RememberObserver implementation into a reusable ComposeLifecycleOwner class.
- ComposeLifecycleOwner synchronizes its lifecycle with the parent Lifecycle and integrates with Compose's memory model via RememberObserver.
- Improves code readability, reusability, and aligns lifecycle management with Compose best practices.

Author: ekeitho

workflow-ui/compose/src/main/java/com/squareup/workflow1/ui/compose/ComposeLifecycleOwner.kt

@@ -0,0 +1,75 @@
+package com.squareup.workflow1.ui.compose
+
+import androidx.compose.runtime.RememberObserver
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleEventObserver
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.LifecycleRegistry
+
+/**
+ * A custom [LifecycleOwner] that synchronizes its lifecycle with a parent [Lifecycle] and
+ * integrates with Jetpack Compose's lifecycle through [RememberObserver].
+ *
+ * **Purpose:**
+ *
+ * - Ensures that any lifecycle-aware components within a composable function have a lifecycle that
+ *   accurately reflects both the parent lifecycle and the composable's own lifecycle.
+ * - Manages lifecycle transitions and observer registration/removal to prevent memory leaks and
+ *   ensure proper cleanup when the composable leaves the composition.
+ *
+ * **Key Features:**
+ *
+ * - **Lifecycle Synchronization:** Mirrors lifecycle events from the provided `parentLifecycle` to
+ *   its own [LifecycleRegistry], ensuring consistent state transitions.
+ * - **Compose Integration:** Implements [RememberObserver] to align with the composable's lifecycle
+ *   in the Compose memory model.
+ * - **Automatic Observer Management:** Adds and removes a [LifecycleEventObserver] to the parent
+ *   lifecycle, preventing leaks and ensuring proper disposal.
+ * - **State Transition Safety:** Carefully manages lifecycle state changes to avoid illegal
+ *   transitions, especially during destruction.
+ *
+ * **Usage Notes:**
+ *
+ * - Should be used in conjunction with `remember` and provided the `parentLifecycle` as a key to
+ *   ensure it updates correctly when the parent lifecycle changes.
+ * - By integrating with Compose's lifecycle, it ensures that resources are properly released when
+ *   the composable leaves the composition.
+ *
+ * @param parentLifecycle The parent [Lifecycle] with which this lifecycle owner should synchronize.
+ */
+internal class ComposeLifecycleOwner(
+  private val parentLifecycle: Lifecycle
+) : LifecycleOwner, RememberObserver {
+
+  private val registry = LifecycleRegistry(this)
+  override val lifecycle: Lifecycle
+    get() = registry
+
+  private val parentObserver = LifecycleEventObserver { _, event ->
+    // Any time the parent lifecycle changes state, perform the same change on our lifecycle.
+    registry.handleLifecycleEvent(event)
+  }
+
+  init {
+    // We do this right away to ensure that the parent lifecycle is observed as soon as possible
+    // to prevent any inconsistencies.
+    parentLifecycle.addObserver(parentObserver)
+  }
+
+  override fun onRemembered() {
+
+  }
+
+  override fun onAbandoned() {
+    onForgotten()
+  }
+
+  override fun onForgotten() {
+    parentLifecycle.removeObserver(parentObserver)
+
+    // If we're leaving the composition, ensure the lifecycle is cleaned up
+    if (registry.currentState != Lifecycle.State.INITIALIZED) {
+      registry.currentState = Lifecycle.State.DESTROYED
+    }
+  }
+}

workflow-ui/compose/src/main/java/com/squareup/workflow1/ui/compose/WorkflowRendering.kt

@@ -3,17 +3,11 @@ package com.squareup.workflow1.ui.compose
 import androidx.compose.foundation.layout.Box
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
-import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.key
 import androidx.compose.runtime.remember
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.platform.LocalLifecycleOwner
-import androidx.lifecycle.Lifecycle
-import androidx.lifecycle.Lifecycle.State.DESTROYED
-import androidx.lifecycle.Lifecycle.State.INITIALIZED
-import androidx.lifecycle.LifecycleEventObserver
 import androidx.lifecycle.LifecycleOwner
-import androidx.lifecycle.LifecycleRegistry
 import com.squareup.workflow1.ui.Compatible
 import com.squareup.workflow1.ui.Screen
 import com.squareup.workflow1.ui.ScreenViewHolder
@@ -101,34 +95,9 @@ public fun WorkflowRendering(
  * bit simpler since we don't need to worry about attachment state.
  */
 @Composable private fun rememberChildLifecycleOwner(): LifecycleOwner {
-  val lifecycleOwner = remember {
-    object : LifecycleOwner {
-      val registry = LifecycleRegistry(this)
-      override val lifecycle: Lifecycle
-        get() = registry
-    }
-  }
   val parentLifecycle = LocalLifecycleOwner.current.lifecycle
-
-  DisposableEffect(parentLifecycle) {
-    val parentObserver = LifecycleEventObserver { _, event ->
-      // Any time the parent lifecycle changes state, perform the same change on our lifecycle.
-      lifecycleOwner.registry.handleLifecycleEvent(event)
-    }
-
-    parentLifecycle.addObserver(parentObserver)
-    onDispose {
-      parentLifecycle.removeObserver(parentObserver)
-
-      // If we're leaving the composition it means the WorkflowRendering is either going away itself
-      // or about to switch to an incompatible rendering – either way, this lifecycle is dead. Note
-      // that we can't transition from INITIALIZED to DESTROYED – the LifecycleRegistry will throw.
-      // WorkflowLifecycleOwner has this same check.
-      if (lifecycleOwner.registry.currentState != INITIALIZED) {
-        lifecycleOwner.registry.currentState = DESTROYED
-      }
-    }
+  val lifecycleOwner = remember(parentLifecycle) {
+    ComposeLifecycleOwner(parentLifecycle)
   }
-
   return lifecycleOwner
 }