feat: add a getRootState method (#119)

packages/core/src/NavigationBuilderContext.tsx

@@ -1,5 +1,10 @@
 import * as React from 'react';
-import { NavigationAction, NavigationHelpers, ParamListBase } from './types';
+import {
+  NavigationAction,
+  NavigationHelpers,
+  NavigationState,
+  ParamListBase,
+} from './types';
 
 export type ChildActionListener = (
   action: NavigationAction,
@@ -14,6 +19,8 @@ export type FocusedNavigationListener = <T>(
   callback: FocusedNavigationCallback<T>
 ) => { handled: boolean; result: T };
 
+export type NavigatorStateGetter = () => NavigationState;
+
 /**
  * Context which holds the required helpers needed to build nested navigators.
  */
@@ -25,6 +32,7 @@ const NavigationBuilderContext = React.createContext<{
   addActionListener?: (listener: ChildActionListener) => void;
   addFocusedListener?: (listener: FocusedNavigationListener) => void;
   onRouteFocus?: (key: string) => void;
+  addStateGetter?: (key: string, getter: NavigatorStateGetter) => void;
   trackAction: (action: NavigationAction) => void;
 }>({
   trackAction: () => undefined,

packages/core/src/NavigationContainer.tsx

@@ -5,6 +5,7 @@ import NavigationBuilderContext from './NavigationBuilderContext';
 import ResetRootContext from './ResetRootContext';
 import useFocusedListeners from './useFocusedListeners';
 import useDevTools from './useDevTools';
+import useStateGetters from './useStateGetters';
 
 import {
   Route,
@@ -108,6 +109,8 @@ const Container = React.forwardRef(function NavigationContainer(
 
   const { listeners, addListener: addFocusedListener } = useFocusedListeners();
 
+  const { getStateForRoute, addStateGetter } = useStateGetters();
+
   const dispatch = (
     action: NavigationAction | ((state: NavigationState) => NavigationAction)
   ) => {
@@ -134,6 +137,10 @@ const Container = React.forwardRef(function NavigationContainer(
     [trackAction]
   );
 
+  const getRootState = () => {
+    return getStateForRoute('root');
+  };
+
   React.useImperativeHandle(ref, () => ({
     ...(Object.keys(CommonActions) as Array<keyof typeof CommonActions>).reduce<
       any
@@ -150,14 +157,16 @@ const Container = React.forwardRef(function NavigationContainer(
     resetRoot,
     dispatch,
     canGoBack,
+    getRootState,
   }));
 
   const builderContext = React.useMemo(
     () => ({
       addFocusedListener,
+      addStateGetter,
       trackAction,
     }),
-    [addFocusedListener, trackAction]
+    [addFocusedListener, trackAction, addStateGetter]
   );
 
   const performTransaction = React.useCallback((callback: () => void) => {

packages/core/src/__tests__/NavigationContainer.test.tsx

@@ -321,3 +321,56 @@ it('handle resetting state with ref', () => {
   expect(onStateChange).toBeCalledTimes(1);
   expect(onStateChange).lastCalledWith(state);
 });
+
+it('handle getRootState', () => {
+  const TestNavigator = (props: any) => {
+    const { state, descriptors } = useNavigationBuilder(MockRouter, props);
+
+    return descriptors[state.routes[state.index].key].render();
+  };
+
+  const ref = React.createRef<NavigationContainerRef>();
+
+  const element = (
+    <NavigationContainer ref={ref}>
+      <TestNavigator initialRouteName="foo">
+        <Screen name="foo">
+          {() => (
+            <TestNavigator>
+              <Screen name="qux" component={() => null} />
+              <Screen name="lex" component={() => null} />
+            </TestNavigator>
+          )}
+        </Screen>
+        <Screen name="bar" component={() => null} />
+      </TestNavigator>
+    </NavigationContainer>
+  );
+
+  render(element);
+
+  let state;
+  if (ref.current) {
+    state = ref.current.getRootState();
+  }
+  expect(state).toEqual({
+    index: 0,
+    key: '7',
+    routeNames: ['foo', 'bar'],
+    routes: [
+      {
+        key: 'foo',
+        name: 'foo',
+        state: {
+          index: 0,
+          key: '8',
+          routeNames: ['qux', 'lex'],
+          routes: [{ key: 'qux', name: 'qux' }, { key: 'lex', name: 'lex' }],
+          stale: false,
+        },
+      },
+      { key: 'bar', name: 'bar' },
+    ],
+    stale: false,
+  });
+});

packages/core/src/types.tsx

@@ -538,6 +538,7 @@ export type NavigationContainerRef =
        * @param state Navigation state object.
        */
       resetRoot(state: PartialState<NavigationState> | NavigationState): void;
+      getRootState(): NavigationState;
     }
   | undefined
   | null;

packages/core/src/useDescriptors.tsx

@@ -3,6 +3,7 @@ import SceneView from './SceneView';
 import NavigationBuilderContext, {
   ChildActionListener,
   FocusedNavigationListener,
+  NavigatorStateGetter,
 } from './NavigationBuilderContext';
 import { NavigationEventEmitter } from './useEventEmitter';
 import useNavigationCache from './useNavigationCache';
@@ -35,6 +36,7 @@ type Options<State extends NavigationState, ScreenOptions extends object> = {
   setState: (state: State) => void;
   addActionListener: (listener: ChildActionListener) => void;
   addFocusedListener: (listener: FocusedNavigationListener) => void;
+  addStateGetter: (key: string, getter: NavigatorStateGetter) => void;
   onRouteFocus: (key: string) => void;
   router: Router<State, NavigationAction>;
   emitter: NavigationEventEmitter;
@@ -61,6 +63,7 @@ export default function useDescriptors<
   setState,
   addActionListener,
   addFocusedListener,
+  addStateGetter,
   onRouteFocus,
   router,
   emitter,
@@ -74,6 +77,7 @@ export default function useDescriptors<
       onAction,
       addActionListener,
       addFocusedListener,
+      addStateGetter,
       onRouteFocus,
       trackAction,
     }),
@@ -83,6 +87,7 @@ export default function useDescriptors<
       addActionListener,
       addFocusedListener,
       onRouteFocus,
+      addStateGetter,
       trackAction,
     ]
   );

packages/core/src/useNavigationBuilder.tsx

@@ -23,6 +23,8 @@ import {
   PrivateValueStore,
   NavigationAction,
 } from './types';
+import useStateGetters from './useStateGetters';
+import useOnGetState from './useOnGetState';
 
 // This is to make TypeScript compiler happy
 // eslint-disable-next-line babel/no-unused-expressions
@@ -227,6 +229,8 @@ export default function useNavigationBuilder<
     addListener: addFocusedListener,
   } = useFocusedListeners();
 
+  const { getStateForRoute, addStateGetter } = useStateGetters();
+
   const onAction = useOnAction({
     router,
     getState,
@@ -254,6 +258,11 @@ export default function useNavigationBuilder<
     focusedListeners,
   });
 
+  useOnGetState({
+    getState,
+    getStateForRoute,
+  });
+
   const descriptors = useDescriptors<State, ScreenOptions>({
     state,
     screens,
@@ -265,6 +274,7 @@ export default function useNavigationBuilder<
     onRouteFocus,
     addActionListener,
     addFocusedListener,
+    addStateGetter,
     router,
     emitter,
   });

packages/core/src/useOnGetState.tsx

@@ -0,0 +1,31 @@
+import * as React from 'react';
+import NavigationBuilderContext from './NavigationBuilderContext';
+import { NavigationState } from './types';
+import NavigationRouteContext from './NavigationRouteContext';
+
+export default function useOnGetState({
+  getStateForRoute,
+  getState,
+}: {
+  getStateForRoute: (routeName: string) => NavigationState | undefined;
+  getState: () => NavigationState;
+}) {
+  const { addStateGetter } = React.useContext(NavigationBuilderContext);
+  const route = React.useContext(NavigationRouteContext);
+  const key = route ? route.key : 'root';
+
+  const getter = React.useCallback(() => {
+    const state = getState();
+    return {
+      ...state,
+      routes: state.routes.map(route => ({
+        ...route,
+        state: getStateForRoute(route.key),
+      })),
+    };
+  }, [getState, getStateForRoute]);
+
+  React.useEffect(() => {
+    return addStateGetter && addStateGetter(key, getter);
+  }, [addStateGetter, getter, key]);
+}

packages/core/src/useStateGetters.tsx

@@ -0,0 +1,35 @@
+import * as React from 'react';
+import { NavigatorStateGetter } from './NavigationBuilderContext';
+
+/**
+ * Hook which lets child navigators add getters to be called for obtaining rehydrated state.
+ */
+
+export default function useStateGetters() {
+  const stateGetters = React.useRef<Record<string, NavigatorStateGetter>>({});
+
+  const getStateForRoute = React.useCallback(
+    routeKey =>
+      stateGetters.current[routeKey] === undefined
+        ? undefined
+        : stateGetters.current[routeKey](),
+    [stateGetters]
+  );
+
+  const addStateGetter = React.useCallback(
+    (key: string, getter: NavigatorStateGetter) => {
+      stateGetters.current[key] = getter;
+
+      return () => {
+        // @ts-ignore
+        stateGetters.current[key] = undefined;
+      };
+    },
+    []
+  );
+
+  return {
+    getStateForRoute,
+    addStateGetter,
+  };
+}