refactor: use a regular action for 'resetRoot'

Previously, 'resetRoot' directly performed a 'setState' on the container instead of dispatching an action. This meant that hooks such as listener for 'preventRemove' won't be notified by it. This commit changes it to dispatch a regular 'RESET' action which will behave the same as other actions.

example/src/Screens/PreventRemove.tsx

@@ -77,18 +77,28 @@ const InputScreen = ({
 
         e.preventDefault();
 
-        Alert.alert(
-          'Discard changes?',
-          'You have unsaved changes. Are you sure to discard them and leave the screen?',
-          [
-            { text: "Don't leave", style: 'cancel', onPress: () => {} },
-            {
-              text: 'Discard',
-              style: 'destructive',
-              onPress: () => navigation.dispatch(action),
-            },
-          ]
-        );
+        if (Platform.OS === 'web') {
+          const discard = confirm(
+            'You have unsaved changes. Discard them and leave the screen?'
+          );
+
+          if (discard) {
+            navigation.dispatch(action);
+          }
+        } else {
+          Alert.alert(
+            'Discard changes?',
+            'You have unsaved changes. Discard them and leave the screen?',
+            [
+              { text: "Don't leave", style: 'cancel', onPress: () => {} },
+              {
+                text: 'Discard',
+                style: 'destructive',
+                onPress: () => navigation.dispatch(action),
+              },
+            ]
+          );
+        }
       }),
     [hasUnsavedChanges, navigation]
   );

packages/compat/src/createCompatNavigationProp.tsx

@@ -147,7 +147,6 @@ export default function createCompatNavigationProp<
       }
     },
     state: {
-      // @ts-expect-error: these properties may actually exist
       key: state.key,
       // @ts-expect-error
       routeName: state.name,
@@ -202,7 +201,6 @@ export default function createCompatNavigationProp<
 
       const { routes } = navigation.dangerouslyGetState();
 
-      // @ts-expect-error
       return routes[0].key === state.key;
     },
     dangerouslyGetParent() {

packages/core/src/BaseNavigationContainer.tsx

@@ -160,9 +160,20 @@ const BaseNavigationContainer = React.forwardRef(
 
     const resetRoot = React.useCallback(
       (state?: PartialState<NavigationState> | NavigationState) => {
-        setState(state);
+        const target = state?.key ?? keyedListeners.getState.root?.().key;
+
+        if (target == null) {
+          throw new Error(NOT_INITIALIZED_ERROR);
+        }
+
+        listeners.focus[0]((navigation) =>
+          navigation.dispatch({
+            ...CommonActions.reset(state),
+            target,
+          })
+        );
       },
-      [setState]
+      [keyedListeners.getState, listeners.focus]
     );
 
     const getRootState = React.useCallback(() => {

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

@@ -1178,3 +1178,149 @@ it("prevents removing by multiple screens with 'beforeRemove' event", () => {
     type: 'stack',
   });
 });
+
+it("prevents removing a child screen with 'beforeRemove' event with 'resetRoot'", () => {
+  const TestNavigator = (props: any) => {
+    const { state, descriptors } = useNavigationBuilder(StackRouter, props);
+
+    return (
+      <React.Fragment>
+        {state.routes.map((route) => descriptors[route.key].render())}
+      </React.Fragment>
+    );
+  };
+
+  const onBeforeRemove = jest.fn();
+
+  let shouldPrevent = true;
+  let shouldContinue = false;
+
+  const TestScreen = (props: any) => {
+    React.useEffect(
+      () =>
+        props.navigation.addListener('beforeRemove', (e: any) => {
+          onBeforeRemove();
+
+          if (shouldPrevent) {
+            e.preventDefault();
+
+            if (shouldContinue) {
+              props.navigation.dispatch(e.data.action);
+            }
+          }
+        }),
+      [props.navigation]
+    );
+
+    return null;
+  };
+
+  const onStateChange = jest.fn();
+
+  const ref = React.createRef<NavigationContainerRef>();
+
+  const element = (
+    <BaseNavigationContainer ref={ref} onStateChange={onStateChange}>
+      <TestNavigator>
+        <Screen name="foo">{() => null}</Screen>
+        <Screen name="bar">{() => null}</Screen>
+        <Screen name="baz">
+          {() => (
+            <TestNavigator>
+              <Screen name="qux" component={TestScreen} />
+              <Screen name="lex">{() => null}</Screen>
+            </TestNavigator>
+          )}
+        </Screen>
+      </TestNavigator>
+    </BaseNavigationContainer>
+  );
+
+  render(element);
+
+  act(() => ref.current?.navigate('baz'));
+
+  expect(onStateChange).toBeCalledTimes(1);
+  expect(onStateChange).toBeCalledWith({
+    index: 1,
+    key: 'stack-2',
+    routeNames: ['foo', 'bar', 'baz'],
+    routes: [
+      { key: 'foo-3', name: 'foo' },
+      {
+        key: 'baz-4',
+        name: 'baz',
+        state: {
+          index: 0,
+          key: 'stack-6',
+          routeNames: ['qux', 'lex'],
+          routes: [{ key: 'qux-7', name: 'qux' }],
+          stale: false,
+          type: 'stack',
+        },
+      },
+    ],
+    stale: false,
+    type: 'stack',
+  });
+
+  act(() =>
+    ref.current?.resetRoot({
+      index: 0,
+      key: 'stack-2',
+      routeNames: ['foo', 'bar', 'baz'],
+      routes: [{ key: 'foo-3', name: 'foo' }],
+      stale: false,
+      type: 'stack',
+    })
+  );
+
+  expect(onStateChange).toBeCalledTimes(1);
+  expect(onBeforeRemove).toBeCalledTimes(1);
+
+  expect(ref.current?.getRootState()).toEqual({
+    index: 1,
+    key: 'stack-2',
+    routeNames: ['foo', 'bar', 'baz'],
+    routes: [
+      { key: 'foo-3', name: 'foo' },
+      {
+        key: 'baz-4',
+        name: 'baz',
+        state: {
+          index: 0,
+          key: 'stack-6',
+          routeNames: ['qux', 'lex'],
+          routes: [{ key: 'qux-7', name: 'qux' }],
+          stale: false,
+          type: 'stack',
+        },
+      },
+    ],
+    stale: false,
+    type: 'stack',
+  });
+
+  shouldPrevent = false;
+
+  act(() =>
+    ref.current?.resetRoot({
+      index: 0,
+      key: 'stack-2',
+      routeNames: ['foo', 'bar', 'baz'],
+      routes: [{ key: 'foo-3', name: 'foo' }],
+      stale: false,
+      type: 'stack',
+    })
+  );
+
+  expect(onStateChange).toBeCalledTimes(2);
+  expect(onStateChange).toBeCalledWith({
+    index: 0,
+    key: 'stack-2',
+    routeNames: ['foo', 'bar', 'baz'],
+    routes: [{ key: 'foo-3', name: 'foo' }],
+    stale: false,
+    type: 'stack',
+  });
+});

packages/core/src/useOnAction.tsx

@@ -90,18 +90,11 @@ export default function useOnAction({
           onDispatchAction(action, state === result);
 
           if (state !== result) {
-            const nextRouteKeys = (result.routes as any[]).map(
-              (route: { key?: string }) => route.key
-            );
-
-            const removedRoutes = state.routes.filter(
-              (route) => !nextRouteKeys.includes(route.key)
-            );
-
             const isPrevented = shouldPreventRemove(
               emitter,
               beforeRemoveListeners,
-              removedRoutes,
+              state.routes,
+              result.routes,
               action
             );
 

packages/core/src/useOnPreventRemove.tsx

@@ -1,7 +1,6 @@
 import * as React from 'react';
 import type {
   NavigationState,
-  Route,
   NavigationAction,
 } from '@react-navigation/routers';
 import NavigationBuilderContext, {
@@ -22,11 +21,16 @@ const VISITED_ROUTE_KEYS = Symbol('VISITED_ROUTE_KEYS');
 export const shouldPreventRemove = (
   emitter: NavigationEventEmitter<EventMapCore<any>>,
   beforeRemoveListeners: Record<string, ChildBeforeRemoveListener | undefined>,
-  routes: Route<string>[],
+  currentRoutes: { key: string }[],
+  nextRoutes: { key?: string | undefined }[],
   action: NavigationAction
 ) => {
+  const nextRouteKeys = nextRoutes.map((route) => route.key);
+
   // Call these in reverse order so last screens handle the event first
-  const reversedRoutes = [...routes].reverse();
+  const removedRoutes = currentRoutes
+    .filter((route) => !nextRouteKeys.includes(route.key))
+    .reverse();
 
   const visitedRouteKeys: Set<string> =
     // @ts-expect-error: add this property to mark that we've already emitted this action
@@ -37,7 +41,7 @@ export const shouldPreventRemove = (
     [VISITED_ROUTE_KEYS]: visitedRouteKeys,
   };
 
-  for (const route of reversedRoutes) {
+  for (const route of removedRoutes) {
     if (visitedRouteKeys.has(route.key)) {
       // Skip if we've already emitted this action for this screen
       continue;
@@ -85,6 +89,7 @@ export default function useOnPreventRemove({
           emitter,
           beforeRemoveListeners,
           state.routes,
+          [],
           action
         );
       });

packages/routers/src/CommonActions.tsx

@@ -23,7 +23,7 @@ export type Action =
     }
   | {
       type: 'RESET';
-      payload: ResetState;
+      payload: ResetState | undefined;
       source?: string;
       target?: string;
     }
@@ -62,7 +62,7 @@ export function navigate(...args: any): Action {
   }
 }
 
-export function reset(state: ResetState): Action {
+export function reset(state: ResetState | undefined): Action {
   return { type: 'RESET', payload: state };
 }
 

packages/routers/src/types.tsx

@@ -56,12 +56,11 @@ export type PartialRoute<R extends Route<string>> = Omit<R, 'key'> & {
 };
 
 export type PartialState<State extends NavigationState> = Partial<
-  Omit<State, 'stale' | 'type' | 'key' | 'routes' | 'routeNames'>
+  Omit<State, 'stale' | 'routes'>
 > &
   Readonly<{
     stale?: true;
-    type?: string;
-    routes: PartialRoute<Route<string>>[];
+    routes: PartialRoute<Route<State['routeNames'][number]>>[];
   }>;
 
 export type Route<