tunjid · tunjid · Nov 18, 2024 · Nov 6, 2024 · Nov 14, 2024 · Nov 17, 2024
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
@@ -12,11 +12,11 @@ androidxTestExt = "1.2.1"
 androidxTestRunner = "1.6.2"
 androidxTestRules = "1.6.1"
 dokka = "1.8.20"
-jetbrainsCompose = "1.7.0"
-jetbrainsLifecycle = "2.8.3"
+jetbrainsCompose = "1.7.1"
+jetbrainsLifecycle = "2.8.4"
 jetbrainsMaterial3Adaptive = "1.0.0"
 junit4 = "4.13.2"
-kotlin = "2.0.20"
+kotlin = "2.0.21"
 kotlinxCoroutines = "1.9.0"
 kotlinxDatetime = "0.6.1"
 lifecycle-runtime = "2.8.6"

diff --git a/...y/compose/src/commonMain/kotlin/com/tunjid/treenav/compose/SavedStatePanedNavHostState.kt b/...y/compose/src/commonMain/kotlin/com/tunjid/treenav/compose/SavedStatePanedNavHostState.kt
@@ -20,6 +20,9 @@ import androidx.compose.runtime.remember
 import androidx.compose.runtime.saveable.SaveableStateHolder
 import androidx.compose.runtime.saveable.rememberSaveableStateHolder
 import androidx.compose.runtime.setValue
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.ViewModelStoreOwner
 import androidx.lifecycle.compose.LocalLifecycleOwner
 import androidx.lifecycle.compose.currentStateAsState
 import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
@@ -64,8 +67,22 @@ interface PanedNavHostScope<Pane, Destination : Node> {
 }
 
 /**
- * An implementation of an [PanedNavHostState] that provides a [SaveableStateHolder] for each
- * navigation destination that shows up in its panes.
+ * An implementation of an [PanedNavHostState] that provides the following for each
+ * navigation [Destination] that shows up in its panes:
+ *
+ * - A single [SaveableStateHolder] for each navigation [Destination] that shows up in its panes.
+ * [SaveableStateHolder.SaveableStateProvider] is keyed on the [Destination]s [Node.id].
+ *
+ * - A [ViewModelStoreOwner] for each [Destination] via [LocalViewModelStoreOwner].
+ * Once present in the navigation tree, a [Destination] will always use the same
+ * [ViewModelStoreOwner], regardless of where in the tree it is, until its is removed from the tree.
+ * [Destination]s are unique based on their [Node.id].
+ *
+ * - A [LifecycleOwner] for each [Destination] via [LocalLifecycleOwner]. This [LifecycleOwner]
+ * follows the [Lifecycle] of its immediate parent, unless it is animating out or placed in the
+ * backstack. This is defined by [PaneScope.isActive], which is a function of the backing
+ * [AnimatedContent] for each [Pane] displayed and if the current [Destination]
+ * matches [PanedNavHostScope.nodeFor] in the visible [Pane].
  *
  * @param panes a list of panes that is possible to show in the [PanedNavHost] in all
  * possible configurations. The panes should consist of enum class instances, or a sealed class
@@ -86,7 +103,7 @@ class SavedStatePanedNavHostState<Pane, Destination : Node>(
         val saveableStateHolder = rememberSaveableStateHolder()
 
         val panedContentScope = remember {
-            SavedStatePanedNavHostScope(
+            NavHostScope(
                 panes = panes,
                 navHostConfiguration = configuration,
                 initialPanesToNodes = panesToNodes,
@@ -106,23 +123,19 @@ class SavedStatePanedNavHostState<Pane, Destination : Node>(
 
     companion object {
         @Stable
-        private class SavedStatePanedNavHostScope<Pane, Destination : Node>(
+        class NavHostScope<Pane, Destination : Node> internal constructor(
             panes: List<Pane>,
             initialPanesToNodes: Map<Pane, Destination?>,
             saveableStateHolder: SaveableStateHolder,
             val navHostConfiguration: PanedNavHostConfiguration<Pane, *, Destination>,
         ) : PanedNavHostScope<Pane, Destination>, SaveableStateHolder by saveableStateHolder {
 
-            private val destinationViewModelStoreCreator = DestinationViewModelStoreCreator(
-                rootNodeProvider = navHostConfiguration.navigationState::value
-            )
-
-            val slots = List(
+            private val slots = List(
                 size = panes.size,
                 init = ::Slot
             ).toSet()
 
-            var panedNavigationState by mutableStateOf(
+            private var panedNavigationState by mutableStateOf(
                 value = SlotBasedPanedNavigationState.initial<Pane, Destination>(slots = slots)
                     .adaptTo(
                         slots = slots,
@@ -131,6 +144,10 @@ class SavedStatePanedNavHostState<Pane, Destination : Node>(
                     )
             )
 
+            private val destinationViewModelStoreCreator = DestinationViewModelStoreCreator(
+                validNodeIdsReader = { panedNavigationState.backStackIds + panedNavigationState.destinationIdsAnimatingOut }
+            )
+
             private val slotsToRoutes =
                 mutableStateMapOf<Slot?, @Composable () -> Unit>().also { map ->
                     map[null] = {}
@@ -139,6 +156,19 @@ class SavedStatePanedNavHostState<Pane, Destination : Node>(
                     }
                 }
 
+            /**
+             * Retrieves the a [ViewModelStoreOwner] for a given [destination]. All destinations
+             * with the same [Node.id] share the same [ViewModelStoreOwner].
+             *
+             * The [destination] must be present in the navigation tree, otherwise an
+             * [IllegalStateException] will be thrown.
+             *
+             * @param destination The destination for which the [ViewModelStoreOwner] should
+             * be retrieved.
+             */
+            fun viewModelStoreOwnerFor(destination: Destination): ViewModelStoreOwner =
+                destinationViewModelStoreCreator.viewModelStoreOwnerFor(destination)
+
             @Composable
             override fun Destination(pane: Pane) {
                 val slot = panedNavigationState.slotFor(pane)
@@ -153,7 +183,7 @@ class SavedStatePanedNavHostState<Pane, Destination : Node>(
                 pane: Pane
             ): Destination? = panedNavigationState.destinationFor(pane)
 
-            fun onNewNavigationState(
+            internal fun onNewNavigationState(
                 navigationState: Node,
                 panesToNodes: Map<Pane, Destination?>,
             ) {
@@ -208,8 +238,10 @@ class SavedStatePanedNavHostState<Pane, Destination : Node>(
                         val destinationLifecycleOwner = rememberDestinationLifecycleOwner(
                             destination
                         )
-                        val destinationViewModelOwner = destinationViewModelStoreCreator
-                            .viewModelStoreOwnerFor(destination)
+                        val destinationViewModelOwner = remember(destination.id) {
+                            destinationViewModelStoreCreator
+                                .viewModelStoreOwnerFor(destination)
+                        }
 
                         CompositionLocalProvider(
                             LocalLifecycleOwner provides destinationLifecycleOwner,
@@ -286,3 +318,13 @@ class SavedStatePanedNavHostState<Pane, Destination : Node>(
             }
     }
 }
+
+fun <Pane, Destination : Node> PanedNavHostScope<
+        Pane,
+        Destination
+        >.requireSavedStatePanedNavHostScope(): SavedStatePanedNavHostState.Companion.NavHostScope<Pane, Destination> {
+    check(this is SavedStatePanedNavHostState.Companion.NavHostScope) {
+        "This PanedNavHostScope instance is not a SavedStatePanedNavHostScope"
+    }
+    return this
+}
diff --git a/...ommonMain/kotlin/com/tunjid/treenav/compose/lifecycle/DestinationViewModelStoreCreator.kt b/...ommonMain/kotlin/com/tunjid/treenav/compose/lifecycle/DestinationViewModelStoreCreator.kt
@@ -17,42 +17,54 @@
 package com.tunjid.treenav.compose.lifecycle
 
 import androidx.compose.runtime.Stable
+import androidx.compose.runtime.mutableStateMapOf
 import androidx.lifecycle.ViewModelStore
 import androidx.lifecycle.ViewModelStoreOwner
 import com.tunjid.treenav.Node
 import com.tunjid.treenav.Order
-import com.tunjid.treenav.flatten
+import com.tunjid.treenav.traverse
 
+/**
+ * A class that lazily loads a [ViewModelStoreOwner] for each destination in the navigation graph.
+ * Each unique destination can only have a single [ViewModelStore], regardless of how many times
+ * it appears in the navigation graph, or its depth at any point.
+ */
 @Stable
 internal class DestinationViewModelStoreCreator(
-    private val rootNodeProvider: () -> Node
+    private val validNodeIdsReader: () -> Set<String>
 ) {
-    private val nodeIdsToViewModelStoreOwner = mutableMapOf<String, ViewModelStoreOwner>()
+    private val nodeIdsToViewModelStoreOwner = mutableStateMapOf<String, ViewModelStoreOwner>()
 
     /**
      * Creates a [ViewModelStoreOwner] for a given [Node]
      */
     fun viewModelStoreOwnerFor(
         node: Node
-    ): ViewModelStoreOwner = nodeIdsToViewModelStoreOwner.getOrPut(
-        node.id
-    ) {
-        object : ViewModelStoreOwner {
-            override val viewModelStore: ViewModelStore = ViewModelStore()
+    ): ViewModelStoreOwner {
+        val existingIds = validNodeIdsReader()
+        check(existingIds.contains(node.id)) {
+            """
+                Attempted to create a ViewModelStoreOwner for Node $node, but the Node is not
+                present in the navigation tree
+            """.trimIndent()
+        }
+        return nodeIdsToViewModelStoreOwner.getOrPut(
+            node.id
+        ) {
+            object : ViewModelStoreOwner {
+                override val viewModelStore: ViewModelStore = ViewModelStore()
+            }
         }
     }
 
     fun clearStoreFor(childNode: Node) {
-        val rootNode = rootNodeProvider()
-        val existingNodeIds = rootNode.flatten(Order.BreadthFirst).mapTo(
-            destination = mutableSetOf(),
-            transform = Node::id
-        )
-        if (existingNodeIds.contains(childNode.id)) {
-            return
+        val existingIds = validNodeIdsReader()
+        childNode.traverse(Order.BreadthFirst) {
+            if (!existingIds.contains(it.id)) {
+                nodeIdsToViewModelStoreOwner.remove(it.id)
+                    ?.viewModelStore
+                    ?.clear()
+            }
         }
-        nodeIdsToViewModelStoreOwner.remove(childNode.id)
-            ?.viewModelStore
-            ?.clear()
     }
 }
diff --git a/...Main/kotlin/com/tunjid/treenav/compose/moveablesharedelement/MovableSharedElementState.kt b/...Main/kotlin/com/tunjid/treenav/compose/moveablesharedelement/MovableSharedElementState.kt
@@ -1,6 +1,5 @@
 package com.tunjid.treenav.compose.moveablesharedelement
 
-import androidx.compose.animation.BoundsTransform
 import androidx.compose.animation.ExperimentalSharedTransitionApi
 import androidx.compose.animation.SharedTransitionScope
 import androidx.compose.runtime.Composable
@@ -9,78 +8,29 @@ import androidx.compose.runtime.Stable
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.movableContentOf
 import androidx.compose.runtime.mutableIntStateOf
-import androidx.compose.runtime.mutableStateMapOf
 import androidx.compose.runtime.mutableStateOf
-import androidx.compose.runtime.produceState
 import androidx.compose.runtime.setValue
-import androidx.compose.runtime.snapshotFlow
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.draw.drawWithContent
-import androidx.compose.ui.graphics.drawscope.ContentDrawScope
-import androidx.compose.ui.graphics.drawscope.translate
-import androidx.compose.ui.graphics.layer.GraphicsLayer
-import androidx.compose.ui.graphics.layer.drawLayer
-import androidx.compose.ui.graphics.rememberGraphicsLayer
-import androidx.compose.ui.unit.toOffset
-import com.tunjid.treenav.Node
-import com.tunjid.treenav.compose.PaneScope
-import com.tunjid.treenav.compose.PaneState
-import com.tunjid.treenav.compose.utilities.AnimatedBoundsState
-import com.tunjid.treenav.compose.utilities.AnimatedBoundsState.Companion.animateBounds
-import kotlinx.coroutines.flow.debounce
-import kotlinx.coroutines.flow.first
 
-@Stable
 @OptIn(ExperimentalSharedTransitionApi::class)
-internal class MovableSharedElementState<State, Pane, Destination : Node>(
-    paneScope: PaneScope<Pane, Destination>,
-    sharedTransitionScope: SharedTransitionScope,
+@Stable
+internal class MovableSharedElementState<State>(
+    sharedContentState: SharedTransitionScope.SharedContentState,
     sharedElement: @Composable (State, Modifier) -> Unit,
-    onRemoved: () -> Unit,
-    boundsTransform: BoundsTransform,
-    private val canAnimateOnStartingFrames: PaneState<Pane, Destination>.() -> Boolean
-) : SharedElementOverlay, SharedTransitionScope by sharedTransitionScope {
-
-    var paneScope by mutableStateOf(paneScope)
+    onRemoved: () -> Unit
+) {
 
+    var sharedContentState by mutableStateOf(sharedContentState)
     private var composedRefCount by mutableIntStateOf(0)
 
-    private var layer: GraphicsLayer? = null
-    var animInProgress by mutableStateOf(false)
-        private set
-
-    private val canDrawInOverlay get() = animInProgress
-    private val panesKeysToSeenCount = mutableStateMapOf<String, Unit>()
-
-    private val animatedBoundsState = AnimatedBoundsState(
-        lookaheadScope = this,
-        boundsTransform = boundsTransform,
-        inProgress = { animInProgress }
-    )
-
     val moveableSharedElement: @Composable (Any?, Modifier) -> Unit =
         movableContentOf { state, modifier ->
-            animInProgress = isInProgress()
-            val layer = rememberGraphicsLayer().also {
-                this.layer = it
-            }
             @Suppress("UNCHECKED_CAST")
             sharedElement(
                 // The shared element composable will be created by the first screen and reused by
                 // subsequent screens. This updates the state from other screens so changes are seen.
                 state as State,
                 modifier
-                    .animateBounds(
-                        state = animatedBoundsState
-                    )
-                    .drawWithContent {
-                        layer.record {
-                            this@drawWithContent.drawContent()
-                        }
-                        if (!canDrawInOverlay) {
-                            drawLayer(layer)
-                        }
-                    },
             )
 
             DisposableEffect(Unit) {
@@ -90,58 +40,5 @@ internal class MovableSharedElementState<State, Pane, Destination : Node>(
                 }
             }
         }
-
-    override fun ContentDrawScope.drawInOverlay() {
-        if (!canDrawInOverlay) return
-        val overlayLayer = layer ?: return
-        val (x, y) = animatedBoundsState.targetOffset.toOffset()
-        translate(x, y) {
-            drawLayer(overlayLayer)
-        }
-    }
-
-    private fun updatePaneStateSeen(
-        paneState: PaneState<*, *>
-    ) {
-        panesKeysToSeenCount[paneState.key] = Unit
-    }
-
-    private val hasBeenShared get() = panesKeysToSeenCount.size > 1
-
-    companion object {
-
-        @Composable
-        private fun <Pane, Destination : Node> MovableSharedElementState<*, Pane, Destination>.isInProgress(): Boolean {
-            val paneState = paneScope.paneState.also(::updatePaneStateSeen)
-
-            val (laggingScopeKey, animationInProgressTillFirstIdle) = produceState(
-                initialValue = Pair(
-                    paneState.key,
-                    paneState.canAnimateOnStartingFrames()
-                ),
-                key1 = paneState.key
-            ) {
-                value = Pair(
-                    paneState.key,
-                    paneState.canAnimateOnStartingFrames()
-                )
-                value = snapshotFlow { animatedBoundsState.isIdle }
-                    .debounce { if (it) 10 else 0 }
-                    .first(true::equals)
-                    .let { value.first to false }
-            }.value
-
-
-            if (!hasBeenShared) return false
-
-            val isLagging = laggingScopeKey != paneScope.paneState.key
-            val canAnimateOnStartingFrames = paneScope.paneState.canAnimateOnStartingFrames()
-
-            if (isLagging) return canAnimateOnStartingFrames
-
-            return animationInProgressTillFirstIdle
-        }
-    }
 }
 
-private val PaneState<*, *>.key get() = "${currentDestination?.id}-$pane"