refactor: add separate method for rehydration and fix types

README.md

@@ -113,15 +113,15 @@ For our example above, we need 2 separate types for stack and tabs:
 
 ```ts
 type StackParamList = {
-  settings: void;
+  settings: undefined;
   profile: { userId: string };
-  home: void;
+  home: undefined;
 };
 
 type TabParamList = {
-  feed: void;
-  article: void;
-  notifications: void;
+  feed: undefined;
+  article: undefined;
+  notifications: undefined;
 };
 ```
 

example/StackNavigator.tsx

@@ -55,27 +55,27 @@ export type StackNavigationProp<
 
 const StackRouter: Router<CommonAction | Action> = {
   getInitialState({
-    screens,
-    partialState,
-    initialRouteName = Object.keys(screens)[0],
+    routeNames,
+    initialRouteName = routeNames[0],
+    initialParamsList,
   }) {
-    const routeNames = Object.keys(screens);
+    const index = routeNames.indexOf(initialRouteName);
 
+    return {
+      key: `stack-${shortid()}`,
+      index,
+      routeNames,
+      routes: routeNames.slice(0, index + 1).map(name => ({
+        name,
+        key: `${name}-${shortid()}`,
+        params: initialParamsList[name],
+      })),
+    };
+  },
+
+  getRehydratedState({ routeNames, partialState }) {
     let state = partialState;
 
-    if (state === undefined) {
-      const index = routeNames.indexOf(initialRouteName);
-
-      state = {
-        index,
-        routes: routeNames.slice(0, index + 1).map(name => ({
-          name,
-          key: `${name}-${shortid()}`,
-          params: screens[name].initialParams,
-        })),
-      };
-    }
-
     if (state.routeNames === undefined || state.key === undefined) {
       state = {
         ...state,

example/TabNavigator.tsx

@@ -39,27 +39,27 @@ export type TabNavigationProp<
 
 const TabRouter: Router<Action | CommonAction> = {
   getInitialState({
-    screens,
-    partialState,
-    initialRouteName = Object.keys(screens)[0],
+    routeNames,
+    initialRouteName = routeNames[0],
+    initialParamsList,
   }) {
-    const routeNames = Object.keys(screens);
+    const index = routeNames.indexOf(initialRouteName);
 
+    return {
+      key: `tab-${shortid()}`,
+      index,
+      routeNames,
+      routes: routeNames.map(name => ({
+        name,
+        key: `${name}-${shortid()}`,
+        params: initialParamsList[name],
+      })),
+    };
+  },
+
+  getRehydratedState({ routeNames, partialState }) {
     let state = partialState;
 
-    if (state === undefined) {
-      const index = routeNames.indexOf(initialRouteName);
-
-      state = {
-        index,
-        routes: routeNames.map(name => ({
-          name,
-          key: `${name}-${shortid()}`,
-          params: screens[name].initialParams,
-        })),
-      };
-    }
-
     if (state.routeNames === undefined || state.key === undefined) {
       state = {
         ...state,

example/index.tsx

@@ -6,20 +6,20 @@ import {
   CompositeNavigationProp,
   TypedNavigator,
   NavigationHelpers,
-  InitialState,
+  PartialState,
 } from '../src';
 import StackNavigator, { StackNavigationProp } from './StackNavigator';
 import TabNavigator, { TabNavigationProp } from './TabNavigator';
 
 type StackParamList = {
   first: { author: string };
-  second: void;
-  third: void;
+  second: undefined;
+  third: undefined;
 };
 
 type TabParamList = {
-  fourth: void;
-  fifth: void;
+  fourth: undefined;
+  fifth: undefined;
 };
 
 const Stack: TypedNavigator<StackParamList, typeof StackNavigator> = {
@@ -134,7 +134,7 @@ const Fifth = ({
 
 const PERSISTENCE_KEY = 'NAVIGATION_STATE';
 
-let initialState: InitialState | undefined;
+let initialState: PartialState | undefined;
 
 try {
   initialState = JSON.parse(localStorage.getItem(PERSISTENCE_KEY) || '');

src/BaseActions.tsx

@@ -1,4 +1,4 @@
-import { InitialState } from './types';
+import { PartialState } from './types';
 
 export type Action =
   | { type: 'GO_BACK' }
@@ -12,7 +12,7 @@ export type Action =
     }
   | {
       type: 'RESET';
-      payload: InitialState & { key?: string };
+      payload: PartialState & { key?: string };
     };
 
 export function goBack(): Action {
@@ -27,6 +27,6 @@ export function replace(name: string, params?: object): Action {
   return { type: 'REPLACE', payload: { name, params } };
 }
 
-export function reset(state: InitialState & { key?: string }): Action {
+export function reset(state: PartialState & { key?: string }): Action {
   return { type: 'RESET', payload: state };
 }

src/NavigationContainer.tsx

@@ -1,10 +1,10 @@
 import * as React from 'react';
 import EnsureSingleNavigator from './EnsureSingleNavigator';
-import { NavigationState, InitialState } from './types';
+import { NavigationState, PartialState } from './types';
 
 type Props = {
-  initialState?: InitialState;
-  onStateChange?: (state: NavigationState | InitialState) => void;
+  initialState?: PartialState;
+  onStateChange?: (state: NavigationState | PartialState) => void;
   children: React.ReactNode;
 };
 
@@ -12,8 +12,8 @@ const MISSING_CONTEXT_ERROR =
   "We couldn't find a navigation context. Have you wrapped your app with 'NavigationContainer'?";
 
 export const NavigationStateContext = React.createContext<{
-  state?: NavigationState | InitialState;
-  getState: () => NavigationState | InitialState | undefined;
+  state?: NavigationState | PartialState;
+  getState: () => NavigationState | PartialState | undefined;
   setState: (state: NavigationState) => void;
 }>({
   get getState(): any {
@@ -30,7 +30,7 @@ export default function NavigationContainer({
   children,
 }: Props) {
   const [state, setState] = React.useState<
-    NavigationState | InitialState | undefined
+    NavigationState | PartialState | undefined
   >(initialState);
 
   const stateRef = React.useRef(state);

src/__tests__/index.test.tsx

@@ -6,24 +6,28 @@ import useNavigationBuilder from '../useNavigationBuilder';
 import { Router } from '../types';
 
 const MockRouter: Router<{ type: string }> = {
-  getInitialState({ screens, partialState, initialRouteName }) {
-    const routeNames = Object.keys(screens);
+  getInitialState({
+    routeNames,
+    initialRouteName = routeNames[0],
+    initialParamsList,
+  }) {
+    const index = routeNames.indexOf(initialRouteName);
 
+    return {
+      key: 'root',
+      index,
+      routeNames,
+      routes: routeNames.map(name => ({
+        name,
+        key: name,
+        params: initialParamsList[name],
+      })),
+    };
+  },
+
+  getRehydratedState({ routeNames, partialState }) {
     let state = partialState;
 
-    if (state === undefined) {
-      const index = routeNames.indexOf(initialRouteName || routeNames[0]);
-
-      state = {
-        index,
-        routes: routeNames.map(name => ({
-          name,
-          key: name,
-          params: screens[name].initialParams,
-        })),
-      };
-    }
-
     if (state.routeNames === undefined || state.key === undefined) {
       state = {
         ...state,
@@ -76,7 +80,11 @@ it('initializes state for a navigator on navigation', () => {
   const element = (
     <NavigationContainer onStateChange={onStateChange}>
       <TestNavigator initialRouteName="foo">
-        <Screen name="foo" component={FooScreen} />
+        <Screen
+          name="foo"
+          component={FooScreen}
+          initialParams={{ count: 10 }}
+        />
         <Screen name="bar" component={jest.fn()} />
         <Screen name="baz">
           {() => (
@@ -97,13 +105,62 @@ it('initializes state for a navigator on navigation', () => {
     key: 'root',
     routeNames: ['foo', 'bar', 'baz'],
     routes: [
-      { key: 'foo', name: 'foo' },
+      { key: 'foo', name: 'foo', params: { count: 10 } },
       { key: 'bar', name: 'bar' },
       { key: 'baz', name: 'baz' },
     ],
   });
 });
 
+it('rehydrates state for a navigator on navigation', () => {
+  expect.assertions(1);
+
+  const TestNavigator = (props: any) => {
+    const { navigation, descriptors } = useNavigationBuilder(MockRouter, props);
+
+    return descriptors[
+      navigation.state.routes[navigation.state.index].key
+    ].render();
+  };
+
+  const BarScreen = (props: any) => {
+    React.useEffect(() => {
+      props.navigation.dispatch({ type: 'UPDATE' });
+      // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, []);
+
+    return null;
+  };
+
+  const initialState = {
+    index: 1,
+    routes: [{ key: 'foo', name: 'foo' }, { key: 'bar', name: 'bar' }],
+  };
+
+  const onStateChange = jest.fn();
+
+  const element = (
+    <NavigationContainer
+      initialState={initialState}
+      onStateChange={onStateChange}
+    >
+      <TestNavigator initialRouteName="foo">
+        <Screen name="foo" component={jest.fn()} />
+        <Screen name="bar" component={BarScreen} />
+      </TestNavigator>
+    </NavigationContainer>
+  );
+
+  render(element).update(element);
+
+  expect(onStateChange).lastCalledWith({
+    index: 1,
+    key: 'root',
+    routeNames: ['foo', 'bar'],
+    routes: [{ key: 'foo', name: 'foo' }, { key: 'bar', name: 'bar' }],
+  });
+});
+
 it('initializes state for nested navigator on navigation', () => {
   expect.assertions(2);
 

src/types.tsx

@@ -21,10 +21,10 @@ export type NavigationState = {
   routes: Array<Route & { state?: NavigationState }>;
 };
 
-export type InitialState = Omit<Omit<NavigationState, 'routeNames'>, 'key'> & {
+export type PartialState = Omit<Omit<NavigationState, 'routeNames'>, 'key'> & {
   key?: undefined;
   routeNames?: undefined;
-  state?: InitialState;
+  state?: PartialState;
 };
 
 export type Route<RouteName = string> = {
@@ -39,7 +39,7 @@ export type Route<RouteName = string> = {
   /**
    * Params for the route.
    */
-  params?: object | void;
+  params?: object;
 };
 
 export type NavigationAction = {
@@ -52,12 +52,20 @@ export type ActionCreators<Action extends NavigationAction = CommonAction> = {
 
 export type Router<Action extends NavigationAction = CommonAction> = {
   /**
-   * Initialize full navigation state with a given partial state.
+   * Initialize the navigation state.
    */
   getInitialState(options: {
-    screens: { [key: string]: ScreenProps };
-    partialState?: NavigationState | InitialState;
-    initialRouteName?: string;
+    routeNames: string[];
+    initialRouteName: string;
+    initialParamsList: ParamListBase;
+  }): NavigationState;
+
+  /**
+   * Rehydrate the full navigation state from a given partial state.
+   */
+  getRehydratedState(options: {
+    routeNames: string[];
+    partialState: NavigationState | PartialState;
   }): NavigationState;
 
   /**
@@ -75,7 +83,7 @@ export type Router<Action extends NavigationAction = CommonAction> = {
   actionCreators: ActionCreators<Action>;
 };
 
-export type ParamListBase = { [key: string]: object | void };
+export type ParamListBase = { [key: string]: object | undefined };
 
 class PrivateValueStore<T> {
   /**
@@ -104,8 +112,8 @@ export type NavigationHelpers<
    * @param [params] Params object for the route.
    */
   navigate<RouteName extends keyof ParamList>(
-    ...args: ParamList[RouteName] extends void
-      ? [RouteName]
+    ...args: ParamList[RouteName] extends undefined
+      ? [RouteName] | [RouteName, undefined]
       : [RouteName, ParamList[RouteName]]
   ): void;
 
@@ -116,8 +124,8 @@ export type NavigationHelpers<
    * @param [params] Params object for the new route.
    */
   replace<RouteName extends keyof ParamList>(
-    ...args: ParamList[RouteName] extends void
-      ? [RouteName]
+    ...args: ParamList[RouteName] extends undefined
+      ? [RouteName] | [RouteName, undefined]
       : [RouteName, ParamList[RouteName]]
   ): void;
 
@@ -125,7 +133,7 @@ export type NavigationHelpers<
    * Reset the navigation state to the provided state.
    * If a key is provided, the state with matching key will be reset.
    */
-  reset(state: InitialState & { key?: string }): void;
+  reset(state: PartialState & { key?: string }): void;
 
   /**
    * Go back to the previous route in history.
@@ -140,9 +148,9 @@ export type NavigationProp<
   /**
    * State for the child navigator.
    */
-  state: Route<RouteName> &
-    (ParamList[RouteName] extends void
-      ? never
+  state: Omit<Route<RouteName>, 'params'> &
+    (ParamList[RouteName] extends undefined
+      ? {}
       : { params: ParamList[RouteName] });
 
   /**

src/useDescriptors.tsx

@@ -1,7 +1,7 @@
 import * as React from 'react';
 import {
   Descriptor,
-  InitialState,
+  PartialState,
   NavigationAction,
   NavigationHelpers,
   NavigationState,
@@ -12,7 +12,7 @@ import SceneView from './SceneView';
 import NavigationBuilderContext from './NavigationBuilderContext';
 
 type Options = {
-  state: NavigationState | InitialState;
+  state: NavigationState | PartialState;
   screens: { [key: string]: ScreenProps<ParamListBase, string> };
   helpers: NavigationHelpers<ParamListBase>;
   onAction: (action: NavigationAction) => boolean;

src/useNavigationBuilder.tsx

@@ -18,37 +18,51 @@ export default function useNavigationBuilder(
 ) {
   useRegisterNavigator();
 
-  const screens = React.Children.map(options.children, child => {
-    if (child === null || child === undefined) {
-      return;
-    }
+  const [screens] = React.useState(() =>
+    React.Children.map(options.children, child => {
+      if (child === null || child === undefined) {
+        return;
+      }
 
-    if (React.isValidElement(child) && child.type === Screen) {
-      return child.props as ScreenProps;
-    }
+      if (React.isValidElement(child) && child.type === Screen) {
+        return child.props as ScreenProps;
+      }
 
-    throw new Error(
-      `A navigator can only contain 'Screen' components as its direct children (found '${
-        // @ts-ignore
-        child.type && child.type.name ? child.type.name : String(child)
-      }')`
-    );
-  })
-    .filter(Boolean)
-    .reduce(
-      (acc, curr) => {
-        acc[curr!.name] = curr as ScreenProps;
-        return acc;
-      },
-      {} as { [key: string]: ScreenProps }
-    );
+      throw new Error(
+        `A navigator can only contain 'Screen' components as its direct children (found '${
+          // @ts-ignore
+          child.type && child.type.name ? child.type.name : String(child)
+        }')`
+      );
+    })
+      .filter(Boolean)
+      .reduce(
+        (acc, curr) => {
+          acc[curr!.name] = curr as ScreenProps;
+          return acc;
+        },
+        {} as { [key: string]: ScreenProps }
+      )
+  );
 
   const routeNames = Object.keys(screens);
+  const initialRouteName =
+    options.initialRouteName !== undefined
+      ? options.initialRouteName
+      : routeNames[0];
+  const initialParamsList = routeNames.reduce(
+    (acc, curr) => {
+      acc[curr] = screens[curr].initialParams;
+      return acc;
+    },
+    {} as { [key: string]: object | undefined }
+  );
 
   const {
     state: currentState = router.getInitialState({
-      screens,
-      initialRouteName: options.initialRouteName,
+      routeNames,
+      initialRouteName,
+      initialParamsList,
     }),
     getState: getCurrentState,
     setState,
@@ -56,13 +70,18 @@ export default function useNavigationBuilder(
 
   const getState = React.useCallback(
     (): NavigationState =>
-      router.getInitialState({
-        screens,
-        partialState: getCurrentState(),
-        initialRouteName: options.initialRouteName,
+      router.getRehydratedState({
+        routeNames,
+        partialState:
+          getCurrentState() ||
+          router.getInitialState({
+            routeNames,
+            initialRouteName,
+            initialParamsList,
+          }),
       }),
     // eslint-disable-next-line react-hooks/exhaustive-deps
-    [getCurrentState, ...routeNames]
+    [getCurrentState, router.getRehydratedState, router.getInitialState]
   );
 
   const onAction = useOnAction({