From 26f3ee8149298cf523bc6e59b105c9877bbb61b8 Mon Sep 17 00:00:00 2001
From: Isaac Udy <isaac.udy@gmail.com>
Date: Mon, 19 Jun 2023 20:02:28 +1200
Subject: [PATCH] Improve remembers related to the local navigation
 context/navigation handle so that the correct instance is returned in
 situations where it may have changed (i.e. when the LifecycleOwner or View
 has also changed)

---
 .../main/java/dev/enro/core/NavigationContext.kt    | 13 ++++++++++++-
 .../enro/core/compose/ComposableNavigationHandle.kt |  4 ++--
 .../core/compose/rememberNavigationContainer.kt     | 10 +++++++++-
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/enro-core/src/main/java/dev/enro/core/NavigationContext.kt b/enro-core/src/main/java/dev/enro/core/NavigationContext.kt
index 2925d1d1..ff853b22 100644
--- a/enro-core/src/main/java/dev/enro/core/NavigationContext.kt
+++ b/enro-core/src/main/java/dev/enro/core/NavigationContext.kt
@@ -5,6 +5,9 @@ import android.os.Bundle
 import androidx.activity.ComponentActivity
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.remember
+import androidx.compose.ui.platform.LocalContext
+import androidx.compose.ui.platform.LocalLifecycleOwner
+import androidx.compose.ui.platform.LocalView
 import androidx.core.os.bundleOf
 import androidx.fragment.app.Fragment
 import androidx.fragment.app.FragmentActivity
@@ -243,7 +246,15 @@ public val containerManager: NavigationContainerManager
     @Composable
     get() {
         val viewModelStoreOwner = LocalViewModelStoreOwner.current!!
-        return remember {
+
+        val context = LocalContext.current
+        val view = LocalView.current
+        val lifecycleOwner = LocalLifecycleOwner.current
+
+        // The navigation context attached to a NavigationHandle may change when the Context, View,
+        // or LifecycleOwner changes, so we're going to re-query the navigation context whenever
+        // any of these change, to ensure the container always has an up-to-date NavigationContext
+        return remember(context, view, lifecycleOwner) {
             viewModelStoreOwner
                 .navigationContext!!
                 .containerManager
diff --git a/enro-core/src/main/java/dev/enro/core/compose/ComposableNavigationHandle.kt b/enro-core/src/main/java/dev/enro/core/compose/ComposableNavigationHandle.kt
index 5a01fe3d..560f72ea 100644
--- a/enro-core/src/main/java/dev/enro/core/compose/ComposableNavigationHandle.kt
+++ b/enro-core/src/main/java/dev/enro/core/compose/ComposableNavigationHandle.kt
@@ -9,7 +9,7 @@ import dev.enro.core.*
 @Composable
 public inline fun <reified T : NavigationKey> navigationHandle(): TypedNavigationHandle<T> {
     val navigationHandle = navigationHandle()
-    return remember {
+    return remember(navigationHandle) {
         navigationHandle.asTyped()
     }
 }
@@ -19,7 +19,7 @@ public fun navigationHandle(): NavigationHandle {
     val localNavigationHandle = LocalNavigationHandle.current
     val localViewModelStoreOwner = LocalViewModelStoreOwner.current
 
-    return remember {
+    return remember(localNavigationHandle, localViewModelStoreOwner) {
         localNavigationHandle ?: localViewModelStoreOwner!!.getNavigationHandle()
     }
 }
diff --git a/enro-core/src/main/java/dev/enro/core/compose/rememberNavigationContainer.kt b/enro-core/src/main/java/dev/enro/core/compose/rememberNavigationContainer.kt
index facdfd1f..08a4a735 100644
--- a/enro-core/src/main/java/dev/enro/core/compose/rememberNavigationContainer.kt
+++ b/enro-core/src/main/java/dev/enro/core/compose/rememberNavigationContainer.kt
@@ -5,6 +5,8 @@ import androidx.compose.runtime.*
 import androidx.compose.runtime.saveable.rememberSaveable
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.platform.LocalContext
+import androidx.compose.ui.platform.LocalLifecycleOwner
+import androidx.compose.ui.platform.LocalView
 import dev.enro.core.*
 import dev.enro.core.compose.container.ComposableNavigationContainer
 import dev.enro.core.compose.container.ContainerRegistrationStrategy
@@ -67,7 +69,13 @@ public fun rememberNavigationContainer(
 ): ComposableNavigationContainer {
     val localNavigationHandle = navigationHandle()
     val context = LocalContext.current
-    val localNavigationContext = remember(context) {
+    val view = LocalView.current
+    val lifecycleOwner = LocalLifecycleOwner.current
+
+    // The navigation context attached to a NavigationHandle may change when the Context, View,
+    // or LifecycleOwner changes, so we're going to re-query the navigation context whenever
+    // any of these change, to ensure the container always has an up-to-date NavigationContext
+    val localNavigationContext = remember(context, view, lifecycleOwner) {
         localNavigationHandle.requireNavigationContext()
     }
     val navigationContainer = remember(localNavigationContext.containerManager) {