Skip route.lazy hydration properties after hydration (#13376)

* Skip `route.lazy` hydration properties after hydration

* Refactor to use `lazyRoutePropertiesToSkip` array

* Refactor to make `initialHydration` an optional arg

* Update tests

Author: markdalgleish

.changeset/happy-spoons-watch.md

@@ -0,0 +1,21 @@
+---
+"react-router": patch
+---
+
+When using the object-based `route.lazy` API, the `HydrateFallback` and `hydrateFallbackElement` properties are now skipped when lazy loading routes after hydration.
+
+If you move the code for these properties into a separate file, you can use this optimization to avoid downloading unused hydration code. For example:
+
+```ts
+createBrowserRouter([
+  {
+    path: "/show/:showId",
+    lazy: {
+      loader: async () => (await import("./show.loader.js")).loader,
+      Component: async () => (await import("./show.component.js")).Component,
+      HydrateFallback: async () =>
+        (await import("./show.hydrate-fallback.js")).HydrateFallback,
+    },
+  },
+]);
+```

packages/react-router/__tests__/router/lazy-test.ts

@@ -1,5 +1,9 @@
 import { createMemoryHistory } from "../../lib/router/history";
 import { createRouter, createStaticHandler } from "../../lib/router/router";
+import {
+  createMemoryRouter,
+  hydrationRouteProperties,
+} from "../../lib/components";
 
 import type {
   TestNonIndexRouteObject,
@@ -561,6 +565,69 @@ describe("lazily loaded route modules", () => {
       expect(t.router.state.matches[0].route.action).toBeUndefined();
     });
 
+    it("only resolves lazy hydration route properties on hydration", async () => {
+      let [lazyLoaderForHydration, lazyLoaderDeferredForHydration] =
+        createAsyncStub();
+      let [lazyLoaderForNavigation, lazyLoaderDeferredForNavigation] =
+        createAsyncStub();
+      let [
+        lazyHydrateFallbackForHydration,
+        lazyHydrateFallbackDeferredForHydration,
+      ] = createAsyncStub();
+      let [
+        lazyHydrateFallbackElementForHydration,
+        lazyHydrateFallbackElementDeferredForHydration,
+      ] = createAsyncStub();
+      let lazyHydrateFallbackForNavigation = jest.fn(async () => null);
+      let lazyHydrateFallbackElementForNavigation = jest.fn(async () => null);
+      let router = createMemoryRouter(
+        [
+          {
+            path: "/hydration",
+            lazy: {
+              HydrateFallback: lazyHydrateFallbackForHydration,
+              hydrateFallbackElement: lazyHydrateFallbackElementForHydration,
+              loader: lazyLoaderForHydration,
+            },
+          },
+          {
+            path: "/navigation",
+            lazy: {
+              HydrateFallback: lazyHydrateFallbackForNavigation,
+              hydrateFallbackElement: lazyHydrateFallbackElementForNavigation,
+              loader: lazyLoaderForNavigation,
+            },
+          },
+        ],
+        {
+          initialEntries: ["/hydration"],
+        }
+      );
+      expect(router.state.initialized).toBe(false);
+
+      expect(lazyHydrateFallbackForHydration).toHaveBeenCalledTimes(1);
+      expect(lazyHydrateFallbackElementForHydration).toHaveBeenCalledTimes(1);
+      expect(lazyLoaderForHydration).toHaveBeenCalledTimes(1);
+      await lazyHydrateFallbackDeferredForHydration.resolve(null);
+      await lazyHydrateFallbackElementDeferredForHydration.resolve(null);
+      await lazyLoaderDeferredForHydration.resolve(null);
+
+      expect(router.state.location.pathname).toBe("/hydration");
+      expect(router.state.navigation.state).toBe("idle");
+      expect(router.state.initialized).toBe(true);
+
+      let navigationPromise = router.navigate("/navigation");
+      expect(router.state.location.pathname).toBe("/hydration");
+      expect(router.state.navigation.state).toBe("loading");
+      expect(lazyHydrateFallbackForNavigation).not.toHaveBeenCalled();
+      expect(lazyHydrateFallbackElementForNavigation).not.toHaveBeenCalled();
+      expect(lazyLoaderForNavigation).toHaveBeenCalledTimes(1);
+      await lazyLoaderDeferredForNavigation.resolve(null);
+      await navigationPromise;
+      expect(router.state.location.pathname).toBe("/navigation");
+      expect(router.state.navigation.state).toBe("idle");
+    });
+
     it("fetches lazy route functions on fetcher.load", async () => {
       let { routes, lazy, lazyDeferred } = createBasicLazyFunctionRoutes();
       let t = setup({ routes });
@@ -606,6 +673,40 @@ describe("lazily loaded route modules", () => {
       expect(lazyLoader).toHaveBeenCalledTimes(1);
     });
 
+    it("skips lazy hydration route properties on fetcher.load", async () => {
+      let [lazyLoader, lazyLoaderDeferred] = createAsyncStub();
+      let lazyHydrateFallback = jest.fn(async () => null);
+      let lazyHydrateFallbackElement = jest.fn(async () => null);
+      let routes = createBasicLazyRoutes({
+        loader: lazyLoader,
+        // @ts-expect-error
+        HydrateFallback: lazyHydrateFallback,
+        hydrateFallbackElement: lazyHydrateFallbackElement,
+      });
+      let t = setup({ routes, hydrationRouteProperties });
+      expect(lazyHydrateFallback).not.toHaveBeenCalled();
+      expect(lazyHydrateFallbackElement).not.toHaveBeenCalled();
+
+      let key = "key";
+      await t.fetch("/lazy", key);
+      expect(t.router.state.fetchers.get(key)?.state).toBe("loading");
+      expect(lazyLoader).toHaveBeenCalledTimes(1);
+      expect(lazyHydrateFallback).not.toHaveBeenCalled();
+      expect(lazyHydrateFallbackElement).not.toHaveBeenCalled();
+
+      let loaderDeferred = createDeferred();
+      lazyLoaderDeferred.resolve(() => loaderDeferred.promise);
+      expect(t.router.state.fetchers.get(key)?.state).toBe("loading");
+
+      await loaderDeferred.resolve("LAZY LOADER");
+      expect(t.fetchers[key].state).toBe("idle");
+      expect(t.fetchers[key].data).toBe("LAZY LOADER");
+
+      expect(lazyLoader).toHaveBeenCalledTimes(1);
+      expect(lazyHydrateFallback).not.toHaveBeenCalled();
+      expect(lazyHydrateFallbackElement).not.toHaveBeenCalled();
+    });
+
     it("fetches lazy route functions on fetcher.submit", async () => {
       let { routes, lazy, lazyDeferred } = createBasicLazyFunctionRoutes();
       let t = setup({ routes });
@@ -666,6 +767,49 @@ describe("lazily loaded route modules", () => {
       expect(lazyAction).toHaveBeenCalledTimes(1);
     });
 
+    it("skips lazy hydration route properties on fetcher.submit", async () => {
+      let [lazyLoaderStub, lazyLoaderDeferred] = createAsyncStub();
+      let [lazyActionStub, lazyActionDeferred] = createAsyncStub();
+      let lazyHydrateFallback = jest.fn(async () => null);
+      let lazyHydrateFallbackElement = jest.fn(async () => null);
+      let routes = createBasicLazyRoutes({
+        loader: lazyLoaderStub,
+        action: lazyActionStub,
+        // @ts-expect-error
+        HydrateFallback: lazyHydrateFallback,
+        hydrateFallbackElement: lazyHydrateFallbackElement,
+      });
+      let t = setup({ routes, hydrationRouteProperties });
+      expect(lazyLoaderStub).not.toHaveBeenCalled();
+      expect(lazyActionStub).not.toHaveBeenCalled();
+
+      let key = "key";
+      await t.fetch("/lazy", key, {
+        formMethod: "post",
+        formData: createFormData({}),
+      });
+      expect(t.router.state.fetchers.get(key)?.state).toBe("submitting");
+      expect(lazyLoaderStub).toHaveBeenCalledTimes(1);
+      expect(lazyActionStub).toHaveBeenCalledTimes(1);
+      expect(lazyHydrateFallback).not.toHaveBeenCalled();
+      expect(lazyHydrateFallbackElement).not.toHaveBeenCalled();
+
+      let actionDeferred = createDeferred();
+      let loaderDeferred = createDeferred();
+      lazyLoaderDeferred.resolve(() => loaderDeferred.promise);
+      lazyActionDeferred.resolve(() => actionDeferred.promise);
+      expect(t.router.state.fetchers.get(key)?.state).toBe("submitting");
+
+      await actionDeferred.resolve("LAZY ACTION");
+      expect(t.fetchers[key]?.state).toBe("idle");
+      expect(t.fetchers[key]?.data).toBe("LAZY ACTION");
+
+      expect(lazyLoaderStub).toHaveBeenCalledTimes(1);
+      expect(lazyActionStub).toHaveBeenCalledTimes(1);
+      expect(lazyHydrateFallback).not.toHaveBeenCalled();
+      expect(lazyHydrateFallbackElement).not.toHaveBeenCalled();
+    });
+
     it("fetches lazy route functions on staticHandler.query()", async () => {
       let { query } = createStaticHandler([
         {
@@ -751,6 +895,35 @@ describe("lazily loaded route modules", () => {
       let data = await response.json();
       expect(data).toEqual({ value: "LAZY LOADER" });
     });
+
+    it("resolves lazy hydration route properties on staticHandler.queryRoute()", async () => {
+      let lazyHydrateFallback = jest.fn(async () => null);
+      let lazyHydrateFallbackElement = jest.fn(async () => null);
+      let { queryRoute } = createStaticHandler(
+        [
+          {
+            id: "lazy",
+            path: "/lazy",
+            lazy: {
+              loader: async () => {
+                await tick();
+                return () => Response.json({ value: "LAZY LOADER" });
+              },
+              // @ts-expect-error
+              HydrateFallback: lazyHydrateFallback,
+              hydrateFallbackElement: lazyHydrateFallbackElement,
+            },
+          },
+        ],
+        { hydrationRouteProperties }
+      );
+
+      let response = await queryRoute(createRequest("/lazy"));
+      let data = await response.json();
+      expect(data).toEqual({ value: "LAZY LOADER" });
+      expect(lazyHydrateFallback).toHaveBeenCalled();
+      expect(lazyHydrateFallbackElement).toHaveBeenCalled();
+    });
   });
 
   describe("statically defined fields", () => {

packages/react-router/__tests__/router/utils/data-router-setup.ts

@@ -139,6 +139,7 @@ type SetupOpts = {
   basename?: string;
   initialEntries?: InitialEntry[];
   initialIndex?: number;
+  hydrationRouteProperties?: string[];
   hydrationData?: HydrationState;
   dataStrategy?: DataStrategyFunction;
 };
@@ -204,6 +205,7 @@ export function setup({
   basename,
   initialEntries,
   initialIndex,
+  hydrationRouteProperties,
   hydrationData,
   dataStrategy,
 }: SetupOpts) {
@@ -319,6 +321,7 @@ export function setup({
     basename,
     history,
     routes: enhanceRoutes(routes),
+    hydrationRouteProperties,
     hydrationData,
     window: testWindow,
     dataStrategy: dataStrategy,

packages/react-router/index.ts

@@ -317,7 +317,10 @@ export {
 } from "./lib/context";
 
 /** @internal */
-export { mapRouteProperties as UNSAFE_mapRouteProperties } from "./lib/components";
+export {
+  hydrationRouteProperties as UNSAFE_hydrationRouteProperties,
+  mapRouteProperties as UNSAFE_mapRouteProperties,
+} from "./lib/components";
 
 /** @internal */
 export { FrameworkContext as UNSAFE_FrameworkContext } from "./lib/dom/ssr/components";

packages/react-router/lib/components.tsx

@@ -131,6 +131,11 @@ export function mapRouteProperties(route: RouteObject) {
   return updates;
 }
 
+export const hydrationRouteProperties: (keyof RouteObject)[] = [
+  "HydrateFallback",
+  "hydrateFallbackElement",
+];
+
 export interface MemoryRouterOpts {
   /**
    * Basename path for the application.
@@ -194,6 +199,7 @@ export function createMemoryRouter(
     }),
     hydrationData: opts?.hydrationData,
     routes,
+    hydrationRouteProperties,
     mapRouteProperties,
     dataStrategy: opts?.dataStrategy,
     patchRoutesOnNavigation: opts?.patchRoutesOnNavigation,

packages/react-router/lib/dom-export/hydrated-router.tsx

@@ -21,6 +21,7 @@ import {
   UNSAFE_shouldHydrateRouteLoader as shouldHydrateRouteLoader,
   UNSAFE_useFogOFWarDiscovery as useFogOFWarDiscovery,
   UNSAFE_mapRouteProperties as mapRouteProperties,
+  UNSAFE_hydrationRouteProperties as hydrationRouteProperties,
   UNSAFE_createClientRoutesWithHMRRevalidationOptOut as createClientRoutesWithHMRRevalidationOptOut,
   matchRoutes,
 } from "react-router";
@@ -201,6 +202,7 @@ function createHydratedRouter({
     basename: ssrInfo.context.basename,
     unstable_getContext,
     hydrationData,
+    hydrationRouteProperties,
     mapRouteProperties,
     future: {
       unstable_middleware: ssrInfo.context.future.unstable_middleware,

packages/react-router/lib/dom/lib.tsx

@@ -66,7 +66,11 @@ import {
   mergeRefs,
   usePrefetchBehavior,
 } from "./ssr/components";
-import { Router, mapRouteProperties } from "../components";
+import {
+  Router,
+  mapRouteProperties,
+  hydrationRouteProperties,
+} from "../components";
 import type {
   RouteObject,
   NavigateOptions,
@@ -186,6 +190,7 @@ export function createBrowserRouter(
     hydrationData: opts?.hydrationData || parseHydrationData(),
     routes,
     mapRouteProperties,
+    hydrationRouteProperties,
     dataStrategy: opts?.dataStrategy,
     patchRoutesOnNavigation: opts?.patchRoutesOnNavigation,
     window: opts?.window,
@@ -209,6 +214,7 @@ export function createHashRouter(
     hydrationData: opts?.hydrationData || parseHydrationData(),
     routes,
     mapRouteProperties,
+    hydrationRouteProperties,
     dataStrategy: opts?.dataStrategy,
     patchRoutesOnNavigation: opts?.patchRoutesOnNavigation,
     window: opts?.window,