refactor: make dispatch accept thunks

packages/compat/src/createCompatNavigationProp.tsx

@@ -28,18 +28,9 @@ export default function createCompatNavigationProp<
     ...navigation,
     ...Object.entries(helpers).reduce<{
       [key: string]: (...args: any[]) => void;
-    }>((acc, [name, method]) => {
+    }>((acc, [name, method]: [string, Function]) => {
       if (name in navigation) {
-        acc[name] = (...args: any[]) => {
-          // @ts-ignore
-          const payload = method(...args);
-
-          navigation.dispatch(
-            typeof payload === 'function'
-              ? payload(navigation.dangerouslyGetState())
-              : payload
-          );
-        };
+        acc[name] = (...args: any[]) => navigation.dispatch(method(...args));
       }
 
       return acc;

packages/core/src/NavigationContainer.tsx

@@ -88,7 +88,9 @@ const Container = React.forwardRef(function NavigationContainer(
 
   const { listeners, addListener: addFocusedListener } = useFocusedListeners();
 
-  const dispatch = (action: NavigationAction) => {
+  const dispatch = (
+    action: NavigationAction | ((state: NavigationState) => NavigationAction)
+  ) => {
     listeners[0](navigation => navigation.dispatch(action));
   };
 

packages/core/src/SceneView.tsx

@@ -19,8 +19,8 @@ type Props<State extends NavigationState, ScreenOptions extends object> = {
   route: Route<string> & {
     state?: NavigationState | PartialState<NavigationState>;
   };
-  getState: () => NavigationState;
-  setState: (state: NavigationState) => void;
+  getState: () => State;
+  setState: (state: State) => void;
 };
 
 /**

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

@@ -323,7 +323,7 @@ it('cleans up state when the navigator unmounts', () => {
   expect(onStateChange).lastCalledWith(undefined);
 });
 
-it('allows arbitrary state updates by dispatching a function', () => {
+it('allows state updates by dispatching a function returning an action', () => {
   const TestNavigator = (props: any) => {
     const { state, descriptors } = useNavigationBuilder(MockRouter, props);
 
@@ -332,11 +332,11 @@ it('allows arbitrary state updates by dispatching a function', () => {
 
   const FooScreen = (props: any) => {
     React.useEffect(() => {
-      props.navigation.dispatch((state: any) => ({
-        ...state,
-        routes: state.routes.slice().reverse(),
-        index: 1,
-      }));
+      props.navigation.dispatch((state: NavigationState) =>
+        state.index === 0
+          ? { type: 'NAVIGATE', payload: { name: state.routeNames[1] } }
+          : { type: 'NOOP' }
+      );
       // eslint-disable-next-line react-hooks/exhaustive-deps
     }, []);
 
@@ -364,7 +364,7 @@ it('allows arbitrary state updates by dispatching a function', () => {
     index: 1,
     key: '0',
     routeNames: ['foo', 'bar'],
-    routes: [{ key: 'bar', name: 'bar' }, { key: 'foo', name: 'foo' }],
+    routes: [{ key: 'foo', name: 'foo' }, { key: 'bar', name: 'bar' }],
   });
 });
 

packages/core/src/types.tsx

@@ -270,7 +270,9 @@ type NavigationHelpersCommon<
    *
    * @param action Action object or update function.
    */
-  dispatch(action: NavigationAction | ((state: State) => State)): void;
+  dispatch(
+    action: NavigationAction | ((state: State) => NavigationAction)
+  ): void;
 
   /**
    * Navigate to a route in current navigation tree.

packages/core/src/useDescriptors.tsx

@@ -17,8 +17,8 @@ import {
   Router,
 } from './types';
 
-type Options<ScreenOptions extends object> = {
-  state: NavigationState;
+type Options<State extends NavigationState, ScreenOptions extends object> = {
+  state: State;
   screens: { [key: string]: RouteConfig<ParamListBase, string, ScreenOptions> };
   navigation: NavigationHelpers<ParamListBase>;
   screenOptions?:
@@ -31,12 +31,12 @@ type Options<ScreenOptions extends object> = {
     action: NavigationAction,
     visitedNavigators?: Set<string>
   ) => boolean;
-  getState: () => NavigationState;
-  setState: (state: NavigationState) => void;
+  getState: () => State;
+  setState: (state: State) => void;
   addActionListener: (listener: ChildActionListener) => void;
   addFocusedListener: (listener: FocusedNavigationListener) => void;
   onRouteFocus: (key: string) => void;
-  router: Router<NavigationState, NavigationAction>;
+  router: Router<State, NavigationAction>;
   emitter: NavigationEventEmitter;
 };
 
@@ -64,7 +64,7 @@ export default function useDescriptors<
   onRouteFocus,
   router,
   emitter,
-}: Options<ScreenOptions>) {
+}: Options<State, ScreenOptions>) {
   const [options, setOptions] = React.useState<{ [key: string]: object }>({});
   const { trackAction } = React.useContext(NavigationBuilderContext);
 

packages/core/src/useNavigationBuilder.tsx

@@ -245,7 +245,6 @@ export default function useNavigationBuilder<
   const navigation = useNavigationHelpers<State, NavigationAction, EventMap>({
     onAction,
     getState,
-    setState,
     emitter,
     router,
   });

packages/core/src/useNavigationCache.tsx

@@ -12,15 +12,15 @@ import {
   Router,
 } from './types';
 
-type Options = {
-  state: NavigationState;
-  getState: () => NavigationState;
+type Options<State extends NavigationState> = {
+  state: State;
+  getState: () => State;
   navigation: NavigationHelpers<ParamListBase> &
     Partial<NavigationProp<ParamListBase, string, any, any, any>>;
   setOptions: (
     cb: (options: { [key: string]: object }) => { [key: string]: object }
   ) => void;
-  router: Router<NavigationState, NavigationAction>;
+  router: Router<State, NavigationAction>;
   emitter: NavigationEventEmitter;
 };
 
@@ -39,7 +39,14 @@ type NavigationCache<
 export default function useNavigationCache<
   State extends NavigationState,
   ScreenOptions extends object
->({ state, getState, navigation, setOptions, router, emitter }: Options) {
+>({
+  state,
+  getState,
+  navigation,
+  setOptions,
+  router,
+  emitter,
+}: Options<State>) {
   // Cache object which holds navigation objects for each screen
   // We use `React.useMemo` instead of `React.useRef` coz we want to invalidate it when deps change
   // In reality, these deps will rarely change, if ever
@@ -70,13 +77,17 @@ export default function useNavigationCache<
         const { emit, ...rest } = navigation;
 
         const dispatch = (
-          action: NavigationAction | ((state: State) => State)
-        ) =>
+          action: NavigationAction | ((state: State) => NavigationAction)
+        ) => {
+          const payload =
+            typeof action === 'function' ? action(getState()) : action;
+
           navigation.dispatch(
-            typeof action === 'object' && action != null
-              ? { source: route.key, ...action }
-              : action
+            typeof payload === 'object' && payload != null
+              ? { source: route.key, ...payload }
+              : payload
           );
+        };
 
         const helpers = Object.keys(actions).reduce(
           (acc, name) => {
@@ -92,7 +103,7 @@ export default function useNavigationCache<
           ...helpers,
           ...emitter.create(route.key),
           dangerouslyGetParent: () => parentNavigation as any,
-          dangerouslyGetState: getState as () => State,
+          dangerouslyGetState: getState,
           dispatch,
           setOptions: (options: object) =>
             setOptions(o => ({

packages/core/src/useNavigationHelpers.tsx

@@ -23,7 +23,6 @@ type Options<State extends NavigationState, Action extends NavigationAction> = {
     visitedNavigators?: Set<string>
   ) => boolean;
   getState: () => State;
-  setState: (state: State) => void;
   emitter: NavigationEventEmitter;
   router: Router<State, Action>;
 };
@@ -36,18 +35,17 @@ export default function useNavigationHelpers<
   State extends NavigationState,
   Action extends NavigationAction,
   EventMap extends { [key: string]: any }
->({ onAction, getState, setState, emitter, router }: Options<State, Action>) {
+>({ onAction, getState, emitter, router }: Options<State, Action>) {
   const parentNavigationHelpers = React.useContext(NavigationContext);
   const { performTransaction } = React.useContext(NavigationStateContext);
 
   return React.useMemo(() => {
-    const dispatch = (action: Action | ((state: State) => State)) =>
+    const dispatch = (action: Action | ((state: State) => Action)) =>
       performTransaction(() => {
-        if (typeof action === 'function') {
-          setState(action(getState()));
-        } else {
-          onAction(action);
-        }
+        const payload =
+          typeof action === 'function' ? action(getState()) : action;
+
+        onAction(payload);
       });
 
     const actions = {
@@ -87,7 +85,6 @@ export default function useNavigationHelpers<
     parentNavigationHelpers,
     emitter.emit,
     performTransaction,
-    setState,
     onAction,
   ]);
 }