From 2009082476983189b9cfe29a7b93b1de21b9d3b9 Mon Sep 17 00:00:00 2001
From: Patrik Oldsberg <poldsberg@gmail.com>
Date: Thu, 26 Nov 2020 16:07:03 +0100
Subject: [PATCH] core-api: add mount point parent discovery

Co-authored-by: blam <ben@blam.sh>
---
 .../core-api/src/routing/discovery.test.tsx   | 33 +++++++++---
 packages/core-api/src/routing/discovery.tsx   | 54 +++++++++++++++++++
 2 files changed, 80 insertions(+), 7 deletions(-)

diff --git a/packages/core-api/src/routing/discovery.test.tsx b/packages/core-api/src/routing/discovery.test.tsx
index c894c8ac6d..02c80297c0 100644
--- a/packages/core-api/src/routing/discovery.test.tsx
+++ b/packages/core-api/src/routing/discovery.test.tsx
@@ -15,7 +15,7 @@
  */
 
 import React, { PropsWithChildren } from 'react';
-import { collectRoutes } from './discovery';
+import { collectRoutes, collectRouteParents } from './discovery';
 import { createRouteRef } from './RouteRef';
 import { createPlugin } from '../plugin';
 import { createRoutableExtension } from '../lib/extensions';
@@ -58,7 +58,7 @@ describe('discovery', () => {
       </div>,
     ];
 
-    const routes = collectRoutes(
+    const root = (
       <MemoryRouter>
         <Routes>
           <Extension1 path="/foo">
@@ -83,10 +83,10 @@ describe('discovery', () => {
             <Route path="/divsoup" element={<Extension4 />} />
           </div>
         </Routes>
-      </MemoryRouter>,
+      </MemoryRouter>
     );
 
-    expect(routes).toEqual(
+    expect(collectRoutes(root)).toEqual(
       new Map([
         [ref1, '/foo'],
         [ref2, '/bar/:id'],
@@ -95,10 +95,20 @@ describe('discovery', () => {
         [ref5, '/blop'],
       ]),
     );
+
+    expect(collectRouteParents(root)).toEqual(
+      new Map([
+        [ref1, undefined],
+        [ref2, ref1],
+        [ref3, ref2],
+        [ref4, undefined],
+        [ref5, ref1],
+      ]),
+    );
   });
 
   it('should handle all react router Route patterns', () => {
-    const routes = collectRoutes(
+    const root = (
       <MemoryRouter>
         <Routes>
           <Route
@@ -116,10 +126,10 @@ describe('discovery', () => {
             <Extension5 path="/blop" />
           </Route>
         </Routes>
-      </MemoryRouter>,
+      </MemoryRouter>
     );
 
-    expect(routes).toEqual(
+    expect(collectRoutes(root)).toEqual(
       new Map([
         [ref1, '/foo'],
         [ref2, '/bar/:id'],
@@ -128,6 +138,15 @@ describe('discovery', () => {
         [ref5, '/blop'],
       ]),
     );
+    expect(collectRouteParents(root)).toEqual(
+      new Map([
+        [ref1, undefined],
+        [ref2, ref1],
+        [ref3, undefined],
+        [ref4, ref3],
+        [ref5, ref3],
+      ]),
+    );
   });
 
   it('should not visit the same element twice', () => {
diff --git a/packages/core-api/src/routing/discovery.tsx b/packages/core-api/src/routing/discovery.tsx
index 61e5df9136..96946e1ca3 100644
--- a/packages/core-api/src/routing/discovery.tsx
+++ b/packages/core-api/src/routing/discovery.tsx
@@ -70,6 +70,60 @@ export const collectRoutes = (tree: ReactNode) => {
   return treeMap;
 };
 
+export const collectRouteParents = (tree: ReactNode) => {
+  const treeMap = new Map<RouteRef, RouteRef | undefined>();
+
+  const nodes = [{ node: tree, parent: undefined as RouteRef | undefined }];
+
+  while (nodes.length !== 0) {
+    const { parent, node } = nodes.shift()!;
+    if (!isIterableElement(node)) {
+      continue;
+    }
+
+    React.Children.forEach(node, child => {
+      if (!isIterableElement(child)) {
+        return;
+      }
+
+      const { path, element, children } = child.props as {
+        path?: string;
+        element?: ReactNode;
+        children?: ReactNode;
+      };
+
+      let nextParent = parent;
+
+      if (path) {
+        const routeRef = getComponentData<RouteRef>(child, 'core.mountPoint');
+        if (routeRef) {
+          treeMap.set(routeRef, parent);
+          nextParent = routeRef;
+        } else if (isIterableElement(element)) {
+          const elementRouteRef = getComponentData<RouteRef>(
+            element,
+            'core.mountPoint',
+          );
+          if (elementRouteRef) {
+            treeMap.set(elementRouteRef, parent);
+
+            nextParent = elementRouteRef;
+            nodes.push({
+              parent: elementRouteRef,
+              node: element.props?.children,
+            });
+          } else {
+            nodes.push({ parent, node: element.props?.children });
+          }
+        }
+      }
+      nodes.push({ parent: nextParent, node: children });
+    });
+  }
+
+  return treeMap;
+};
+
 function isIterableElement(node: ReactNode): node is JSX.Element {
   return isValidElement(node) || Array.isArray(node);
 }