refactor: remove code for cleaning up navigator

It messes up with other state changes. We'll figure something out later

example/index.tsx

@@ -58,7 +58,7 @@ const First = ({
       Navigate with params
     </button>
     <button type="button" onClick={() => navigation.pop()}>
-      Go back
+      Pop
     </button>
   </div>
 );
@@ -80,7 +80,7 @@ const Second = ({
       Push first
     </button>
     <button type="button" onClick={() => navigation.pop()}>
-      Go back
+      Pop
     </button>
   </div>
 );
@@ -104,7 +104,7 @@ const Fourth = ({
     >
       Push first
     </button>
-    <button type="button" onClick={() => navigation.pop()}>
+    <button type="button" onClick={() => navigation.goBack()}>
       Go back
     </button>
   </div>
@@ -127,7 +127,7 @@ const Fifth = ({
       Push second
     </button>
     <button type="button" onClick={() => navigation.pop()}>
-      Go back
+      Pop
     </button>
   </div>
 );

src/NavigationContainer.tsx

@@ -33,16 +33,19 @@ export default function NavigationContainer({
     NavigationState | PartialState | undefined
   >(initialState);
 
-  const initialMountRef = React.useRef(true);
+  const firstRenderRef = React.useRef(true);
+  const initialStateRef = React.useRef(initialState);
   const stateRef = React.useRef(state);
 
-  stateRef.current = state;
+  React.useLayoutEffect(() => {
+    stateRef.current = state;
+  });
 
   React.useEffect(() => {
-    if (initialMountRef.current) {
-      initialMountRef.current = false;
+    if (firstRenderRef.current) {
+      firstRenderRef.current = false;
 
-      if (state === undefined) {
+      if (state === initialStateRef.current) {
         // Don't call the listener if we haven't initialized any state
         return;
       }

src/SceneView.tsx

@@ -57,7 +57,7 @@ export default function SceneView(props: Props) {
         ),
       });
     },
-    [getState, route.key, setState]
+    [getState, route, setState]
   );
 
   const context = React.useMemo(

src/__tests__/index.test.tsx

@@ -237,7 +237,7 @@ it("doesn't update state if nothing changed", () => {
 
   const onStateChange = jest.fn();
 
-  const element = (
+  render(
     <NavigationContainer onStateChange={onStateChange}>
       <TestNavigator initialRouteName="foo">
         <Screen name="foo" component={FooScreen} />
@@ -246,8 +246,6 @@ it("doesn't update state if nothing changed", () => {
     </NavigationContainer>
   );
 
-  render(element).update(element);
-
   expect(onStateChange).toBeCalledTimes(0);
 });
 
@@ -271,7 +269,7 @@ it("doesn't update state if action wasn't handled", () => {
 
   const onStateChange = jest.fn();
 
-  const element = (
+  render(
     <NavigationContainer onStateChange={onStateChange}>
       <TestNavigator initialRouteName="foo">
         <Screen name="foo" component={FooScreen} />
@@ -280,60 +278,9 @@ it("doesn't update state if action wasn't handled", () => {
     </NavigationContainer>
   );
 
-  render(element).update(element);
-
   expect(onStateChange).toBeCalledTimes(0);
 });
 
-it('cleans up state when the navigator unmounts', () => {
-  const TestNavigator = (props: any) => {
-    const { navigation, descriptors } = useNavigationBuilder(MockRouter, props);
-
-    return descriptors[
-      navigation.state.routes[navigation.state.index].key
-    ].render();
-  };
-
-  const FooScreen = (props: any) => {
-    React.useEffect(() => {
-      props.navigation.dispatch({ type: 'UPDATE' });
-      // eslint-disable-next-line react-hooks/exhaustive-deps
-    }, []);
-
-    return null;
-  };
-
-  const onStateChange = jest.fn();
-
-  const element = (
-    <NavigationContainer onStateChange={onStateChange}>
-      <TestNavigator>
-        <Screen name="foo" component={FooScreen} />
-        <Screen name="bar" component={jest.fn()} />
-      </TestNavigator>
-    </NavigationContainer>
-  );
-
-  const root = render(element);
-
-  root.update(element);
-
-  expect(onStateChange).toBeCalledTimes(1);
-  expect(onStateChange).lastCalledWith({
-    index: 0,
-    key: 'root',
-    routeNames: ['foo', 'bar'],
-    routes: [{ key: 'foo', name: 'foo' }, { key: 'bar', name: 'bar' }],
-  });
-
-  root.update(
-    <NavigationContainer onStateChange={onStateChange} children={null} />
-  );
-
-  expect(onStateChange).toBeCalledTimes(2);
-  expect(onStateChange).lastCalledWith(undefined);
-});
-
 it("lets parent handle the action if child didn't", () => {
   const ParentRouter: Router<{ type: string }> = {
     ...MockRouter,
@@ -372,6 +319,7 @@ it("lets parent handle the action if child didn't", () => {
   const TestScreen = (props: any) => {
     React.useEffect(() => {
       props.navigation.dispatch({ type: 'REVERSE' });
+
       // eslint-disable-next-line react-hooks/exhaustive-deps
     }, []);
 
@@ -380,7 +328,7 @@ it("lets parent handle the action if child didn't", () => {
 
   const onStateChange = jest.fn();
 
-  const element = (
+  render(
     <NavigationContainer onStateChange={onStateChange}>
       <ParentNavigator initialRouteName="baz">
         <Screen name="foo">{() => null}</Screen>
@@ -396,9 +344,7 @@ it("lets parent handle the action if child didn't", () => {
     </NavigationContainer>
   );
 
-  render(element).update(element);
-
-  expect(onStateChange).toBeCalledTimes(2);
+  expect(onStateChange).toBeCalledTimes(1);
   expect(onStateChange).lastCalledWith({
     index: 2,
     key: 'root',
@@ -476,7 +422,7 @@ it('updates route params with setParams', () => {
 
   const onStateChange = jest.fn();
 
-  const element = (
+  render(
     <NavigationContainer onStateChange={onStateChange}>
       <TestNavigator initialRouteName="foo">
         <Screen name="foo" component={FooScreen} />
@@ -485,8 +431,6 @@ it('updates route params with setParams', () => {
     </NavigationContainer>
   );
 
-  render(element);
-
   act(() => setParams({ username: 'alice' }));
 
   expect(onStateChange).toBeCalledTimes(1);
@@ -601,7 +545,7 @@ it("doesn't throw when direct children is Screen or empty element", () => {
     return null;
   };
 
-  const element = (
+  render(
     <NavigationContainer>
       <TestNavigator>
         <Screen name="foo" component={jest.fn()} />
@@ -612,6 +556,4 @@ it("doesn't throw when direct children is Screen or empty element", () => {
       </TestNavigator>
     </NavigationContainer>
   );
-
-  render(element).update(element);
 });

src/useNavigationBuilder.tsx

@@ -68,12 +68,6 @@ export default function useNavigationBuilder(
     setState,
   } = React.useContext(NavigationStateContext);
 
-  React.useEffect(() => {
-    // We need to clean up the state object when the navigator unmounts
-    return () => setState(undefined);
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, []);
-
   const getState = React.useCallback(
     (): NavigationState =>
       router.getRehydratedState({