refactor: mark initial state as stale to determine when to rehydrate

example/StackNavigator.tsx

@@ -86,9 +86,10 @@ const StackRouter: Router<CommonAction | Action> = {
   getRehydratedState({ routeNames, partialState }) {
     let state = partialState;
 
-    if (state.routeNames === undefined || state.key === undefined) {
+    if (state.stale) {
       state = {
         ...state,
+        stale: false,
         routeNames,
         key: `stack-${shortid()}`,
       };

example/TabNavigator.tsx

@@ -72,9 +72,10 @@ const TabRouter: Router<Action | CommonAction> = {
   getRehydratedState({ routeNames, partialState }) {
     let state = partialState;
 
-    if (state.routeNames === undefined || state.key === undefined) {
+    if (state.stale) {
       state = {
         ...state,
+        stale: false,
         routeNames,
         key: `tab-${shortid()}`,
       };

example/index.tsx

@@ -3,9 +3,9 @@ import { render } from 'react-dom';
 import {
   NavigationContainer,
   CompositeNavigationProp,
-  PartialState,
   NavigationProp,
   RouteProp,
+  InitialState,
 } from '../src';
 import StackNavigator, { StackNavigationProp } from './StackNavigator';
 import TabNavigator, { TabNavigationProp } from './TabNavigator';
@@ -143,7 +143,7 @@ const Fifth = ({
 
 const PERSISTENCE_KEY = 'NAVIGATION_STATE';
 
-let initialState: PartialState | undefined;
+let initialState: InitialState | undefined;
 
 try {
   initialState = JSON.parse(localStorage.getItem(PERSISTENCE_KEY) || '');

src/NavigationContainer.tsx

@@ -1,9 +1,9 @@
 import * as React from 'react';
 import EnsureSingleNavigator from './EnsureSingleNavigator';
-import { NavigationState, PartialState } from './types';
+import { Route, NavigationState, InitialState, PartialState } from './types';
 
 type Props = {
-  initialState?: PartialState;
+  initialState?: InitialState;
   onStateChange?: (state: NavigationState | PartialState | undefined) => void;
   children: React.ReactNode;
 };
@@ -31,9 +31,37 @@ export const NavigationStateContext = React.createContext<{
   },
 });
 
-export default function NavigationContainer(props: Props) {
-  const { onStateChange } = props;
-  const [state, setState] = React.useState<State>(props.initialState);
+const getPartialState = (
+  state: InitialState | undefined
+): PartialState | undefined => {
+  if (state === undefined) {
+    return;
+  }
+
+  // @ts-ignore
+  return {
+    ...state,
+    stale: true,
+    key: undefined,
+    routeNames: undefined,
+    routes: state.routes.map(route => {
+      if (route.state === undefined) {
+        return route as Route<string> & { state?: PartialState };
+      }
+
+      return { ...route, state: getPartialState(route.state) };
+    }),
+  };
+};
+
+export default function NavigationContainer({
+  initialState,
+  onStateChange,
+  children,
+}: Props) {
+  const [state, setState] = React.useState<State>(() =>
+    getPartialState(initialState)
+  );
 
   const navigationStateRef = React.useRef<State | null>(null);
   const isTransactionActiveRef = React.useRef<boolean>(false);
@@ -92,7 +120,7 @@ export default function NavigationContainer(props: Props) {
 
   return (
     <NavigationStateContext.Provider value={context}>
-      <EnsureSingleNavigator>{props.children}</EnsureSingleNavigator>
+      <EnsureSingleNavigator>{children}</EnsureSingleNavigator>
     </NavigationStateContext.Provider>
   );
 }

src/__tests__/__fixtures__/MockRouter.tsx

@@ -1,7 +1,9 @@
 import { Router, CommonAction } from '../../types';
 import { BaseRouter } from '../../index';
 
-export type MockActions = CommonAction & { type: 'NOOP' | 'REVERSE' | 'UPDATE' };
+export type MockActions = CommonAction & {
+  type: 'NOOP' | 'REVERSE' | 'UPDATE';
+};
 
 const MockRouter: Router<MockActions> & { key: number } = {
   key: 0,
@@ -28,9 +30,10 @@ const MockRouter: Router<MockActions> & { key: number } = {
   getRehydratedState({ routeNames, partialState }) {
     let state = partialState;
 
-    if (state.routeNames === undefined || state.key === undefined) {
+    if (state.stale) {
       state = {
         ...state,
+        stale: false,
         routeNames,
         key: String(MockRouter.key++),
       };

src/__tests__/index.test.tsx

@@ -102,6 +102,7 @@ it('rehydrates state for a navigator on navigation', () => {
     key: '2',
     routeNames: ['foo', 'bar'],
     routes: [{ key: 'foo', name: 'foo' }, { key: 'bar', name: 'bar' }],
+    stale: false,
   });
 });
 
@@ -445,7 +446,10 @@ it('updates another route params with setParams', () => {
     index: 0,
     key: '0',
     routeNames: ['foo', 'bar'],
-    routes: [{ key: 'foo', name: 'foo', params: undefined }, { key: 'bar', name: 'bar', params: { username: 'alice' } }],
+    routes: [
+      { key: 'foo', name: 'foo', params: undefined },
+      { key: 'bar', name: 'bar', params: { username: 'alice' } },
+    ],
   });
 
   act(() => setParams({ age: 25 }, { name: 'bar' }));

src/__tests__/useOnAction.test.tsx

@@ -174,16 +174,18 @@ it("lets children handle the action if parent didn't", () => {
 
   expect(onStateChange).toBeCalledTimes(1);
   expect(onStateChange).lastCalledWith({
+    stale: false,
     index: 0,
-    key: '5',
+    key: '7',
     routeNames: ['foo', 'bar', 'baz'],
     routes: [
       {
         key: 'baz',
         name: 'baz',
         state: {
+          stale: false,
           index: 0,
-          key: '4',
+          key: '6',
           routeNames: ['qux', 'lex'],
           routes: [{ key: 'lex', name: 'lex' }, { key: 'qux', name: 'qux' }],
         },

src/types.tsx

@@ -23,13 +23,26 @@ export type NavigationState = {
   /**
    * List of rendered routes.
    */
-  routes: Array<Route<string> & { state?: NavigationState }>;
+  routes: Array<Route<string> & { state?: NavigationState | PartialState }>;
+  /**
+   * Whether the navigation state has been rehydrated.
+   */
+  stale?: false;
+};
+
+export type InitialState = Omit<
+  Omit<Omit<NavigationState, 'key'>, 'routes'>,
+  'routeNames'
+> & {
+  key?: string;
+  routeNames?: string[];
+  routes: Array<Route<string> & { state?: InitialState }>;
 };
 
-export type PartialState = Omit<Omit<NavigationState, 'routeNames'>, 'key'> & {
+export type PartialState = NavigationState & {
+  stale: true;
   key?: undefined;
   routeNames?: undefined;
-  state?: PartialState;
 };
 
 export type Route<RouteName extends string> = {

src/useDescriptors.tsx

@@ -61,7 +61,7 @@ export default function useDescriptors<ScreenOptions extends object>({
       acc[route.key] = {
         render() {
           return (
-            <NavigationBuilderContext.Provider value={context}>
+            <NavigationBuilderContext.Provider key={route.key} value={context}>
               <SceneView
                 navigation={navigation}
                 route={route}