feat: expose container ref in useNavigation

packages/core/src/BaseNavigationContainer.tsx

@@ -29,6 +29,7 @@ import type {
   NavigationContainerRef,
   NavigationContainerProps,
 } from './types';
+import NavigationContainerRefContext from './NavigationContainerRefContext';
 
 type State = NavigationState | PartialState<NavigationState> | undefined;
 
@@ -136,17 +137,22 @@ const BaseNavigationContainer = React.forwardRef(
 
     const { keyedListeners, addKeyedListener } = useKeyedChildListeners();
 
-    const dispatch = (
-      action: NavigationAction | ((state: NavigationState) => NavigationAction)
-    ) => {
-      if (listeners.focus[0] == null) {
-        console.error(NOT_INITIALIZED_ERROR);
-      } else {
-        listeners.focus[0]((navigation) => navigation.dispatch(action));
-      }
-    };
+    const dispatch = React.useCallback(
+      (
+        action:
+          | NavigationAction
+          | ((state: NavigationState) => NavigationAction)
+      ) => {
+        if (listeners.focus[0] == null) {
+          console.error(NOT_INITIALIZED_ERROR);
+        } else {
+          listeners.focus[0]((navigation) => navigation.dispatch(action));
+        }
+      },
+      [listeners.focus]
+    );
 
-    const canGoBack = () => {
+    const canGoBack = React.useCallback(() => {
       if (listeners.focus[0] == null) {
         return false;
       }
@@ -160,7 +166,7 @@ const BaseNavigationContainer = React.forwardRef(
       } else {
         return false;
       }
-    };
+    }, [listeners.focus]);
 
     const resetRoot = React.useCallback(
       (state?: PartialState<NavigationState> | NavigationState) => {
@@ -200,24 +206,38 @@ const BaseNavigationContainer = React.forwardRef(
 
     const { addOptionsGetter, getCurrentOptions } = useOptionsGetters({});
 
-    React.useImperativeHandle(ref, () => ({
-      ...Object.keys(CommonActions).reduce<any>((acc, name) => {
-        acc[name] = (...args: any[]) =>
-          // @ts-expect-error: this is ok
-          dispatch(CommonActions[name](...args));
-        return acc;
-      }, {}),
-      ...emitter.create('root'),
-      resetRoot,
-      dispatch,
-      canGoBack,
-      getRootState,
-      getState: () => state,
-      getParent: () => undefined,
-      getCurrentRoute,
-      getCurrentOptions,
-      isReady: () => listeners.focus[0] != null,
-    }));
+    const navigation: NavigationContainerRef<ParamListBase> = React.useMemo(
+      () => ({
+        ...Object.keys(CommonActions).reduce<any>((acc, name) => {
+          acc[name] = (...args: any[]) =>
+            // @ts-expect-error: this is ok
+            dispatch(CommonActions[name](...args));
+          return acc;
+        }, {}),
+        ...emitter.create('root'),
+        resetRoot,
+        dispatch,
+        canGoBack,
+        getRootState,
+        getState: () => stateRef.current,
+        getParent: () => undefined,
+        getCurrentRoute,
+        getCurrentOptions,
+        isReady: () => listeners.focus[0] != null,
+      }),
+      [
+        canGoBack,
+        dispatch,
+        emitter,
+        getCurrentOptions,
+        getCurrentRoute,
+        getRootState,
+        listeners.focus,
+        resetRoot,
+      ]
+    );
+
+    React.useImperativeHandle(ref, () => navigation, [navigation]);
 
     const onDispatchAction = React.useCallback(
       (action: NavigationAction, noop: boolean) => {
@@ -285,10 +305,12 @@ const BaseNavigationContainer = React.forwardRef(
     );
 
     const onStateChangeRef = React.useRef(onStateChange);
+    const stateRef = React.useRef(state);
 
     React.useEffect(() => {
       isInitialRef.current = false;
       onStateChangeRef.current = onStateChange;
+      stateRef.current = state;
     });
 
     React.useEffect(() => {
@@ -415,17 +437,19 @@ const BaseNavigationContainer = React.forwardRef(
     );
 
     let element = (
-      <ScheduleUpdateContext.Provider value={scheduleContext}>
-        <NavigationBuilderContext.Provider value={builderContext}>
-          <NavigationStateContext.Provider value={context}>
-            <UnhandledActionContext.Provider
-              value={onUnhandledAction ?? defaultOnUnhandledAction}
-            >
-              <EnsureSingleNavigator>{children}</EnsureSingleNavigator>
-            </UnhandledActionContext.Provider>
-          </NavigationStateContext.Provider>
-        </NavigationBuilderContext.Provider>
-      </ScheduleUpdateContext.Provider>
+      <NavigationContainerRefContext.Provider value={navigation}>
+        <ScheduleUpdateContext.Provider value={scheduleContext}>
+          <NavigationBuilderContext.Provider value={builderContext}>
+            <NavigationStateContext.Provider value={context}>
+              <UnhandledActionContext.Provider
+                value={onUnhandledAction ?? defaultOnUnhandledAction}
+              >
+                <EnsureSingleNavigator>{children}</EnsureSingleNavigator>
+              </UnhandledActionContext.Provider>
+            </NavigationStateContext.Provider>
+          </NavigationBuilderContext.Provider>
+        </ScheduleUpdateContext.Provider>
+      </NavigationContainerRefContext.Provider>
     );
 
     if (independent) {

packages/core/src/NavigationContainerRefContext.tsx

@@ -0,0 +1,12 @@
+import * as React from 'react';
+import type { ParamListBase } from '@react-navigation/routers';
+import type { NavigationContainerRef } from './types';
+
+/**
+ * Context which holds the route prop for a screen.
+ */
+const NavigationContainerRefContext = React.createContext<
+  NavigationContainerRef<ParamListBase> | undefined
+>(undefined);
+
+export default NavigationContainerRefContext;

packages/core/src/NavigationRouteContext.tsx

@@ -4,8 +4,8 @@ import type { Route } from '@react-navigation/routers';
 /**
  * Context which holds the route prop for a screen.
  */
-const NavigationContext = React.createContext<Route<string> | undefined>(
+const NavigationRouteContext = React.createContext<Route<string> | undefined>(
   undefined
 );
 
-export default NavigationContext;
+export default NavigationRouteContext;

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

@@ -106,13 +106,40 @@ it("gets navigation's parent's parent from context", () => {
   );
 });
 
+it('gets navigation from container from context', () => {
+  expect.assertions(1);
+
+  const TestNavigator = (props: any): any => {
+    const { state, descriptors } = useNavigationBuilder(MockRouter, props);
+
+    return state.routes.map((route) => descriptors[route.key].render());
+  };
+
+  const Test = () => {
+    const navigation = useNavigation();
+
+    expect(navigation.navigate).toBeDefined();
+
+    return null;
+  };
+
+  render(
+    <BaseNavigationContainer>
+      <Test />
+      <TestNavigator>
+        <Screen name="foo">{() => null}</Screen>
+      </TestNavigator>
+    </BaseNavigationContainer>
+  );
+});
+
 it('throws if called outside a navigation context', () => {
   expect.assertions(1);
 
   const Test = () => {
     // eslint-disable-next-line react-hooks/rules-of-hooks
     expect(() => useNavigation()).toThrow(
-      "Couldn't find a navigation object. Is your component inside a screen in a navigator?"
+      "Couldn't find a navigation object. Is your component inside NavigationContainer?"
     );
 
     return null;

packages/core/src/index.tsx

@@ -6,6 +6,7 @@ export { default as createNavigatorFactory } from './createNavigatorFactory';
 export { default as createNavigationContainerRef } from './createNavigationContainerRef';
 export { default as useNavigationContainerRef } from './useNavigationContainerRef';
 
+export { default as NavigationContainerRefContext } from './NavigationContainerRefContext';
 export { default as NavigationHelpersContext } from './NavigationHelpersContext';
 export { default as NavigationContext } from './NavigationContext';
 export { default as NavigationRouteContext } from './NavigationRouteContext';

packages/core/src/useNavigation.tsx

@@ -1,4 +1,5 @@
 import * as React from 'react';
+import NavigationContainerRefContext from './NavigationContainerRefContext';
 import NavigationContext from './NavigationContext';
 import type { NavigationProp } from './types';
 
@@ -10,14 +11,15 @@ import type { NavigationProp } from './types';
 export default function useNavigation<
   T = NavigationProp<ReactNavigation.RootParamList>
 >(): T {
+  const root = React.useContext(NavigationContainerRefContext);
   const navigation = React.useContext(NavigationContext);
 
-  if (navigation === undefined) {
+  if (navigation === undefined && root === undefined) {
     throw new Error(
-      "Couldn't find a navigation object. Is your component inside a screen in a navigator?"
+      "Couldn't find a navigation object. Is your component inside NavigationContainer?"
     );
   }
 
   // FIXME: Figure out a better way to do this
-  return (navigation as unknown) as T;
+  return ((navigation ?? root) as unknown) as T;
 }

packages/native/src/useLinkProps.tsx

@@ -2,6 +2,7 @@ import * as React from 'react';
 import { Platform, GestureResponderEvent } from 'react-native';
 import {
   NavigationAction,
+  NavigationContainerRefContext,
   NavigationHelpersContext,
 } from '@react-navigation/core';
 import useLinkTo, { To } from './useLinkTo';
@@ -20,6 +21,7 @@ type Props<ParamList extends ReactNavigation.RootParamList> = {
 export default function useLinkProps<
   ParamList extends ReactNavigation.RootParamList
 >({ to, action }: Props<ParamList>) {
+  const root = React.useContext(NavigationContainerRefContext);
   const navigation = React.useContext(NavigationHelpersContext);
   const linkTo = useLinkTo<ParamList>();
 
@@ -47,8 +49,12 @@ export default function useLinkProps<
       if (action) {
         if (navigation) {
           navigation.dispatch(action);
+        } else if (root) {
+          root.dispatch(action);
         } else {
-          throw new Error("Couldn't find a navigation object.");
+          throw new Error(
+            "Couldn't find a navigation object. Is your component inside NavigationContainer?"
+          );
         }
       } else {
         linkTo(to);

packages/native/src/useLinkTo.tsx

@@ -1,8 +1,8 @@
 import * as React from 'react';
 import {
+  NavigationContainerRefContext,
   getStateFromPath,
   getActionFromState,
-  NavigationContext,
 } from '@react-navigation/core';
 import LinkingContext from './LinkingContext';
 
@@ -24,25 +24,17 @@ export type To<
 export default function useLinkTo<
   ParamList extends ReactNavigation.RootParamList
 >() {
-  const navigation = React.useContext(NavigationContext);
+  const navigation = React.useContext(NavigationContainerRefContext);
   const linking = React.useContext(LinkingContext);
 
   const linkTo = React.useCallback(
     (to: To<ParamList>) => {
       if (navigation === undefined) {
         throw new Error(
-          "Couldn't find a navigation object. Is your component inside a screen in a navigator?"
+          "Couldn't find a navigation object. Is your component inside NavigationContainer?"
         );
       }
 
-      let root = navigation;
-      let current;
-
-      // Traverse up to get the root navigation
-      while ((current = root.getParent())) {
-        root = current;
-      }
-
       if (typeof to !== 'string') {
         // @ts-expect-error: This is fine
         root.navigate(to.screen, to.params);
@@ -63,9 +55,9 @@ export default function useLinkTo<
         const action = getActionFromState(state, options?.config);
 
         if (action !== undefined) {
-          root.dispatch(action);
+          navigation.dispatch(action);
         } else {
-          root.reset(state);
+          navigation.reset(state);
         }
       } else {
         throw new Error('Failed to parse the path to a navigation state.');