From 953e2ab89b67999574ba5953255318a4d3bf8636 Mon Sep 17 00:00:00 2001
From: Keith Abdulla <ekeitho@users.noreply.github.com>
Date: Wed, 20 Nov 2024 13:35:26 -0800
Subject: [PATCH] ensure to keep a singule compose lifecycle owner and attach
 new parent lifecycles if new ones are provided

---
 .../ui/compose/ComposeLifecycleOwner.kt         | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/workflow-ui/compose/src/main/java/com/squareup/workflow1/ui/compose/ComposeLifecycleOwner.kt b/workflow-ui/compose/src/main/java/com/squareup/workflow1/ui/compose/ComposeLifecycleOwner.kt
index d3be64c63..0856028c6 100644
--- a/workflow-ui/compose/src/main/java/com/squareup/workflow1/ui/compose/ComposeLifecycleOwner.kt
+++ b/workflow-ui/compose/src/main/java/com/squareup/workflow1/ui/compose/ComposeLifecycleOwner.kt
@@ -18,8 +18,11 @@ import androidx.lifecycle.LifecycleRegistry
 @Composable internal fun rememberChildLifecycleOwner(
   parentLifecycle: Lifecycle = LocalLifecycleOwner.current.lifecycle
 ): LifecycleOwner {
-  val lifecycleOwner = remember {
-    ComposeLifecycleOwner.installOn(parentLifecycle)
+  val owner = remember { ComposeLifecycleOwner.installOn(parentLifecycle) }
+  val lifecycleOwner = remember(parentLifecycle) {
+    owner.apply {
+      updateParentLifecycle(parentLifecycle)
+    }
   }
   return lifecycleOwner
 }
@@ -56,9 +59,11 @@ import androidx.lifecycle.LifecycleRegistry
  * @param parentLifecycle The parent [Lifecycle] with which this lifecycle owner should synchronize.
  */
 private class ComposeLifecycleOwner(
-  private val parentLifecycle: Lifecycle
+  initialParentLifecycle: Lifecycle
 ) : LifecycleOwner, RememberObserver, LifecycleEventObserver {
 
+  private var parentLifecycle: Lifecycle = initialParentLifecycle
+
   private val registry = LifecycleRegistry(this)
   override val lifecycle: Lifecycle
     get() = registry
@@ -79,6 +84,12 @@ private class ComposeLifecycleOwner(
     }
   }
 
+  fun updateParentLifecycle(lifecycle: Lifecycle) {
+    parentLifecycle.removeObserver(this)
+    parentLifecycle = lifecycle
+    parentLifecycle.addObserver(this)
+  }
+
   override fun onStateChanged(
     source: LifecycleOwner,
     event: Event