refactor: packages/

2 files changed, 176 insertions, 0 deletions

diff --git a/packages/excalidraw/components/hoc/withInternalFallback.test.tsx b/packages/excalidraw/components/hoc/withInternalFallback.test.tsx
new file mode 100644
index 0000000..5543133
--- /dev/null
+++ b/packages/excalidraw/components/hoc/withInternalFallback.test.tsx
@@ -0,0 +1,101 @@
+import React from "react";
+import { render, queryAllByTestId } from "../../tests/test-utils";
+import { Excalidraw, MainMenu } from "../../index";
+
+describe("Test internal component fallback rendering", () => {
+  it("should render only one menu per excalidraw instance (custom menu first scenario)", async () => {
+    const { container } = await render(
+      <div>
+        <Excalidraw>
+          <MainMenu>test</MainMenu>
+        </Excalidraw>
+        <Excalidraw />
+      </div>,
+    );
+
+    expect(queryAllByTestId(container, "main-menu-trigger")?.length).toBe(2);
+
+    const excalContainers = container.querySelectorAll<HTMLDivElement>(
+      ".excalidraw-container",
+    );
+
+    expect(
+      queryAllByTestId(excalContainers[0], "main-menu-trigger")?.length,
+    ).toBe(1);
+    expect(
+      queryAllByTestId(excalContainers[1], "main-menu-trigger")?.length,
+    ).toBe(1);
+  });
+
+  it("should render only one menu per excalidraw instance (default menu first scenario)", async () => {
+    const { container } = await render(
+      <div>
+        <Excalidraw />
+        <Excalidraw>
+          <MainMenu>test</MainMenu>
+        </Excalidraw>
+      </div>,
+    );
+
+    expect(queryAllByTestId(container, "main-menu-trigger")?.length).toBe(2);
+
+    const excalContainers = container.querySelectorAll<HTMLDivElement>(
+      ".excalidraw-container",
+    );
+
+    expect(
+      queryAllByTestId(excalContainers[0], "main-menu-trigger")?.length,
+    ).toBe(1);
+    expect(
+      queryAllByTestId(excalContainers[1], "main-menu-trigger")?.length,
+    ).toBe(1);
+  });
+
+  it("should render only one menu per excalidraw instance (two custom menus scenario)", async () => {
+    const { container } = await render(
+      <div>
+        <Excalidraw>
+          <MainMenu>test</MainMenu>
+        </Excalidraw>
+        <Excalidraw>
+          <MainMenu>test</MainMenu>
+        </Excalidraw>
+      </div>,
+    );
+
+    expect(queryAllByTestId(container, "main-menu-trigger")?.length).toBe(2);
+
+    const excalContainers = container.querySelectorAll<HTMLDivElement>(
+      ".excalidraw-container",
+    );
+
+    expect(
+      queryAllByTestId(excalContainers[0], "main-menu-trigger")?.length,
+    ).toBe(1);
+    expect(
+      queryAllByTestId(excalContainers[1], "main-menu-trigger")?.length,
+    ).toBe(1);
+  });
+
+  it("should render only one menu per excalidraw instance (two default menus scenario)", async () => {
+    const { container } = await render(
+      <div>
+        <Excalidraw />
+        <Excalidraw />
+      </div>,
+    );
+
+    expect(queryAllByTestId(container, "main-menu-trigger")?.length).toBe(2);
+
+    const excalContainers = container.querySelectorAll<HTMLDivElement>(
+      ".excalidraw-container",
+    );
+
+    expect(
+      queryAllByTestId(excalContainers[0], "main-menu-trigger")?.length,
+    ).toBe(1);
+    expect(
+      queryAllByTestId(excalContainers[1], "main-menu-trigger")?.length,
+    ).toBe(1);
+  });
+});
diff --git a/packages/excalidraw/components/hoc/withInternalFallback.tsx b/packages/excalidraw/components/hoc/withInternalFallback.tsx
new file mode 100644
index 0000000..5906b30
--- /dev/null
+++ b/packages/excalidraw/components/hoc/withInternalFallback.tsx
@@ -0,0 +1,75 @@
+import React, { useLayoutEffect, useRef } from "react";
+import { useTunnels } from "../../context/tunnels";
+import { atom } from "../../editor-jotai";
+
+export const withInternalFallback = <P,>(
+  componentName: string,
+  Component: React.FC<P>,
+) => {
+  const renderAtom = atom(0);
+
+  const WrapperComponent: React.FC<
+    P & {
+      __fallback?: boolean;
+    }
+  > = (props) => {
+    const {
+      tunnelsJotai: { useAtom },
+    } = useTunnels();
+    // for rerenders
+    const [, setCounter] = useAtom(renderAtom);
+    // for initial & subsequent renders. Tracked as component state
+    // due to excalidraw multi-instance scanerios.
+    const metaRef = useRef({
+      // flag set on initial render to tell the fallback component to skip the
+      // render until mount counter are initialized. This is because the counter
+      // is initialized in an effect, and thus we could end rendering both
+      // components at the same time until counter is initialized.
+      preferHost: false,
+      counter: 0,
+    });
+
+    useLayoutEffect(() => {
+      const meta = metaRef.current;
+      setCounter((c) => {
+        const next = c + 1;
+        meta.counter = next;
+
+        return next;
+      });
+      return () => {
+        setCounter((c) => {
+          const next = c - 1;
+          meta.counter = next;
+          if (!next) {
+            meta.preferHost = false;
+          }
+          return next;
+        });
+      };
+    }, [setCounter]);
+
+    if (!props.__fallback) {
+      metaRef.current.preferHost = true;
+    }
+
+    // ensure we don't render fallback and host components at the same time
+    if (
+      // either before the counters are initialized
+      (!metaRef.current.counter &&
+        props.__fallback &&
+        metaRef.current.preferHost) ||
+      // or after the counters are initialized, and both are rendered
+      // (this is the default when host renders as well)
+      (metaRef.current.counter > 1 && props.__fallback)
+    ) {
+      return null;
+    }
+
+    return <Component {...props} />;
+  };
+
+  WrapperComponent.displayName = componentName;
+
+  return WrapperComponent;
+};