refactor: let router specify its state shape

example/StackNavigator.tsx

@@ -5,6 +5,7 @@ import shortid from 'shortid';
 import {
   useNavigationBuilder,
   NavigationProp,
+  NavigationState,
   CommonAction,
   ParamListBase,
   Router,
@@ -38,7 +39,12 @@ export type StackNavigationOptions = {
 export type StackNavigationProp<
   ParamList extends ParamListBase,
   RouteName extends keyof ParamList = string
-> = NavigationProp<ParamList, RouteName, StackNavigationOptions> & {
+> = NavigationProp<
+  ParamList,
+  RouteName,
+  NavigationState,
+  StackNavigationOptions
+> & {
   /**
    * Push a new screen onto the stack.
    *
@@ -62,7 +68,7 @@ export type StackNavigationProp<
   popToTop(): void;
 };
 
-const StackRouter: Router<CommonAction | Action> = {
+const StackRouter: Router<NavigationState, CommonAction | Action> = {
   ...BaseRouter,
 
   getInitialState({
@@ -245,10 +251,10 @@ const StackRouter: Router<CommonAction | Action> = {
 };
 
 export function StackNavigator(props: Props) {
-  const { state, descriptors } = useNavigationBuilder<StackNavigationOptions>(
-    StackRouter,
-    props
-  );
+  const { state, descriptors } = useNavigationBuilder<
+    NavigationState,
+    StackNavigationOptions
+  >(StackRouter, props);
 
   return (
     <div style={{ position: 'relative' }}>

example/TabNavigator.tsx

@@ -10,6 +10,7 @@ import {
   Router,
   createNavigator,
   BaseRouter,
+  NavigationState,
 } from '../src/index';
 
 type Props = {
@@ -32,7 +33,12 @@ export type TabNavigationOptions = {
 export type TabNavigationProp<
   ParamList extends ParamListBase,
   RouteName extends keyof ParamList = string
-> = NavigationProp<ParamList, RouteName, TabNavigationOptions> & {
+> = NavigationProp<
+  ParamList,
+  RouteName,
+  NavigationState,
+  TabNavigationOptions
+> & {
   /**
    * Jump to an existing tab.
    *
@@ -46,7 +52,7 @@ export type TabNavigationProp<
   ): void;
 };
 
-const TabRouter: Router<Action | CommonAction> = {
+const TabRouter: Router<NavigationState, Action | CommonAction> = {
   ...BaseRouter,
 
   getInitialState({
@@ -169,7 +175,10 @@ const TabRouter: Router<Action | CommonAction> = {
 };
 
 export function TabNavigator(props: Props) {
-  const { state, descriptors } = useNavigationBuilder(TabRouter, props);
+  const { state, descriptors } = useNavigationBuilder<
+    NavigationState,
+    TabNavigationOptions
+  >(TabRouter, props);
 
   return (
     <div style={{ display: 'flex', flexDirection: 'row', height: '100%' }}>

src/BaseActions.tsx

@@ -1,4 +1,4 @@
-import { PartialState, TargetRoute } from './types';
+import { PartialState, NavigationState, TargetRoute } from './types';
 
 export type Action =
   | { type: 'GO_BACK' }
@@ -14,7 +14,7 @@ export type Action =
     }
   | {
       type: 'RESET';
-      payload: PartialState & { key?: string };
+      payload: PartialState<NavigationState> & { key?: string };
     }
   | {
       type: 'SET_PARAMS';
@@ -47,7 +47,9 @@ export function replace(name: string, params?: object): Action {
   return { type: 'REPLACE', payload: { name, params } };
 }
 
-export function reset(state: PartialState & { key?: string }): Action {
+export function reset(
+  state: PartialState<NavigationState> & { key?: string }
+): Action {
   return { type: 'RESET', payload: state };
 }
 

src/BaseRouter.tsx

@@ -2,7 +2,10 @@ import shortid from 'shortid';
 import { CommonAction, NavigationState } from './types';
 
 const BaseRouter = {
-  getStateForAction(state: NavigationState, action: CommonAction) {
+  getStateForAction<State extends NavigationState>(
+    state: State,
+    action: CommonAction
+  ): State | null {
     switch (action.type) {
       case 'REPLACE': {
         return {
@@ -35,7 +38,7 @@ const BaseRouter = {
           action.payload.key === state.key
         ) {
           return {
-            ...action.payload,
+            ...(action.payload as any),
             key: state.key,
             routeNames: state.routeNames,
           };

src/NavigationContainer.tsx

@@ -4,18 +4,20 @@ import { Route, NavigationState, InitialState, PartialState } from './types';
 
 type Props = {
   initialState?: InitialState;
-  onStateChange?: (state: NavigationState | PartialState | undefined) => void;
+  onStateChange?: (
+    state: NavigationState | PartialState<NavigationState> | undefined
+  ) => void;
   children: React.ReactNode;
 };
 
-type State = NavigationState | PartialState | undefined;
+type State = NavigationState | PartialState<NavigationState> | undefined;
 
 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 | PartialState;
-  getState: () => NavigationState | PartialState | undefined;
+  state?: NavigationState | PartialState<NavigationState>;
+  getState: () => NavigationState | PartialState<NavigationState> | undefined;
   setState: (state: NavigationState | undefined) => void;
   key?: string;
   performTransaction: (action: () => void) => void;
@@ -33,7 +35,7 @@ export const NavigationStateContext = React.createContext<{
 
 const getPartialState = (
   state: InitialState | undefined
-): PartialState | undefined => {
+): PartialState<NavigationState> | undefined => {
   if (state === undefined) {
     return;
   }
@@ -46,7 +48,9 @@ const getPartialState = (
     routeNames: undefined,
     routes: state.routes.map(route => {
       if (route.state === undefined) {
-        return route as Route<string> & { state?: PartialState };
+        return route as Route<string> & {
+          state?: PartialState<NavigationState>;
+        };
       }
 
       return { ...route, state: getPartialState(route.state) };

src/__tests__/__fixtures__/MockRouter.tsx

@@ -1,11 +1,11 @@
-import { Router, CommonAction } from '../../types';
+import { Router, CommonAction, NavigationState } from '../../types';
 import { BaseRouter } from '../../index';
 
 export type MockActions = CommonAction & {
   type: 'NOOP' | 'REVERSE' | 'UPDATE';
 };
 
-const MockRouter: Router<MockActions> & { key: number } = {
+const MockRouter: Router<NavigationState, MockActions> & { key: number } = {
   key: 0,
 
   getInitialState({

src/__tests__/useDescriptors.test.tsx

@@ -4,6 +4,7 @@ import useNavigationBuilder from '../useNavigationBuilder';
 import NavigationContainer from '../NavigationContainer';
 import Screen from '../Screen';
 import MockRouter from './__fixtures__/MockRouter';
+import { NavigationState } from '../types';
 
 jest.useFakeTimers();
 
@@ -11,10 +12,10 @@ beforeEach(() => (MockRouter.key = 0));
 
 it('sets options with options prop as an object', () => {
   const TestNavigator = (props: any) => {
-    const { state, descriptors } = useNavigationBuilder<{ title?: string }>(
-      MockRouter,
-      props
-    );
+    const { state, descriptors } = useNavigationBuilder<
+      NavigationState,
+      { title?: string }
+    >(MockRouter, props);
     const { render, options } = descriptors[state.routes[state.index].key];
 
     return (
@@ -54,10 +55,10 @@ it('sets options with options prop as an object', () => {
 
 it('sets options with options prop as a fuction', () => {
   const TestNavigator = (props: any) => {
-    const { state, descriptors } = useNavigationBuilder<{ title?: string }>(
-      MockRouter,
-      props
-    );
+    const { state, descriptors } = useNavigationBuilder<
+      NavigationState,
+      { title?: string }
+    >(MockRouter, props);
     const { render, options } = descriptors[state.routes[state.index].key];
 
     return (
@@ -98,10 +99,13 @@ it('sets options with options prop as a fuction', () => {
 
 it('sets initial options with setOptions', () => {
   const TestNavigator = (props: any) => {
-    const { state, descriptors } = useNavigationBuilder<{
-      title?: string;
-      color?: string;
-    }>(MockRouter, props);
+    const { state, descriptors } = useNavigationBuilder<
+      NavigationState,
+      {
+        title?: string;
+        color?: string;
+      }
+    >(MockRouter, props);
     const { render, options } = descriptors[state.routes[state.index].key];
 
     return (
@@ -147,7 +151,10 @@ it('sets initial options with setOptions', () => {
 
 it('updates options with setOptions', () => {
   const TestNavigator = (props: any) => {
-    const { state, descriptors } = useNavigationBuilder<any>(MockRouter, props);
+    const { state, descriptors } = useNavigationBuilder<NavigationState, any>(
+      MockRouter,
+      props
+    );
     const { render, options } = descriptors[state.routes[state.index].key];
 
     return (

src/__tests__/useOnAction.test.tsx

@@ -1,15 +1,15 @@
 import * as React from 'react';
 import { render } from 'react-native-testing-library';
-import { Router } from '../types';
 import useNavigationBuilder from '../useNavigationBuilder';
 import NavigationContainer from '../NavigationContainer';
 import Screen from '../Screen';
 import MockRouter, { MockActions } from './__fixtures__/MockRouter';
+import { Router, NavigationState } from '../types';
 
 beforeEach(() => (MockRouter.key = 0));
 
 it("lets parent handle the action if child didn't", () => {
-  const ParentRouter: Router<MockActions> = {
+  const ParentRouter: Router<NavigationState, MockActions> = {
     ...MockRouter,
 
     getStateForAction(state, action) {
@@ -78,7 +78,7 @@ it("lets parent handle the action if child didn't", () => {
 });
 
 it("lets children handle the action if parent didn't", () => {
-  const ParentRouter: Router<MockActions> = {
+  const ParentRouter: Router<NavigationState, MockActions> = {
     ...MockRouter,
 
     shouldActionPropagateToChildren() {
@@ -86,7 +86,7 @@ it("lets children handle the action if parent didn't", () => {
     },
   };
 
-  const ChildRouter: Router<MockActions> = {
+  const ChildRouter: Router<NavigationState, MockActions> = {
     ...MockRouter,
 
     shouldActionChangeFocus() {

src/types.tsx

@@ -23,7 +23,9 @@ export type NavigationState = {
   /**
    * List of rendered routes.
    */
-  routes: Array<Route<string> & { state?: NavigationState | PartialState }>;
+  routes: Array<
+    Route<string> & { state?: NavigationState | PartialState<NavigationState> }
+  >;
   /**
    * Whether the navigation state has been rehydrated.
    */
@@ -39,7 +41,7 @@ export type InitialState = Omit<
   routes: Array<Route<string> & { state?: InitialState }>;
 };
 
-export type PartialState = NavigationState & {
+export type PartialState<State extends NavigationState> = State & {
   stale: true;
   key?: undefined;
   routeNames?: undefined;
@@ -68,7 +70,10 @@ export type ActionCreators<Action extends NavigationAction> = {
   [key: string]: (...args: any) => Action;
 };
 
-export type Router<Action extends NavigationAction> = {
+export type Router<
+  State extends NavigationState,
+  Action extends NavigationAction
+> = {
   /**
    * Initialize the navigation state.
    *
@@ -80,7 +85,7 @@ export type Router<Action extends NavigationAction> = {
     routeNames: string[];
     initialRouteName: string;
     initialParamsList: ParamListBase;
-  }): NavigationState;
+  }): State;
 
   /**
    * Rehydrate the full navigation state from a given partial state.
@@ -90,8 +95,8 @@ export type Router<Action extends NavigationAction> = {
    */
   getRehydratedState(options: {
     routeNames: string[];
-    partialState: NavigationState | PartialState;
-  }): NavigationState;
+    partialState: State | PartialState<State>;
+  }): State;
 
   /**
    * Take the current state and updated list of route names, and return a new state.
@@ -102,13 +107,13 @@ export type Router<Action extends NavigationAction> = {
    * @param options.initialParamsList Object containing initial params for each route.
    */
   getStateForRouteNamesChange(
-    state: NavigationState,
+    state: State,
     options: {
       routeNames: string[];
       initialRouteName: string;
       initialParamsList: ParamListBase;
     }
-  ): NavigationState;
+  ): State;
 
   /**
    * Take the current state and key of a route, and return a new state with the route focused
@@ -116,7 +121,7 @@ export type Router<Action extends NavigationAction> = {
    * @param state State object to apply the action on.
    * @param key Key of the route to focus.
    */
-  getStateForRouteFocus(state: NavigationState, key: string): NavigationState;
+  getStateForRouteFocus(state: State, key: string): State;
 
   /**
    * Take the current state and action, and return a new state.
@@ -125,10 +130,7 @@ export type Router<Action extends NavigationAction> = {
    * @param state State object to apply the action on.
    * @param action Action object to apply.
    */
-  getStateForAction(
-    state: NavigationState,
-    action: Action
-  ): NavigationState | null;
+  getStateForAction(state: State, action: Action): State | null;
 
   /**
    * Whether the action bubbles to other navigators
@@ -162,16 +164,17 @@ class PrivateValueStore<T> {
   private __private_value_type?: T;
 }
 
-type NavigationHelpersCommon<ParamList extends ParamListBase> = {
+type NavigationHelpersCommon<
+  ParamList extends ParamListBase,
+  State extends NavigationState = NavigationState
+> = {
   /**
    * Dispatch an action or an update function to the router.
    * The update function will receive the current state,
    *
    * @param action Action object or update function.
    */
-  dispatch(
-    action: NavigationAction | ((state: NavigationState) => NavigationState)
-  ): void;
+  dispatch(action: NavigationAction | ((state: State) => State)): void;
 
   /**
    * Navigate to a route in current navigation tree.
@@ -203,7 +206,7 @@ type NavigationHelpersCommon<ParamList extends ParamListBase> = {
    *
    * @param state Navigation state object.
    */
-  reset(state: PartialState & { key?: string }): void;
+  reset(state: PartialState<State> & { key?: string }): void;
 
   /**
    * Go back to the previous route in history.
@@ -230,8 +233,9 @@ export type NavigationHelpers<
 export type NavigationProp<
   ParamList extends ParamListBase,
   RouteName extends keyof ParamList = string,
+  State extends NavigationState = NavigationState,
   ScreenOptions extends object = {}
-> = NavigationHelpersCommon<ParamList> & {
+> = NavigationHelpersCommon<ParamList, State> & {
   /**
    * Update the param object for the route.
    * The new params will be shallow merged with the old one.

src/useDescriptors.tsx

@@ -14,7 +14,7 @@ import NavigationBuilderContext, {
 } from './NavigationBuilderContext';
 
 type Options<ScreenOptions extends object> = {
-  state: NavigationState | PartialState;
+  state: NavigationState | PartialState<NavigationState>;
   screens: { [key: string]: RouteConfig<ParamListBase, string, ScreenOptions> };
   navigation: NavigationHelpers<ParamListBase>;
   onAction: (action: NavigationAction, sourceNavigatorKey?: string) => boolean;

src/useNavigationBuilder.tsx

@@ -5,9 +5,9 @@ import useRegisterNavigator from './useRegisterNavigator';
 import useDescriptors from './useDescriptors';
 import useNavigationHelpers from './useNavigationHelpers';
 import useOnAction from './useOnAction';
-import { Router, NavigationState, RouteConfig, ParamListBase } from './types';
 import useOnRouteFocus from './useOnRouteFocus';
 import useChildActionListeners from './useChildActionListeners';
+import { Router, NavigationState, RouteConfig, ParamListBase } from './types';
 
 type Options = {
   initialRouteName?: string;
@@ -49,10 +49,10 @@ const getRouteConfigsFromChildren = <ScreenOptions extends object>(
     );
   }, []);
 
-export default function useNavigationBuilder<ScreenOptions extends object>(
-  router: Router<any>,
-  options: Options
-) {
+export default function useNavigationBuilder<
+  State extends NavigationState,
+  ScreenOptions extends object
+>(router: Router<State, any>, options: Options) {
   useRegisterNavigator();
 
   const screens = getRouteConfigsFromChildren<ScreenOptions>(
@@ -96,7 +96,7 @@ export default function useNavigationBuilder<ScreenOptions extends object>(
 
   let state = router.getRehydratedState({
     routeNames,
-    partialState: currentState,
+    partialState: currentState as any,
   });
 
   if (!isArrayEqual(state.routeNames, routeNames)) {
@@ -130,10 +130,10 @@ export default function useNavigationBuilder<ScreenOptions extends object>(
   }, []);
 
   const getState = React.useCallback(
-    (): NavigationState =>
+    (): State =>
       router.getRehydratedState({
         routeNames,
-        partialState: getCurrentState() || state,
+        partialState: (getCurrentState() as any) || state,
       }),
     // eslint-disable-next-line react-hooks/exhaustive-deps
     [getCurrentState, router.getRehydratedState, router.getInitialState]

src/useOnAction.tsx

@@ -5,7 +5,7 @@ import NavigationBuilderContext, {
 import { NavigationAction, NavigationState, Router } from './types';
 
 type Options = {
-  router: Router<NavigationAction>;
+  router: Router<NavigationState, NavigationAction>;
   key?: string;
   getState: () => NavigationState;
   setState: (state: NavigationState) => void;

src/useOnRouteFocus.tsx

@@ -3,7 +3,7 @@ import { NavigationAction, NavigationState, Router } from './types';
 import NavigationBuilderContext from './NavigationBuilderContext';
 
 type Options<Action extends NavigationAction> = {
-  router: Router<Action>;
+  router: Router<NavigationState, Action>;
   onAction: (action: NavigationAction, sourceNavigatorKey?: string) => boolean;
   getState: () => NavigationState;
   setState: (state: NavigationState) => void;