From 049d03f32114ceb682f406d025ebe9b6c4595986 Mon Sep 17 00:00:00 2001
From: "satyajit.happy" <satyajit.happy@gmail.com>
Date: Wed, 24 Jul 2019 11:20:13 +0200
Subject: [PATCH] refactor: rehydrate only once

---
 example/StackNavigator.tsx                |  4 +-
 example/TabNavigator.tsx                  |  4 +-
 src/__tests__/__fixtures__/MockRouter.tsx |  4 +-
 src/types.tsx                             |  1 -
 src/useNavigationBuilder.tsx              | 47 ++++++++++++-----------
 5 files changed, 30 insertions(+), 30 deletions(-)

diff --git a/example/StackNavigator.tsx b/example/StackNavigator.tsx
index 8c86d86c..bc5dac85 100644
--- a/example/StackNavigator.tsx
+++ b/example/StackNavigator.tsx
@@ -91,7 +91,7 @@ function StackRouter(options: DefaultRouterOptions) {
       };
     },
 
-    getRehydratedState({ key, routeNames, partialState }) {
+    getRehydratedState({ routeNames, partialState }) {
       let state = partialState;
 
       if (state.stale) {
@@ -99,7 +99,7 @@ function StackRouter(options: DefaultRouterOptions) {
           ...state,
           stale: false,
           routeNames,
-          key,
+          key: `stack-${shortid()}`,
         };
       }
 
diff --git a/example/TabNavigator.tsx b/example/TabNavigator.tsx
index 6dffeb96..47d42849 100644
--- a/example/TabNavigator.tsx
+++ b/example/TabNavigator.tsx
@@ -75,7 +75,7 @@ function TabRouter(options: DefaultRouterOptions) {
       };
     },
 
-    getRehydratedState({ key, routeNames, partialState }) {
+    getRehydratedState({ routeNames, partialState }) {
       let state = partialState;
 
       if (state.stale) {
@@ -83,7 +83,7 @@ function TabRouter(options: DefaultRouterOptions) {
           ...state,
           stale: false,
           routeNames,
-          key,
+          key: `tab-${shortid()}`,
         };
       }
 
diff --git a/src/__tests__/__fixtures__/MockRouter.tsx b/src/__tests__/__fixtures__/MockRouter.tsx
index de2404cc..b129b81b 100644
--- a/src/__tests__/__fixtures__/MockRouter.tsx
+++ b/src/__tests__/__fixtures__/MockRouter.tsx
@@ -32,7 +32,7 @@ export default function MockRouter(options: DefaultRouterOptions) {
       };
     },
 
-    getRehydratedState({ key, routeNames, partialState }) {
+    getRehydratedState({ routeNames, partialState }) {
       let state = partialState;
 
       if (state.stale) {
@@ -40,7 +40,7 @@ export default function MockRouter(options: DefaultRouterOptions) {
           ...state,
           stale: false,
           routeNames,
-          key,
+          key: String(MockRouterKey.current++),
         };
       }
 
diff --git a/src/types.tsx b/src/types.tsx
index fbe53ddb..c9dfd012 100644
--- a/src/types.tsx
+++ b/src/types.tsx
@@ -103,7 +103,6 @@ export type Router<
    * @param options.partialState Navigation state to rehydrate from.
    */
   getRehydratedState(options: {
-    key: string;
     routeNames: string[];
     partialState: State | PartialState<State>;
   }): State;
diff --git a/src/useNavigationBuilder.tsx b/src/useNavigationBuilder.tsx
index e815507e..27f22de3 100644
--- a/src/useNavigationBuilder.tsx
+++ b/src/useNavigationBuilder.tsx
@@ -87,26 +87,30 @@ export default function useNavigationBuilder<
     {} as { [key: string]: object | undefined }
   );
 
-  const [initialState] = React.useState(() =>
-    router.getInitialState({
-      routeNames,
-      routeParamList,
-    })
-  );
-
   const {
-    state: currentState = initialState,
+    state: currentState,
     getState: getCurrentState,
     setState,
     key,
     performTransaction,
   } = React.useContext(NavigationStateContext);
 
-  let state = router.getRehydratedState({
-    key: initialState.key,
-    routeNames,
-    partialState: currentState as any,
-  });
+  const [initialState] = React.useState(() =>
+    currentState === undefined
+      ? router.getInitialState({
+          routeNames,
+          routeParamList,
+        })
+      : router.getRehydratedState({
+          routeNames,
+          partialState: currentState as any,
+        })
+  );
+
+  let state =
+    currentState === undefined || currentState.stale
+      ? initialState
+      : (currentState as State);
 
   if (!isArrayEqual(state.routeNames, routeNames)) {
     // When the list of route names change, the router should handle it to remove invalid routes
@@ -137,16 +141,13 @@ export default function useNavigationBuilder<
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);
 
-  const getState = React.useCallback(
-    (): State =>
-      router.getRehydratedState({
-        key: initialState.key,
-        routeNames,
-        partialState: (getCurrentState() as any) || state,
-      }),
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-    [getCurrentState, router.getRehydratedState, router.getInitialState]
-  );
+  const getState = React.useCallback((): State => {
+    const currentState = getCurrentState();
+
+    return currentState === undefined || currentState.stale
+      ? initialState
+      : (currentState as State);
+  }, [getCurrentState, initialState]);
 
   const {
     listeners: actionListeners,