refactor: memoize the context object in container (facebook#14)

src/NavigationContainer.tsx

@@ -8,9 +8,7 @@ type Props = {
  children: React.ReactNode;
 };
 
-type State = {
- navigationState: NavigationState | PartialState | undefined;
-};
+type State = NavigationState | PartialState | undefined;
 
 const MISSING_CONTEXT_ERROR =
  "We couldn't find a navigation context. Have you wrapped your app with 'NavigationContainer'?";
@@ -34,65 +32,66 @@ export const NavigationStateContext = React.createContext<{
 });
 
 export default function NavigationContainer(props: Props) {
- const [state, setState] = React.useState<State>({
- navigationState: props.initialState,
- });
-
- const navigationState = React.useRef<
- NavigationState | PartialState | undefined | null
- >(null);
-
- const {
- performTransaction,
- getNavigationState,
- setNavigationState,
- }: {
- performTransaction: (action: () => void) => void;
- getNavigationState: () => PartialState | NavigationState | undefined;
- setNavigationState: (
- newNavigationState: NavigationState | undefined
- ) => void;
- } = React.useMemo(
+ const { onStateChange } = props;
+ const [state, setState] = React.useState<State>(props.initialState);
+
+ const navigationStateRef = React.useRef<State | null>(null);
+ const isTransactionActiveRef = React.useRef<boolean>(false);
+ const isFirstMountRef = React.useRef<boolean>(true);
+
+ const context = React.useMemo(
  () => ({
- performTransaction: action => {
- setState((state: State) => {
- navigationState.current = state.navigationState;
- action();
- return { navigationState: navigationState.current };
+ state,
+
+ performTransaction: (callback: () => void) => {
+ if (isTransactionActiveRef.current) {
+ throw new Error(
+ "Only one transaction can be active at a time. Did you accidentally nest 'performTransaction'?"
+ );
+ }
+
+ setState((navigationState: State) => {
+ isTransactionActiveRef.current = true;
+ navigationStateRef.current = navigationState;
+
+ callback();
+
+ isTransactionActiveRef.current = false;
+
+ return navigationStateRef.current;
  });
  },
- getNavigationState: () =>
- navigationState.current || state.navigationState,
- setNavigationState: newNavigationState => {
- if (navigationState.current === null) {
+
+ getState: () =>
+ navigationStateRef.current !== null
+ ? navigationStateRef.current
+ : state,
+
+ setState: (navigationState: State) => {
+ if (navigationStateRef.current === null) {
  throw new Error(
- 'setState need to be wrapped in a performTransaction'
+ "Any 'setState' calls need to be done inside 'performTransaction'"
  );
  }
- navigationState.current = newNavigationState;
+
+ navigationStateRef.current = navigationState;
  },
  }),
- []
+ [state]
  );
 
- const isFirstMount = React.useRef<boolean>(true);
  React.useEffect(() => {
- navigationState.current = null;
- if (!isFirstMount.current && props.onStateChange) {
- props.onStateChange(state.navigationState);
+ navigationStateRef.current = null;
+
+ if (!isFirstMountRef.current && onStateChange) {
+ onStateChange(state);
  }
- isFirstMount.current = false;
- }, [state.navigationState, props.onStateChange]);
+
+ isFirstMountRef.current = false;
+ }, [state, onStateChange]);
 
  return (
- <NavigationStateContext.Provider
- value={{
- state: state.navigationState,
- getState: getNavigationState,
- setState: setNavigationState,
- performTransaction: performTransaction,
- }}
- >
+ <NavigationStateContext.Provider value={context}>
  <EnsureSingleNavigator>{props.children}</EnsureSingleNavigator>
  </NavigationStateContext.Provider>
  );

src/__tests__/NavigationContainer.test.tsx

@@ -66,6 +66,7 @@ it('throws when setState is called outside performTransaction', () => {
 
  const Test = () => {
  const { setState } = React.useContext(NavigationStateContext);
+
  React.useEffect(() => {
  setState(undefined);
  // eslint-disable-next-line react-hooks/exhaustive-deps
@@ -81,6 +82,33 @@ it('throws when setState is called outside performTransaction', () => {
  );
 
  expect(() => render(element).update(element)).toThrowError(
- 'setState need to be wrapped in a performTransaction'
+ "Any 'setState' calls need to be done inside 'performTransaction'"
+ );
+});
+
+it('throws when nesting performTransaction', () => {
+ expect.assertions(1);
+
+ const Test = () => {
+ const { performTransaction } = React.useContext(NavigationStateContext);
+
+ React.useEffect(() => {
+ performTransaction(() => {
+ performTransaction(() => {});
+ });
+ // eslint-disable-next-line react-hooks/exhaustive-deps
+ }, []);
+
+ return null;
+ };
+
+ const element = (
+ <NavigationContainer>
+ <Test />
+ </NavigationContainer>
+ );
+
+ expect(() => render(element).update(element)).toThrowError(
+ "Only one transaction can be active at a time. Did you accidentally nest 'performTransaction'?"
  );
 });