refactor: add transaction support (#12)

src/NavigationContainer.tsx

@@ -18,8 +18,9 @@ const MISSING_CONTEXT_ERROR =
 export const NavigationStateContext = React.createContext<{
   state?: NavigationState | PartialState;
   getState: () => NavigationState | PartialState | undefined;
-  setState: (state: NavigationState | undefined) => void;
+  setState: (state: NavigationState | undefined, dangerously?: boolean) => void;
   key?: string;
+  performTransaction: (action: () => void) => void;
 }>({
   get getState(): any {
     throw new Error(MISSING_CONTEXT_ERROR);
@@ -27,6 +28,9 @@ export const NavigationStateContext = React.createContext<{
   get setState(): any {
     throw new Error(MISSING_CONTEXT_ERROR);
   },
+  get performTransaction(): any {
+    throw new Error(MISSING_CONTEXT_ERROR);
+  },
 });
 
 export default class NavigationContainer extends React.Component<Props, State> {
@@ -42,10 +46,39 @@ export default class NavigationContainer extends React.Component<Props, State> {
     }
   }
 
-  private getNavigationState = () => this.state.navigationState;
+  private navigationState:
+    | NavigationState
+    | PartialState
+    | undefined
+    | null = null;
 
-  private setNavigationState = (navigationState: NavigationState | undefined) =>
-    this.setState({ navigationState });
+  private performTransaction = (action: () => void) => {
+    this.setState(
+      state => {
+        this.navigationState = state.navigationState;
+        action();
+        return { navigationState: this.navigationState };
+      },
+      () => (this.navigationState = null)
+    );
+  };
+
+  private getNavigationState = () =>
+    this.navigationState || this.state.navigationState;
+
+  private setNavigationState = (
+    navigationState: NavigationState | undefined,
+    dangerously = false
+  ) => {
+    if (this.navigationState === null && !dangerously) {
+      throw new Error('setState need to be wrapped in a performTransaction');
+    }
+    if (dangerously) {
+      this.setState({ navigationState });
+    } else {
+      this.navigationState = navigationState;
+    }
+  };
 
   render() {
     return (
@@ -54,6 +87,7 @@ export default class NavigationContainer extends React.Component<Props, State> {
           state: this.state.navigationState,
           getState: this.getNavigationState,
           setState: this.setNavigationState,
+          performTransaction: this.performTransaction,
         }}
       >
         <EnsureSingleNavigator>{this.props.children}</EnsureSingleNavigator>

src/SceneView.tsx

@@ -19,24 +19,27 @@ type Props = {
 
 export default function SceneView(props: Props) {
   const { screen, route, navigation: helpers, getState, setState } = props;
+  const { performTransaction } = React.useContext(NavigationStateContext);
 
   const navigation = React.useMemo(
     () => ({
       ...helpers,
       setParams: (params: object) => {
-        const state = getState();
+        performTransaction(() => {
+          const state = getState();
 
-        setState({
-          ...state,
-          routes: state.routes.map(r =>
-            r.key === route.key
-              ? { ...r, params: { ...r.params, ...params } }
-              : r
-          ),
+          setState({
+            ...state,
+            routes: state.routes.map(r =>
+              r.key === route.key
+                ? { ...r, params: { ...r.params, ...params } }
+                : r
+            ),
+          });
         });
       },
     }),
-    [getState, helpers, route.key, setState]
+    [getState, helpers, performTransaction, route.key, setState]
   );
 
   const getCurrentState = React.useCallback(() => {
@@ -65,9 +68,16 @@ export default function SceneView(props: Props) {
       state: route.state,
       getState: getCurrentState,
       setState: setCurrentState,
+      performTransaction,
       key: route.key,
     }),
-    [getCurrentState, route.key, route.state, setCurrentState]
+    [
+      getCurrentState,
+      performTransaction,
+      route.key,
+      route.state,
+      setCurrentState,
+    ]
   );
 
   return (

src/__tests__/NavigationContainer.test.tsx

@@ -1,6 +1,8 @@
 import * as React from 'react';
 import { render } from 'react-native-testing-library';
-import { NavigationStateContext } from '../NavigationContainer';
+import NavigationContainer, {
+  NavigationStateContext,
+} from '../NavigationContainer';
 
 it('throws when getState is accessed without a container', () => {
   expect.assertions(1);
@@ -39,3 +41,46 @@ it('throws when setState is accessed without a container', () => {
     "We couldn't find a navigation context. Have you wrapped your app with 'NavigationContainer'?"
   );
 });
+
+it('throws when performTransaction is accessed without a container', () => {
+  expect.assertions(1);
+
+  const Test = () => {
+    const { performTransaction } = React.useContext(NavigationStateContext);
+
+    // eslint-disable-next-line babel/no-unused-expressions
+    performTransaction;
+
+    return null;
+  };
+
+  const element = <Test />;
+
+  expect(() => render(element).update(element)).toThrowError(
+    "We couldn't find a navigation context. Have you wrapped your app with 'NavigationContainer'?"
+  );
+});
+
+it('throws when setState is called outside performTransaction', () => {
+  expect.assertions(1);
+
+  const Test = () => {
+    const { setState } = React.useContext(NavigationStateContext);
+    React.useEffect(() => {
+      setState(undefined);
+      // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, []);
+
+    return null;
+  };
+
+  const element = (
+    <NavigationContainer>
+      <Test />
+    </NavigationContainer>
+  );
+
+  expect(() => render(element).update(element)).toThrowError(
+    'setState need to be wrapped in a performTransaction'
+  );
+});

src/useNavigationBuilder.tsx

@@ -74,6 +74,7 @@ export default function useNavigationBuilder(
     getState: getCurrentState,
     setState,
     key,
+    performTransaction,
   } = React.useContext(NavigationStateContext);
 
   let state = router.getRehydratedState({
@@ -93,7 +94,7 @@ export default function useNavigationBuilder(
       // If the state needs to be updated, we'll schedule an update with React
       // setState in render seems hacky, but that's how React docs implement getDerivedPropsFromState
       // https://reactjs.org/docs/hooks-faq.html#how-do-i-implement-getderivedstatefromprops
-      setState(nextState);
+      setState(nextState, true);
     }
 
     state = nextState;
@@ -102,7 +103,9 @@ export default function useNavigationBuilder(
   React.useEffect(() => {
     return () => {
       // We need to clean up state for this navigator on unmount
-      getCurrentState() !== undefined && setState(undefined);
+      performTransaction(
+        () => getCurrentState() !== undefined && setState(undefined)
+      );
     };
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);

src/useNavigationHelpers.tsx

@@ -7,6 +7,7 @@ import {
   NavigationState,
   ActionCreators,
 } from './types';
+import { NavigationStateContext } from './NavigationContainer';
 
 type Options = {
   onAction: (action: NavigationAction, sourceNavigatorKey?: string) => boolean;
@@ -25,15 +26,19 @@ export default function useNavigationHelpers({
     NavigationBuilderContext
   );
 
+  const { performTransaction } = React.useContext(NavigationStateContext);
+
   return React.useMemo((): NavigationHelpers => {
     const dispatch = (
       action: NavigationAction | ((state: NavigationState) => NavigationState)
     ) => {
-      if (typeof action === 'function') {
-        setState(action(getState()));
-      } else {
-        onAction(action);
-      }
+      performTransaction(() => {
+        if (typeof action === 'function') {
+          setState(action(getState()));
+        } else {
+          onAction(action);
+        }
+      });
     };
 
     const actions = {
@@ -54,5 +59,12 @@ export default function useNavigationHelpers({
       ),
       dispatch,
     };
-  }, [getState, onAction, parentNavigationHelpers, actionCreators, setState]);
+  }, [
+    actionCreators,
+    parentNavigationHelpers,
+    performTransaction,
+    setState,
+    getState,
+    onAction,
+  ]);
 }