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,17 +4,18 @@ import useNavigationBuilder from '../useNavigationBuilder';
 import NavigationContainer from '../NavigationContainer';
 import Screen from '../Screen';
 import MockRouter from './__fixtures__/MockRouter';
+import { NavigationState } from '../types';
 
 jest.useFakeTimers();
 
 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,15 +78,15 @@ 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() {
       return true;
     },
   };
 
-  const ChildRouter: Router<MockActions> = {
+  const ChildRouter: Router<NavigationState, MockActions> = {
     ...MockRouter,
 
     shouldActionChangeFocus() {