fix: prevent remount of underlying component (#22)

Co-authored-by: Simen Bekkhus <[email protected]>
Co-authored-by: Eric Parsons <[email protected]>

Author: 3 people

src/__tests__/index.test.tsx

@@ -1,26 +1,45 @@
-import React from "react";
-import { act, render, screen, waitFor } from "@testing-library/react";
+import React, { memo, useEffect } from "react";
+import {
+    act,
+    render,
+    RenderResult,
+    screen,
+    waitFor,
+} from "@testing-library/react";
 import lazy, { lazyWithPreload as namedExport } from "../index";
 
 function getTestComponentModule() {
     const TestComponent = React.forwardRef<
         HTMLDivElement,
         { foo: string; children: React.ReactNode }
     >(function TestComponent(props, ref) {
+        renders++;
+        useEffect(() => {
+            mounts++;
+        }, []);
         return <div ref={ref}>{`${props.foo} ${props.children}`}</div>;
     });
     let loaded = false;
     let loadCalls = 0;
+    let renders = 0;
+    let mounts = 0;
 
     return {
         isLoaded: () => loaded,
         loadCalls: () => loadCalls,
+        renders: () => renders,
+        mounts: () => mounts,
         OriginalComponent: TestComponent,
         TestComponent: async () => {
             loaded = true;
             loadCalls++;
             return { default: TestComponent };
         },
+        TestMemoizedComponent: async () => {
+            loaded = true;
+            loadCalls++;
+            return { default: memo(TestComponent) };
+        },
     };
 }
 
@@ -158,4 +177,66 @@ describe("lazy", () => {
     it("exports named export as well", () => {
         expect(lazy).toBe(namedExport);
     });
+
+    it("does not re-render memoized base component when passed same props after preload", async () => {
+        const { TestMemoizedComponent, renders } = getTestComponentModule();
+        const LazyTestComponent = lazy(TestMemoizedComponent);
+
+        expect(renders()).toBe(0);
+
+        let rerender: RenderResult["rerender"] | undefined;
+        await act(async () => {
+            const result = render(
+                <React.Suspense fallback={null}>
+                    <LazyTestComponent foo="bar">baz</LazyTestComponent>
+                </React.Suspense>
+            );
+            rerender = result.rerender;
+        });
+
+        expect(renders()).toBe(1);
+
+        await LazyTestComponent.preload();
+
+        await act(async () => {
+            rerender?.(
+                <React.Suspense fallback={null}>
+                    <LazyTestComponent foo="bar">baz</LazyTestComponent>
+                </React.Suspense>
+            );
+        });
+
+        expect(renders()).toBe(1);
+    });
+
+    it("does not re-mount base component after preload", async () => {
+        const { TestComponent, mounts } = getTestComponentModule();
+        const LazyTestComponent = lazy(TestComponent);
+
+        expect(mounts()).toBe(0);
+
+        let rerender: RenderResult["rerender"] | undefined;
+        await act(async () => {
+            const result = render(
+                <React.Suspense fallback={null}>
+                    <LazyTestComponent foo="bar">baz</LazyTestComponent>
+                </React.Suspense>
+            );
+            rerender = result.rerender;
+        });
+
+        expect(mounts()).toBe(1);
+
+        await LazyTestComponent.preload();
+
+        await act(async () => {
+            rerender?.(
+                <React.Suspense fallback={null}>
+                    <LazyTestComponent foo="updated">updated</LazyTestComponent>
+                </React.Suspense>
+            );
+        });
+
+        expect(mounts()).toBe(1);
+    });
 });

src/index.ts

@@ -1,4 +1,4 @@
-import { ComponentType, createElement, forwardRef, lazy } from "react";
+import { ComponentType, createElement, forwardRef, lazy, useRef } from "react";
 
 export type PreloadableComponent<T extends ComponentType<any>> = T & {
     preload: () => Promise<T>;
@@ -7,29 +7,37 @@ export type PreloadableComponent<T extends ComponentType<any>> = T & {
 export function lazyWithPreload<T extends ComponentType<any>>(
     factory: () => Promise<{ default: T }>
 ): PreloadableComponent<T> {
-    const LazyComponent = lazy(factory);
+    const ReactLazyComponent = lazy(factory);
+    let PreloadedComponent: T | undefined;
     let factoryPromise: Promise<T> | undefined;
-    let LoadedComponent: T | undefined;
 
     const Component = forwardRef(function LazyWithPreload(props, ref) {
+        // Once one of these is chosen, we must ensure that it continues to be
+        // used for all subsequent renders, otherwise it can cause the
+        // underlying component to be unmounted and remounted.
+        const ComponentToRender = useRef(
+            PreloadedComponent ?? ReactLazyComponent
+        );
         return createElement(
-            LoadedComponent ?? LazyComponent,
+            ComponentToRender.current,
             Object.assign(ref ? { ref } : {}, props) as any
         );
-    }) as any as PreloadableComponent<T>;
+    });
+
+    const LazyWithPreload = Component as any as PreloadableComponent<T>;
 
-    Component.preload = () => {
+    LazyWithPreload.preload = () => {
         if (!factoryPromise) {
             factoryPromise = factory().then((module) => {
-                LoadedComponent = module.default;
-                return LoadedComponent;
+                PreloadedComponent = module.default;
+                return PreloadedComponent;
             });
         }
 
         return factoryPromise;
     };
 
-    return Component;
+    return LazyWithPreload;
 }
 
 export default lazyWithPreload;