feat: Multiple ScreenLifecycleOwner support for Screen

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

@@ -219,7 +219,9 @@ public class AndroidScreenLifecycleOwner private constructor() :
         private val disposeEvents = arrayOf(
             Lifecycle.Event.ON_DESTROY
         )
-        public fun get(screen: Screen): ScreenLifecycleOwner =
-            ScreenLifecycleStore.get(screen) { AndroidScreenLifecycleOwner() }
+        public fun get(screen: Screen): ScreenLifecycleOwner {
+            return ScreenLifecycleStore.register(screen) { AndroidScreenLifecycleOwner() }
+        }
+
     }
 }

voyager-core/src/commonMain/kotlin/cafe/adriel/voyager/core/lifecycle/ScreenLifecycleOwner.kt

@@ -1,87 +1,29 @@
 package cafe.adriel.voyager.core.lifecycle
 
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.remember
-import cafe.adriel.voyager.core.annotation.ExperimentalVoyagerApi
 import cafe.adriel.voyager.core.annotation.InternalVoyagerApi
 import cafe.adriel.voyager.core.screen.Screen
 
 public interface ScreenLifecycleOwner {
+
+    /**
+     * Called before rendering the Screen Content.
+     *
+     * IMPORTANT: This is only called when ScreenLifecycleOwner is provided by [ScreenLifecycleProvider] or [NavigatorScreenLifecycleProvider].
+     */
     @Composable
     public fun ProvideBeforeScreenContent(
         provideSaveableState: @Composable (suffixKey: String, content: @Composable () -> Unit) -> Unit,
         content: @Composable () -> Unit
     ): Unit = content()
 
+    /**
+     * Called on the Screen leaves the stack.
+     */
     public fun onDispose(screen: Screen) {}
 }
 
 @InternalVoyagerApi
 public object DefaultScreenLifecycleOwner : ScreenLifecycleOwner
 
-@ExperimentalVoyagerApi
-@InternalVoyagerApi
-public class ComposedScreenLifecycleOwner(
-    private val screenLifecycleOwners: List<ScreenLifecycleOwner>
-) : ScreenLifecycleOwner {
-
-    @Composable
-    public fun RecursiveProvideBeforeScreenContent(
-        screenLifecycleOwner: ScreenLifecycleOwner,
-        provideSaveableState: @Composable (suffixKey: String, content: @Composable () -> Unit) -> Unit,
-        content: @Composable () -> Unit,
-        nextOrNull: () -> ScreenLifecycleOwner?,
-    ) {
-        val next = remember(screenLifecycleOwner, provideSaveableState, content, nextOrNull) { nextOrNull() }
-        if(next != null) {
-            val recursiveContent = @Composable {
-                RecursiveProvideBeforeScreenContent(
-                    screenLifecycleOwner = next,
-                    provideSaveableState = provideSaveableState,
-                    content = content,
-                    nextOrNull = nextOrNull,
-                )
-            }
-            screenLifecycleOwner.ProvideBeforeScreenContent(
-                provideSaveableState = { suffixKey, _ ->
-                    provideSaveableState(suffixKey, recursiveContent)
-                }
-            ) {
-                recursiveContent()
-            }
-
-        } else {
-            screenLifecycleOwner.ProvideBeforeScreenContent(
-                provideSaveableState = { suffixKey, content ->
-                    provideSaveableState(suffixKey, content)
-                }
-            ) {
-                content()
-            }
-        }
-    }
 
-    @Composable
-    override fun ProvideBeforeScreenContent(
-        provideSaveableState: @Composable (suffixKey: String, content: @Composable () -> Unit) -> Unit,
-        content: @Composable () -> Unit
-    ) {
-        if(screenLifecycleOwners.isNotEmpty()) {
-            val copy = screenLifecycleOwners.toMutableList()
-            RecursiveProvideBeforeScreenContent(
-                screenLifecycleOwner = copy.removeFirst(),
-                provideSaveableState = provideSaveableState,
-                content = content,
-                nextOrNull = { copy.removeFirstOrNull() }
-            )
-        } else {
-            content()
-        }
-    }
-
-    override fun onDispose(screen: Screen) {
-        for (screenLifecycleOwner in screenLifecycleOwners) {
-            screenLifecycleOwner.onDispose(screen)
-        }
-    }
-}

voyager-core/src/commonMain/kotlin/cafe/adriel/voyager/core/lifecycle/ScreenLifecycleStore.kt

@@ -3,18 +3,52 @@ package cafe.adriel.voyager.core.lifecycle
 import cafe.adriel.voyager.core.concurrent.ThreadSafeMap
 import cafe.adriel.voyager.core.screen.Screen
 import cafe.adriel.voyager.core.screen.ScreenKey
+import kotlin.reflect.KType
+import kotlin.reflect.typeOf
 
 public object ScreenLifecycleStore {
 
     private val owners = ThreadSafeMap<ScreenKey, ScreenLifecycleOwner>()
+    private val newOwners = ThreadSafeMap<ScreenKey, ThreadSafeMap<KType, ScreenLifecycleOwner>>()
 
+    @Deprecated(
+        message = "Use `register` instead. Will be removed in 1.0.0.",
+        replaceWith = ReplaceWith("ScreenLifecycleStore.register<T>(screen, factory)"),
+    )
     public fun get(
         screen: Screen,
         factory: (ScreenKey) -> ScreenLifecycleOwner
     ): ScreenLifecycleOwner =
         owners.getOrPut(screen.key) { factory(screen.key) }
 
+    /**
+     * Register a ScreenLifecycleOwner that will be called `onDispose` on the
+     * [screen] leaves the Navigation stack.
+     */
+    public inline fun <reified T : ScreenLifecycleOwner> register(
+        screen: Screen,
+        noinline factory: (ScreenKey) -> T,
+    ): T {
+        return register(screen, typeOf<T>(), factory) as T
+    }
+
+    @PublishedApi
+    internal fun <T : ScreenLifecycleOwner> register(
+        screen: Screen,
+        screenLifecycleOwnerType: KType,
+        factory: (ScreenKey) -> T,
+    ): ScreenLifecycleOwner {
+        return newOwners.getOrPut(screen.key) {
+            ThreadSafeMap<KType, ScreenLifecycleOwner>().apply {
+                put(screenLifecycleOwnerType, factory(screen.key))
+            }
+        }.getOrPut(screenLifecycleOwnerType) {
+            factory(screen.key)
+        }
+    }
+
     public fun remove(screen: Screen) {
         owners.remove(screen.key)?.onDispose(screen)
+        newOwners.remove(screen.key)?.forEach { it.value.onDispose(screen) }
     }
 }

voyager-core/src/commonMain/kotlin/cafe/adriel/voyager/core/lifecycle/multipleScreenLifecycleOwnerUtil.kt

@@ -0,0 +1,63 @@
+package cafe.adriel.voyager.core.lifecycle
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.remember
+import cafe.adriel.voyager.core.annotation.ExperimentalVoyagerApi
+import cafe.adriel.voyager.core.annotation.InternalVoyagerApi
+
+@ExperimentalVoyagerApi
+@InternalVoyagerApi
+@Composable
+public fun MultipleProvideBeforeScreenContent(
+    screenLifecycleOwners: List<ScreenLifecycleOwner>,
+    provideSaveableState: @Composable (suffixKey: String, content: @Composable () -> Unit) -> Unit,
+    content: @Composable () -> Unit,
+) {
+    if(screenLifecycleOwners.isNotEmpty()) {
+        val copy = screenLifecycleOwners.toMutableList()
+        RecursiveProvideBeforeScreenContent(
+            screenLifecycleOwner = copy.removeFirst(),
+            provideSaveableState = provideSaveableState,
+            content = content,
+            nextOrNull = { copy.removeFirstOrNull() }
+        )
+    } else {
+        content()
+    }
+}
+
+@Composable
+private fun RecursiveProvideBeforeScreenContent(
+    screenLifecycleOwner: ScreenLifecycleOwner,
+    provideSaveableState: @Composable (suffixKey: String, content: @Composable () -> Unit) -> Unit,
+    content: @Composable () -> Unit,
+    nextOrNull: () -> ScreenLifecycleOwner?,
+) {
+    val next = remember(screenLifecycleOwner, provideSaveableState, content, nextOrNull) { nextOrNull() }
+    if(next != null) {
+        val recursiveContent = @Composable {
+            RecursiveProvideBeforeScreenContent(
+                screenLifecycleOwner = next,
+                provideSaveableState = provideSaveableState,
+                content = content,
+                nextOrNull = nextOrNull,
+            )
+        }
+        screenLifecycleOwner.ProvideBeforeScreenContent(
+            provideSaveableState = { suffixKey, _ ->
+                provideSaveableState(suffixKey, recursiveContent)
+            }
+        ) {
+            recursiveContent()
+        }
+
+    } else {
+        screenLifecycleOwner.ProvideBeforeScreenContent(
+            provideSaveableState = { suffixKey, content ->
+                provideSaveableState(suffixKey, content)
+            }
+        ) {
+            content()
+        }
+    }
+}

voyager-navigator/src/commonMain/kotlin/cafe/adriel/voyager/navigator/Navigator.kt

@@ -5,12 +5,13 @@ import androidx.compose.runtime.CompositionLocalProvider
 import androidx.compose.runtime.ProvidableCompositionLocal
 import androidx.compose.runtime.derivedStateOf
 import androidx.compose.runtime.getValue
+import androidx.compose.runtime.remember
 import androidx.compose.runtime.saveable.SaveableStateHolder
 import androidx.compose.runtime.saveable.rememberSaveableStateHolder
 import androidx.compose.runtime.staticCompositionLocalOf
 import cafe.adriel.voyager.core.annotation.InternalVoyagerApi
 import cafe.adriel.voyager.core.concurrent.ThreadSafeSet
-import cafe.adriel.voyager.core.lifecycle.ComposedScreenLifecycleOwner
+import cafe.adriel.voyager.core.lifecycle.MultipleProvideBeforeScreenContent
 import cafe.adriel.voyager.core.lifecycle.ScreenLifecycleStore
 import cafe.adriel.voyager.core.lifecycle.getNavigatorScreenLifecycleOwner
 import cafe.adriel.voyager.core.lifecycle.rememberScreenLifecycleOwner
@@ -135,8 +136,11 @@ public class Navigator @InternalVoyagerApi constructor(
         val lifecycleOwner = rememberScreenLifecycleOwner(screen)
         val navigatorScreenLifecycleOwners = getNavigatorScreenLifecycleOwner(screen)
 
-        val composed = ComposedScreenLifecycleOwner(listOf(lifecycleOwner) + navigatorScreenLifecycleOwners)
-        composed.ProvideBeforeScreenContent(
+        val composed = remember(lifecycleOwner, navigatorScreenLifecycleOwners) {
+            listOf(lifecycleOwner) + navigatorScreenLifecycleOwners
+        }
+        MultipleProvideBeforeScreenContent(
+            screenLifecycleOwners = composed,
             provideSaveableState = { suffix, content -> provideSaveableState(suffix, content) },
             content = {
                 stateHolder.SaveableStateProvider(stateKey, content)