refactor: rework how sync state is handled

This simplifies the sync store for navigation state: - Use a external store to store the state instead of `useState` - so the store can be updated synchronously - Use `useSyncExternalStore` API to subscribe to the state changes - Use an effect when a child navigator needs to update state due to list of screens changing etc. The changes should make the code easier to understand, as well as should make the code compatible with concurrent mode as we're not doing anything weird with the state.
react-navigation · Nov 5, 2023 · c346099 · c346099
1 parent ad45df6
commit c346099
Show file tree

Hide file tree

Showing 8 changed files with 155 additions and 161 deletions.
diff --git a/packages/core/src/BaseNavigationContainer.tsx b/packages/core/src/BaseNavigationContainer.tsx
@@ -33,7 +33,6 @@ import { useEventEmitter } from './useEventEmitter';
 import { useKeyedChildListeners } from './useKeyedChildListeners';
 import { useNavigationIndependentTree } from './useNavigationIndependentTree';
 import { useOptionsGetters } from './useOptionsGetters';
-import { ScheduleUpdateContext } from './useScheduleUpdate';
 import { useSyncState } from './useSyncState';
 
 type State = NavigationState | PartialState<NavigationState> | undefined;
@@ -103,10 +102,9 @@ export const BaseNavigationContainer = React.forwardRef(
       );
     }
 
-    const [state, getState, setState, scheduleUpdate, flushUpdates] =
-      useSyncState<State>(() =>
-        getPartialState(initialState == null ? undefined : initialState)
-      );
+    const [state, getState, setState] = useSyncState<State>(() =>
+      getPartialState(initialState == null ? undefined : initialState)
+    );
 
     const isFirstMountRef = React.useRef<boolean>(true);
 
@@ -205,7 +203,7 @@ export const BaseNavigationContainer = React.forwardRef(
         isFocused: () => true,
         canGoBack,
         getParent: () => undefined,
-        getState: () => stateRef.current,
+        getState,
         getRootState,
         getCurrentRoute,
         getCurrentOptions,
@@ -218,6 +216,7 @@ export const BaseNavigationContainer = React.forwardRef(
         getCurrentOptions,
         getCurrentRoute,
         getRootState,
+        getState,
         isReady,
         resetRoot,
       ]
@@ -262,11 +261,6 @@ export const BaseNavigationContainer = React.forwardRef(
       [addListener, addKeyedListener, onDispatchAction, onOptionsChange]
     );
 
-    const scheduleContext = React.useMemo(
-      () => ({ scheduleUpdate, flushUpdates }),
-      [scheduleUpdate, flushUpdates]
-    );
-
     const isInitialRef = React.useRef(true);
 
     const getIsInitial = React.useCallback(() => isInitialRef.current, []);
@@ -294,13 +288,11 @@ export const BaseNavigationContainer = React.forwardRef(
 
     const onReadyRef = React.useRef(onReady);
     const onStateChangeRef = React.useRef(onStateChange);
-    const stateRef = React.useRef(state);
 
     React.useEffect(() => {
       isInitialRef.current = false;
       onStateChangeRef.current = onStateChange;
       onReadyRef.current = onReady;
-      stateRef.current = state;
     });
 
     const onReadyCalledRef = React.useRef(false);
@@ -448,26 +440,25 @@ export const BaseNavigationContainer = React.forwardRef(
     if (state !== previousState && stateForNextRouteNamesChange !== null) {
       setStateForNextRouteNamesChange(null);
     }
+
     return (
       <NavigationIndependentTreeContext.Provider value={false}>
         <NavigationContainerRefContext.Provider value={navigation}>
-          <ScheduleUpdateContext.Provider value={scheduleContext}>
-            <NavigationBuilderContext.Provider value={builderContext}>
-              <NavigationStateContext.Provider value={context}>
-                <SetNextStateContext.Provider value={setNextStateContext}>
-                  <UnhandledActionContext.Provider
-                    value={onUnhandledAction ?? defaultOnUnhandledAction}
+          <NavigationBuilderContext.Provider value={builderContext}>
+            <NavigationStateContext.Provider value={context}>
+              <SetNextStateContext.Provider value={setNextStateContext}>
+                <UnhandledActionContext.Provider
+                  value={onUnhandledAction ?? defaultOnUnhandledAction}
+                >
+                  <DeprecatedNavigationInChildContext.Provider
+                    value={navigationInChildEnabled}
                   >
-                    <DeprecatedNavigationInChildContext.Provider
-                      value={navigationInChildEnabled}
-                    >
-                      <EnsureSingleNavigator>{children}</EnsureSingleNavigator>
-                    </DeprecatedNavigationInChildContext.Provider>
-                  </UnhandledActionContext.Provider>
-                </SetNextStateContext.Provider>
-              </NavigationStateContext.Provider>
-            </NavigationBuilderContext.Provider>
-          </ScheduleUpdateContext.Provider>
+                    <EnsureSingleNavigator>{children}</EnsureSingleNavigator>
+                  </DeprecatedNavigationInChildContext.Provider>
+                </UnhandledActionContext.Provider>
+              </SetNextStateContext.Provider>
+            </NavigationStateContext.Provider>
+          </NavigationBuilderContext.Provider>
         </NavigationContainerRefContext.Provider>
       </NavigationIndependentTreeContext.Provider>
     );

diff --git a/packages/core/src/__tests__/index.test.tsx b/packages/core/src/__tests__/index.test.tsx
@@ -430,8 +430,6 @@ it("doesn't update state if action wasn't handled", () => {
 });
 
 it('cleans up state when the navigator unmounts', () => {
-  jest.useFakeTimers();
-
   const TestNavigator = (props: any) => {
     const { state, descriptors } = useNavigationBuilder(MockRouter, props);
 
@@ -481,8 +479,6 @@ it('cleans up state when the navigator unmounts', () => {
     </BaseNavigationContainer>
   );
 
-  act(() => jest.runAllTimers());
-
   expect(onStateChange).toHaveBeenCalledTimes(2);
   expect(onStateChange).toHaveBeenLastCalledWith(undefined);
 });
@@ -536,6 +532,72 @@ it('allows state updates by dispatching a function returning an action', () => {
   });
 });
 
+it('re-initializes state once for conditional rendering', () => {
+  const TestNavigatorA = (props: any) => {
+    const { state, descriptors } = useNavigationBuilder(MockRouter, props);
+
+    return descriptors[state.routes[state.index].key].render();
+  };
+
+  const TestNavigatorB = (props: any) => {
+    const { state, descriptors } = useNavigationBuilder(MockRouter, props);
+
+    return descriptors[state.routes[state.index].key].render();
+  };
+
+  const onStateChange = jest.fn();
+
+  const navigation = createNavigationContainerRef<ParamListBase>();
+
+  const Test = ({ condition }: { condition: boolean }) => {
+    return (
+      <BaseNavigationContainer ref={navigation} onStateChange={onStateChange}>
+        {condition ? (
+          <TestNavigatorA>
+            <Screen name="foo">{() => null}</Screen>
+            <Screen name="bar">{() => null}</Screen>
+          </TestNavigatorA>
+        ) : (
+          <TestNavigatorB>
+            <Screen name="bar">{() => null}</Screen>
+            <Screen name="baz">{() => null}</Screen>
+          </TestNavigatorB>
+        )}
+      </BaseNavigationContainer>
+    );
+  };
+
+  const root = render(<Test condition />);
+
+  expect(onStateChange).toHaveBeenCalledTimes(0);
+  expect(navigation.getRootState()).toEqual({
+    stale: false,
+    type: 'test',
+    index: 0,
+    key: '0',
+    routeNames: ['foo', 'bar'],
+    routes: [
+      { key: 'foo', name: 'foo' },
+      { key: 'bar', name: 'bar' },
+    ],
+  });
+
+  root.update(<Test condition={false} />);
+
+  expect(onStateChange).toHaveBeenCalledTimes(1);
+  expect(onStateChange).toHaveBeenCalledWith({
+    stale: false,
+    type: 'test',
+    index: 0,
+    key: '1',
+    routeNames: ['bar', 'baz'],
+    routes: [
+      { key: 'bar', name: 'bar' },
+      { key: 'baz', name: 'baz' },
+    ],
+  });
+});
+
 it('updates route params with setParams', () => {
   const TestNavigator = (props: any) => {
     const { state, descriptors } = useNavigationBuilder(MockRouter, props);

diff --git a/packages/core/src/useIsomorphicLayoutEffect.native.tsx b/packages/core/src/useIsomorphicLayoutEffect.native.tsx
@@ -0,0 +1,3 @@
+import { useLayoutEffect } from 'react';
+
+export const useIsomorphicLayoutEffect = useLayoutEffect;
diff --git a/packages/core/src/useIsomorphicLayoutEffect.tsx b/packages/core/src/useIsomorphicLayoutEffect.tsx
@@ -0,0 +1,7 @@
+import { useEffect, useLayoutEffect } from 'react';
+
+/**
+ * Use `useEffect` during SSR and `useLayoutEffect` in the browser to avoid warnings.
+ */
+export const useIsomorphicLayoutEffect =
+  typeof document !== 'undefined' ? useLayoutEffect : useEffect;
diff --git a/packages/core/src/useNavigationBuilder.tsx b/packages/core/src/useNavigationBuilder.tsx
@@ -38,13 +38,13 @@ import { type ScreenConfigWithParent, useDescriptors } from './useDescriptors';
 import { useEventEmitter } from './useEventEmitter';
 import { useFocusedListenersChildrenAdapter } from './useFocusedListenersChildrenAdapter';
 import { useFocusEvents } from './useFocusEvents';
+import { useIsomorphicLayoutEffect } from './useIsomorphicLayoutEffect';
 import { useKeyedChildListeners } from './useKeyedChildListeners';
 import { useNavigationHelpers } from './useNavigationHelpers';
 import { useOnAction } from './useOnAction';
 import { useOnGetState } from './useOnGetState';
 import { useOnRouteFocus } from './useOnRouteFocus';
 import { useRegisterNavigator } from './useRegisterNavigator';
-import { useScheduleUpdate } from './useScheduleUpdate';
 
 // This is to make TypeScript compiler happy
 PrivateValueStore;
@@ -339,22 +339,17 @@ export function useNavigationBuilder<
 
   const stateCleanedUp = React.useRef(false);
 
-  const cleanUpState = React.useCallback(() => {
-    setCurrentState(undefined);
-    stateCleanedUp.current = true;
-  }, [setCurrentState]);
-
-  const setState = React.useCallback(
+  const setState = useLatestCallback(
     (state: NavigationState | PartialState<NavigationState> | undefined) => {
       if (stateCleanedUp.current) {
         // State might have been already cleaned up due to unmount
         // We do not want to expose API allowing to override this
         // This would lead to old data preservation on main navigator unmount
         return;
       }
+
       setCurrentState(state);
-    },
-    [setCurrentState]
+    }
   );
 
   const [initializedState, isFirstStateInitialization] = React.useMemo(() => {
@@ -547,7 +542,7 @@ export function useNavigationBuilder<
 
   const shouldUpdate = state !== nextState;
 
-  useScheduleUpdate(() => {
+  useIsomorphicLayoutEffect(() => {
     if (shouldUpdate) {
       // If the state needs to be updated, we'll schedule an update
       setState(nextState);
@@ -571,31 +566,21 @@ export function useNavigationBuilder<
 
     return () => {
       // We need to clean up state for this navigator on unmount
-      // We do it in a timeout because we need to detect if another navigator mounted in the meantime
-      // For example, if another navigator has started rendering, we should skip cleanup
-      // Otherwise, our cleanup step will cleanup state for the other navigator and re-initialize it
-      setTimeout(() => {
-        if (getCurrentState() !== undefined && getKey() === navigatorKey) {
-          cleanUpState();
-        }
-      }, 0);
+      if (getCurrentState() !== undefined && getKey() === navigatorKey) {
+        setCurrentState(undefined);
+        stateCleanedUp.current = true;
+      }
     };
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);
 
-  // We initialize this ref here to avoid a new getState getting initialized
-  // whenever initializedState changes. We want getState to have access to the
-  // latest initializedState, but don't need it to change when that happens
-  const initializedStateRef = React.useRef<State>();
-  initializedStateRef.current = initializedState;
-
-  const getState = React.useCallback((): State => {
-    const currentState = getCurrentState();
+  const getState = useLatestCallback((): State => {
+    const currentState = shouldUpdate ? nextState : getCurrentState();
 
-    return isStateInitialized(currentState)
-      ? (currentState as State)
-      : (initializedStateRef.current as State);
-  }, [getCurrentState, isStateInitialized]);
+    return (
+      isStateInitialized(currentState) ? currentState : initializedState
+    ) as State;
+  });
 
   const emitter = useEventEmitter<EventMapCore<State>>((e) => {
     const routeNames = [];

diff --git a/packages/core/src/useScheduleUpdate.tsx b/packages/core/src/useScheduleUpdate.tsx