Suspend nodes when required

core/src/main/kotlin/com/bumble/appyx/core/children/ChildNodeCreationManager.kt

@@ -16,10 +16,16 @@ import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.flow.updateAndGet
 import kotlinx.coroutines.launch
 
+/**
+ * Initializes and removes nodes based on parent node routing source.
+ *
+ * Lifecycle of these nodes is managed in [com.bumble.appyx.core.lifecycle.ChildNodeLifecycleManager].
+ */
 internal class ChildNodeCreationManager<Routing : Any>(
     private var savedStateMap: SavedStateMap?,
     private val customisations: NodeCustomisationDirectory,
     private val childMode: ChildEntry.ChildMode,
+    private val keepMode: ChildEntry.KeepMode,
 ) {
     private val _children =
         MutableStateFlow<Map<RoutingKey<Routing>, ChildEntry<Routing>>>(emptyMap())
@@ -32,29 +38,80 @@ internal class ChildNodeCreationManager<Routing : Any>(
             _children.update { restoredMap }
             savedStateMap = null
         }
+        syncNavModelWithChildren(parentNode)
+    }
+
+    private fun syncNavModelWithChildren(parentNode: ParentNode<Routing>) {
         parentNode.lifecycle.coroutineScope.launch {
-            parentNode.navModel.elements.collect { elements ->
-                _children.update { map ->
-                    val navModelKeys = elements
-                        .mapTo(HashSet(elements.size, 1f)) { element -> element.key }
-                    val localKeys = map.keys
-                    val newKeys = navModelKeys - localKeys
-                    val removedKeys = localKeys - navModelKeys
-                    val mutableMap = map.toMutableMap()
-                    newKeys.forEach { key ->
-                        mutableMap[key] =
-                            childEntry(
-                                key = key,
-                                savedState = null,
-                                suspended = childMode == ChildEntry.ChildMode.LAZY,
-                            )
+            parentNode.navModel.screenState.collect { state ->
+                val navModelOnScreenKeys: Set<RoutingKey<Routing>>
+                val navModelOffScreenKeys: Set<RoutingKey<Routing>>
+                val navModelKeys: Set<RoutingKey<Routing>>
+                when (keepMode) {
+                    ChildEntry.KeepMode.KEEP -> {
+                        // Consider everything as on-screen for keep mode
+                        navModelOnScreenKeys =
+                            (state.onScreen + state.offScreen).mapNotNullToSet { element -> element.key }
+                        navModelOffScreenKeys = emptySet()
+                        navModelKeys = navModelOnScreenKeys
                     }
-                    removedKeys.forEach { key ->
-                        mutableMap.remove(key)
+                    ChildEntry.KeepMode.SUSPEND -> {
+                        navModelOnScreenKeys =
+                            state.onScreen.mapNotNullToSet { element -> element.key }
+                        navModelOffScreenKeys =
+                            state.offScreen.mapNotNullToSet { element -> element.key }
+                        navModelKeys = navModelOnScreenKeys + navModelOffScreenKeys
                     }
-                    mutableMap
                 }
+                updateChildren(navModelKeys, navModelOnScreenKeys, navModelOffScreenKeys)
+            }
+        }
+    }
+
+    // TODO: Does not work with CHILD_MODE = LAZY
+    private fun updateChildren(
+        navModelKeys: Set<RoutingKey<Routing>>,
+        navModelOnScreenKeys: Set<RoutingKey<Routing>>,
+        navModelOffScreenKeys: Set<RoutingKey<Routing>>,
+    ) {
+        _children.update { map ->
+            val localKeys = map.keys
+            val localOnScreenKeys = map.entries.mapNotNullToSet { entry ->
+                entry.key.takeIf { entry.value is ChildEntry.Initialized }
+            }
+            val localOffScreenKeys = map.entries.mapNotNullToSet { entry ->
+                entry.key.takeIf { entry.value is ChildEntry.Suspended }
+            }
+            val newKeys = navModelKeys - localKeys
+            val removedKeys = localKeys - navModelKeys
+            val offToOnScreenKeys = localOffScreenKeys.intersect(navModelOnScreenKeys)
+            val onToOffScreenKeys = localOnScreenKeys.intersect(navModelOffScreenKeys)
+            val noKeysChanges = newKeys.isEmpty() && removedKeys.isEmpty()
+            val noScreenChanges = offToOnScreenKeys.isEmpty() && onToOffScreenKeys.isEmpty()
+            if (noKeysChanges && noScreenChanges) {
+                return@update map
+            }
+            val mutableMap = map.toMutableMap()
+            newKeys.forEach { key ->
+                val shouldSuspend = keepMode == ChildEntry.KeepMode.SUSPEND &&
+                        navModelOffScreenKeys.contains(key)
+                mutableMap[key] =
+                    childEntry(
+                        key = key,
+                        savedState = null,
+                        suspended = childMode == ChildEntry.ChildMode.LAZY || shouldSuspend,
+                    )
+            }
+            removedKeys.forEach { key ->
+                mutableMap.remove(key)
             }
+            offToOnScreenKeys.forEach { key ->
+                mutableMap[key] = requireNotNull(mutableMap[key]).initialize()
+            }
+            onToOffScreenKeys.forEach { key ->
+                mutableMap[key] = requireNotNull(mutableMap[key]).suspend()
+            }
+            mutableMap
         }
     }
 
@@ -74,16 +131,8 @@ internal class ChildNodeCreationManager<Routing : Any>(
                         ?: error("Requested child $routingKey disappeared")
                     when (updateChild) {
                         is ChildEntry.Initialized -> map
-                        is ChildEntry.Suspended -> {
-                            val initialized = ChildEntry.Initialized(
-                                key = updateChild.key,
-                                node = parentNode.resolve(
-                                    routing = updateChild.key.routing,
-                                    buildContext = childBuildContext(child.savedState),
-                                ).build()
-                            )
-                            map.plus(routingKey to initialized)
-                        }
+                        is ChildEntry.Suspended ->
+                            map.plus(routingKey to updateChild.initialize())
                     }
                 }[routingKey] as ChildEntry.Initialized
         }
@@ -134,8 +183,35 @@ internal class ChildNodeCreationManager<Routing : Any>(
             )
         }
 
+    private fun ChildEntry<Routing>.initialize(): ChildEntry.Initialized<Routing> =
+        when (this) {
+            is ChildEntry.Initialized -> this
+            is ChildEntry.Suspended ->
+                ChildEntry.Initialized(
+                    key = key,
+                    node = parentNode.resolve(
+                        routing = key.routing,
+                        buildContext = childBuildContext(savedState),
+                    ).build()
+                )
+        }
+
+    private fun ChildEntry<Routing>.suspend(): ChildEntry.Suspended<Routing> =
+        when (this) {
+            is ChildEntry.Suspended -> this
+            is ChildEntry.Initialized ->
+                ChildEntry.Suspended(
+                    key = key,
+                    // TODO: Not able to get a scope from Compose here, providing fake one
+                    savedState = node.saveInstanceState { true },
+                )
+        }
+
     private companion object {
         const val KEY_CHILDREN_STATE = "ChildrenState"
+
+        private fun <T, R : Any> Collection<T>.mapNotNullToSet(mapper: (T) -> R?): Set<R> =
+            mapNotNullTo(HashSet(size, 1f), mapper)
     }
 
 }

core/src/main/kotlin/com/bumble/appyx/core/navigation/onscreen/OnScreenStateResolver.kt

@@ -1,6 +1,6 @@
 package com.bumble.appyx.core.navigation.onscreen
 
-interface OnScreenStateResolver<State> {
+fun interface OnScreenStateResolver<State> {
 
     fun isOnScreen(state: State): Boolean
 }

core/src/main/kotlin/com/bumble/appyx/core/node/ParentNode.kt

@@ -48,6 +48,7 @@ abstract class ParentNode<Routing : Any>(
     buildContext: BuildContext,
     view: ParentNodeView<Routing> = EmptyParentNodeView(),
     childMode: ChildEntry.ChildMode = ChildEntry.ChildMode.EAGER,
+    keepMode: ChildEntry.KeepMode = ChildEntry.KeepMode.KEEP,
     private val childAware: ChildAware<ParentNode<Routing>> = ChildAwareImpl(),
     plugins: List<Plugin> = listOf(),
 ) : Node(
@@ -66,6 +67,7 @@ abstract class ParentNode<Routing : Any>(
         savedStateMap = buildContext.savedStateMap,
         customisations = buildContext.customisations,
         childMode = childMode,
+        keepMode = keepMode,
     )
     val children: StateFlow<ChildEntryMap<Routing>>
         get() = childNodeCreationManager.children