From 6955ae9d25be5fc61c2d4fcbc0a435af05adc02e Mon Sep 17 00:00:00 2001
From: "satyajit.happy" <satyajit.happy@gmail.com>
Date: Thu, 18 Jul 2019 16:49:58 +0200
Subject: [PATCH] feat: handle screens nested in React.Fragment

---
 src/__tests__/index.test.tsx | 45 ++++++++++++++++++++++++++++++++++
 src/useNavigationBuilder.tsx | 47 ++++++++++++++++++++----------------
 2 files changed, 71 insertions(+), 21 deletions(-)

diff --git a/src/__tests__/index.test.tsx b/src/__tests__/index.test.tsx
index 5ca18208..395ccda6 100644
--- a/src/__tests__/index.test.tsx
+++ b/src/__tests__/index.test.tsx
@@ -189,6 +189,51 @@ it('rehydrates state for a navigator on navigation', () => {
   });
 });
 
+it('initializes state for nested screens in React.Fragment', () => {
+  const TestNavigator = (props: any) => {
+    const { state, descriptors } = useNavigationBuilder(MockRouter, props);
+
+    return descriptors[state.routes[state.index].key].render();
+  };
+
+  const TestScreen = (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={TestScreen} />
+        <React.Fragment>
+          <Screen name="bar" component={jest.fn()} />
+          <Screen name="baz" component={jest.fn()} />
+        </React.Fragment>
+      </TestNavigator>
+    </NavigationContainer>
+  );
+
+  render(element).update(element);
+
+  expect(onStateChange).toBeCalledTimes(1);
+  expect(onStateChange).toBeCalledWith({
+    index: 0,
+    key: '0',
+    routeNames: ['foo', 'bar', 'baz'],
+    routes: [
+      { key: 'foo', name: 'foo' },
+      { key: 'bar', name: 'bar' },
+      { key: 'baz', name: 'baz' },
+    ],
+  });
+});
+
 it('initializes state for nested navigator on navigation', () => {
   const TestNavigator = (props: any) => {
     const { state, descriptors } = useNavigationBuilder(MockRouter, props);
diff --git a/src/useNavigationBuilder.tsx b/src/useNavigationBuilder.tsx
index 2cac2f10..1ae60faa 100644
--- a/src/useNavigationBuilder.tsx
+++ b/src/useNavigationBuilder.tsx
@@ -17,19 +17,18 @@ type Options = {
 const isArrayEqual = (a: any[], b: any[]) =>
   a.length === b.length && a.every((it, index) => it === b[index]);
 
-export default function useNavigationBuilder(
-  router: Router<any>,
-  options: Options
-) {
-  useRegisterNavigator();
+const getRouteConfigsFromChildren = (children: React.ReactNode) =>
+  React.Children.toArray(children).reduce<RouteConfig[]>((acc, child) => {
+    if (React.isValidElement(child)) {
+      if (child.type === Screen) {
+        acc.push(child.props as RouteConfig);
+        return acc;
+      }
 
-  const screens = React.Children.map(options.children, child => {
-    if (child === null || child === undefined) {
-      return;
-    }
-
-    if (React.isValidElement(child) && child.type === Screen) {
-      return child.props as RouteConfig;
+      if (child.type === React.Fragment) {
+        acc.push(...getRouteConfigsFromChildren(child.props.children));
+        return acc;
+      }
     }
 
     throw new Error(
@@ -38,15 +37,21 @@ export default function useNavigationBuilder(
         child.type && child.type.name ? child.type.name : String(child)
       }')`
     );
-  })
-    .filter(Boolean)
-    .reduce(
-      (acc, curr) => {
-        acc[curr!.name] = curr as RouteConfig;
-        return acc;
-      },
-      {} as { [key: string]: RouteConfig }
-    );
+  }, []);
+
+export default function useNavigationBuilder(
+  router: Router<any>,
+  options: Options
+) {
+  useRegisterNavigator();
+
+  const screens = getRouteConfigsFromChildren(options.children).reduce(
+    (acc, curr) => {
+      acc[curr.name] = curr;
+      return acc;
+    },
+    {} as { [key: string]: RouteConfig }
+  );
 
   const routeNames = Object.keys(screens);
   const initialRouteName =