fix: propagating lifecycle events to AndroidScreenLifecycleOwner

voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/androidx/AndroidScreenLifecycleOwner.kt

@@ -53,11 +53,10 @@ public class AndroidScreenLifecycleOwner private constructor() :
     override val viewModelStore: ViewModelStore = ViewModelStore()
 
     private val atomicContext = AtomicReference<Context>()
+    internal val atomicParentLifecycleOwner = AtomicReference<LifecycleOwner>()
 
     private val controller = SavedStateRegistryController.create(this)
 
-    private var deactivateLifecycleListener: (() -> Unit)? = null
-
     private var isCreated: Boolean by mutableStateOf(false)
 
     override val savedStateRegistry: SavedStateRegistry
@@ -97,15 +96,13 @@ public class AndroidScreenLifecycleOwner private constructor() :
         }
     }
 
-    private fun onStart() {
+    private fun emitOnStartEvents() {
         startEvents.forEach {
             lifecycle.handleLifecycleEvent(it)
         }
     }
 
-    private fun onStop() {
-        deactivateLifecycleListener?.invoke()
-        deactivateLifecycleListener = null
+    private fun emitOnStopEvents() {
         stopEvents.forEach {
             lifecycle.handleLifecycleEvent(it)
         }
@@ -129,10 +126,11 @@ public class AndroidScreenLifecycleOwner private constructor() :
 
     override fun onDispose(screen: Screen) {
         val context = atomicContext.getAndSet(null) ?: return
-        if (context is Activity && context.isChangingConfigurations) return
+        val activity = context.getActivity()
+        if (activity != null && activity.isChangingConfigurations) return
         viewModelStore.clear()
-        disposeEvents.forEach {
-            lifecycle.handleLifecycleEvent(it)
+        disposeEvents.forEach { event ->
+            lifecycle.handleLifecycleEvent(event)
         }
     }
 
@@ -143,6 +141,7 @@ public class AndroidScreenLifecycleOwner private constructor() :
     @Composable
     private fun getHooks(): List<ProvidedValue<*>> {
         atomicContext.compareAndSet(null, LocalContext.current)
+        atomicParentLifecycleOwner.compareAndSet(null, LocalLifecycleOwner.current)
 
         return remember(this) {
             listOf(
@@ -153,17 +152,38 @@ public class AndroidScreenLifecycleOwner private constructor() :
         }
     }
 
-    private fun registerLifecycleListener(outState: Bundle) {
-        val activity = atomicContext.get()?.getActivity()
-        if (activity != null && activity is LifecycleOwner) {
+    /**
+     * Returns a unregister callback
+     */
+    private fun registerLifecycleListener(outState: Bundle): () -> Unit {
+        val lifecycleOwner = atomicParentLifecycleOwner.get()
+        if (lifecycleOwner != null) {
             val observer = object : DefaultLifecycleObserver {
+                override fun onPause(owner: LifecycleOwner) {
+                    lifecycle.handleLifecycleEvent(Lifecycle.Event.ON_PAUSE)
+                }
+
+                override fun onResume(owner: LifecycleOwner) {
+                    lifecycle.handleLifecycleEvent(Lifecycle.Event.ON_RESUME)
+                }
+
+                override fun onStart(owner: LifecycleOwner) {
+                    lifecycle.handleLifecycleEvent(Lifecycle.Event.ON_START)
+                }
+
                 override fun onStop(owner: LifecycleOwner) {
+                    lifecycle.handleLifecycleEvent(Lifecycle.Event.ON_STOP)
+
+                    // when the Application goes to background, perform save
                     performSave(outState)
                 }
             }
-            val lifecycle = activity.lifecycle
+            val lifecycle = lifecycleOwner.lifecycle
             lifecycle.addObserver(observer)
-            deactivateLifecycleListener = { lifecycle.removeObserver(observer) }
+
+            return { lifecycle.removeObserver(observer) }
+        } else {
+            return { }
         }
     }
 
@@ -175,11 +195,17 @@ public class AndroidScreenLifecycleOwner private constructor() :
         }
 
         DisposableEffect(this) {
-            registerLifecycleListener(savedState)
-            onStart()
+            val unregisterLifecycle = registerLifecycleListener(savedState)
+            emitOnStartEvents()
+
             onDispose {
+                unregisterLifecycle()
+
+                // when the screen goes to stack, perform save
                 performSave(savedState)
-                onStop()
+
+                // notify lifecycle screen listeners
+                emitOnStopEvents()
             }
         }
     }