refactor: rehydrate only once

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()}`,
         };
       }
 

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()}`,
         };
       }
 

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++),
         };
       }
 

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;

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,