fix: handle partial initial state better when rehydrating

packages/core/src/SceneView.tsx

@@ -9,12 +9,15 @@ import {
   NavigationState,
   NavigationProp,
   RouteConfig,
+  PartialState,
 } from './types';
 
 type Props<State extends NavigationState, ScreenOptions extends object> = {
   screen: RouteConfig<ParamListBase, string, ScreenOptions>;
   navigation: NavigationProp<ParamListBase, string, State, ScreenOptions>;
-  route: Route<string> & { state?: NavigationState };
+  route: Route<string> & {
+    state?: NavigationState | PartialState<NavigationState>;
+  };
   getState: () => NavigationState;
   setState: (state: NavigationState) => void;
 };

packages/core/src/__tests__/__fixtures__/MockRouter.tsx

@@ -3,6 +3,7 @@ import {
   Router,
   CommonAction,
   NavigationState,
+  Route,
   DefaultRouterOptions,
 } from '../../types';
 
@@ -30,19 +31,40 @@ export default function MockRouter(options: DefaultRouterOptions) {
       };
     },
 
-    getRehydratedState({ routeNames, partialState }) {
+    getRehydratedState(partialState, { routeNames, routeParamList }) {
       let state = partialState;
 
-      if (state.stale) {
-        state = {
-          ...state,
-          stale: false,
-          routeNames,
-          key: String(MockRouterKey.current++),
-        };
+      if (!state.stale) {
+        return state as NavigationState;
       }
 
-      return state;
+      const routes = state.routes
+        .filter(route => routeNames.includes(route.name))
+        .map(
+          route =>
+            ({
+              ...route,
+              key: route.key || `${route.name}-${MockRouterKey.current++}`,
+              params:
+                routeParamList[route.name] !== undefined
+                  ? {
+                      ...routeParamList[route.name],
+                      ...route.params,
+                    }
+                  : route.params,
+            } as Route<string>)
+        );
+
+      return {
+        stale: false,
+        key: String(MockRouterKey.current++),
+        index:
+          typeof state.index === 'number' && state.index < routes.length
+            ? state.index
+            : 0,
+        routeNames,
+        routes,
+      };
     },
 
     getStateForRouteNamesChange(state, { routeNames }) {

packages/core/src/types.tsx

@@ -27,19 +27,17 @@ export type NavigationState = {
   stale?: false;
 };
 
-export type InitialState = Omit<
-  NavigationState,
-  'key' | 'routes' | 'routeNames'
-> & {
-  key?: string;
-  routeNames?: string[];
-  routes: Array<Route<string> & { state?: InitialState }>;
+export type InitialState = Partial<Omit<NavigationState, 'routes'>> & {
+  routes: Array<Omit<Route<string>, 'key'> & { state?: InitialState }>;
 };
 
-export type PartialState<State extends NavigationState> = State & {
-  stale: true;
-  key?: undefined;
-  routeNames?: undefined;
+export type PartialState<State extends NavigationState> = Partial<
+  Omit<State, 'key' | 'routes' | 'routeNames'>
+> & {
+  stale: boolean;
+  routes: Array<
+    Omit<Route<string>, 'key'> & { key?: string; state?: InitialState }
+  >;
 };
 
 export type Route<RouteName extends string> = {
@@ -122,13 +120,17 @@ export type Router<
   /**
    * Rehydrate the full navigation state from a given partial state.
    *
+   * @param partialState Navigation state to rehydrate from.
    * @param options.routeNames List of valid route names as defined in the screen components.
-   * @param options.partialState Navigation state to rehydrate from.
+   * @param options.routeParamsList Object containing params for each route.
    */
-  getRehydratedState(options: {
-    routeNames: string[];
-    partialState: State | PartialState<State>;
-  }): State;
+  getRehydratedState(
+    partialState: PartialState<State>,
+    options: {
+      routeNames: string[];
+      routeParamList: ParamListBase;
+    }
+  ): State;
 
   /**
    * Take the current state and updated list of route names, and return a new state.

packages/core/src/useDescriptors.tsx

@@ -8,7 +8,6 @@ import { NavigationEventEmitter } from './useEventEmitter';
 import useNavigationCache from './useNavigationCache';
 import {
   Descriptor,
-  PartialState,
   NavigationAction,
   NavigationHelpers,
   NavigationState,
@@ -18,7 +17,7 @@ import {
 } from './types';
 
 type Options<ScreenOptions extends object> = {
-  state: NavigationState | PartialState<NavigationState>;
+  state: NavigationState;
   screens: { [key: string]: RouteConfig<ParamListBase, string, ScreenOptions> };
   navigation: NavigationHelpers<ParamListBase>;
   screenOptions?: ScreenOptions;

packages/core/src/useNavigationBuilder.tsx

@@ -19,6 +19,7 @@ import {
   RouteConfig,
   Router,
   RouterFactory,
+  PartialState,
 } from './types';
 
 /**
@@ -137,9 +138,9 @@ export default function useNavigationBuilder<
           routeNames,
           routeParamList,
         })
-      : router.getRehydratedState({
+      : router.getRehydratedState(currentState as PartialState<State>, {
           routeNames,
-          partialState: currentState as any,
+          routeParamList,
         })
   );
 

packages/routers/src/DrawerRouter.tsx

@@ -68,6 +68,25 @@ export default function DrawerRouter(
       };
     },
 
+    getRehydratedState(partialState, { routeNames, routeParamList }) {
+      if (!partialState.stale) {
+        return partialState as DrawerNavigationState;
+      }
+
+      const state = router.getRehydratedState(partialState, {
+        routeNames,
+        routeParamList,
+      });
+
+      return {
+        ...state,
+        isDrawerOpen:
+          typeof partialState.isDrawerOpen === 'boolean'
+            ? partialState.isDrawerOpen
+            : false,
+      };
+    },
+
     getStateForRouteFocus(state, key) {
       const index = state.routes.findIndex(r => r.key === key);
 

packages/routers/src/StackRouter.tsx

@@ -5,6 +5,7 @@ import {
   Router,
   BaseRouter,
   DefaultRouterOptions,
+  Route,
 } from '@navigation-ex/core';
 
 export type StackActionType =
@@ -47,36 +48,69 @@ export default function StackRouter(options: StackRouterOptions) {
     ...BaseRouter,
 
     getInitialState({ routeNames, routeParamList }) {
-      const index =
-        options.initialRouteName === undefined
-          ? 0
-          : routeNames.indexOf(options.initialRouteName);
+      const initialRouteName =
+        options.initialRouteName !== undefined
+          ? options.initialRouteName
+          : routeNames[0];
 
       return {
         key: `stack-${shortid()}`,
-        index,
+        index: 0,
         routeNames,
-        routes: routeNames.slice(0, index + 1).map(name => ({
-          name,
-          key: `${name}-${shortid()}`,
-          params: routeParamList[name],
-        })),
+        routes: [
+          {
+            key: `${initialRouteName}-${shortid()}`,
+            name: initialRouteName,
+            params: routeParamList[initialRouteName],
+          },
+        ],
       };
     },
 
-    getRehydratedState({ routeNames, partialState }) {
+    getRehydratedState(partialState, { routeNames, routeParamList }) {
       let state = partialState;
 
-      if (state.stale) {
-        state = {
-          ...state,
-          stale: false,
-          routeNames,
-          key: `stack-${shortid()}`,
-        };
+      if (!state.stale) {
+        return state as StackNavigationState;
       }
 
-      return state;
+      const routes = state.routes
+        .filter(route => routeNames.includes(route.name))
+        .map(
+          route =>
+            ({
+              ...route,
+              key: route.key || `${route.name}-${shortid()}`,
+              params:
+                routeParamList[route.name] !== undefined
+                  ? {
+                      ...routeParamList[route.name],
+                      ...route.params,
+                    }
+                  : route.params,
+            } as Route<string>)
+        );
+
+      if (routes.length === 0) {
+        const initialRouteName =
+          options.initialRouteName !== undefined
+            ? options.initialRouteName
+            : routeNames[0];
+
+        routes.push({
+          key: `${initialRouteName}-${shortid()}`,
+          name: initialRouteName,
+          params: routeParamList[initialRouteName],
+        });
+      }
+
+      return {
+        stale: false,
+        key: `stack-${shortid()}`,
+        index: routes.length - 1,
+        routeNames,
+        routes,
+      };
     },
 
     getStateForRouteNamesChange(state, { routeNames }) {

packages/routers/src/TabRouter.tsx

@@ -71,19 +71,46 @@ export default function TabRouter({
       };
     },
 
-    getRehydratedState({ routeNames, partialState }) {
+    getRehydratedState(partialState, { routeNames, routeParamList }) {
       let state = partialState;
 
-      if (state.stale) {
-        state = {
-          ...state,
-          stale: false,
-          routeNames,
-          key: `tab-${shortid()}`,
-        };
+      if (!state.stale) {
+        return state as TabNavigationState;
       }
 
-      return state;
+      const routes = routeNames.map(name => {
+        const route = state.routes.find(r => r.name === name);
+
+        return {
+          name,
+          key: route && route.key ? route.key : `${name}-${shortid()}`,
+          params:
+            routeParamList[name] !== undefined
+              ? {
+                  ...routeParamList[name],
+                  ...(route ? route.params : undefined),
+                }
+              : route
+              ? route.params
+              : undefined,
+        };
+      });
+
+      const routeKeyHistory = state.routeKeyHistory
+        ? state.routeKeyHistory.filter(key => routes.find(r => r.key === key))
+        : [];
+
+      return {
+        stale: false,
+        key: `tab-${shortid()}`,
+        index:
+          typeof state.index === 'number' && state.index < routes.length
+            ? state.index
+            : 0,
+        routeNames,
+        routeKeyHistory,
+        routes,
+      };
     },
 
     getStateForRouteNamesChange(state, { routeNames, routeParamList }) {