refactor: add separate method for rehydration and fix types

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) || '');