feat: attempt at enabling full transition support for the rsc router

Author: jacob-ebey

packages/react-router/lib/components.tsx

@@ -291,6 +291,137 @@ export interface RouterProviderProps {
   unstable_onError?: unstable_ClientOnErrorFunction;
 }
 
+function shallowDiff(a: any, b: any) {
+  if (a === b) {
+    return false;
+  }
+  let aKeys = Object.keys(a);
+  let bKeys = Object.keys(b);
+  if (aKeys.length !== bKeys.length) {
+    return true;
+  }
+  for (let key of aKeys) {
+    if (a[key] !== b[key]) {
+      return true;
+    }
+  }
+  return false;
+}
+
+export function UNSTABLE_TransitionEnabledRouterProvider({
+  router,
+  flushSync: reactDomFlushSyncImpl,
+  unstable_onError,
+}: RouterProviderProps) {
+  let fetcherData = React.useRef<Map<string, any>>(new Map());
+  let [revalidating, startRevalidation] = React.useTransition();
+  (router as any).__startRevalidation = startRevalidation;
+  let [_state, setState] = React.useState(router.state);
+  let [state, setOptimisticState] = (
+    (React as any).useOptimistic as typeof React.useState
+  )(_state);
+
+  let navigator = React.useMemo((): Navigator => {
+    return {
+      createHref: router.createHref,
+      encodeLocation: router.encodeLocation,
+      go: (n) => router.navigate(n),
+      push: (to, state, opts) =>
+        router.navigate(to, {
+          state,
+          preventScrollReset: opts?.preventScrollReset,
+        }),
+      replace: (to, state, opts) =>
+        router.navigate(to, {
+          replace: true,
+          state,
+          preventScrollReset: opts?.preventScrollReset,
+        }),
+    };
+  }, [router]);
+
+  let basename = router.basename || "/";
+
+  let dataRouterContext = React.useMemo(
+    () => ({
+      router,
+      navigator,
+      static: false,
+      basename,
+      unstable_onError,
+    }),
+    [router, navigator, basename, unstable_onError],
+  );
+
+  React.useLayoutEffect(() => {
+    return router.subscribe(
+      (newState, { deletedFetchers, flushSync, viewTransitionOpts }) => {
+        newState.fetchers.forEach((fetcher, key) => {
+          if (fetcher.data !== undefined) {
+            fetcherData.current.set(key, fetcher.data);
+          }
+        });
+        deletedFetchers.forEach((key) => fetcherData.current.delete(key));
+
+        const diff = shallowDiff(state, newState);
+
+        if (!diff) return;
+
+        if (flushSync) {
+          if (reactDomFlushSyncImpl) {
+            reactDomFlushSyncImpl(() => setState(newState));
+          } else {
+            setState(newState);
+          }
+        } else {
+          React.startTransition(() => {
+            setOptimisticState(newState);
+            setState(newState);
+          });
+        }
+      },
+    );
+  }, [router, reactDomFlushSyncImpl, state, setOptimisticState]);
+
+  // The fragment and {null} here are important!  We need them to keep React 18's
+  // useId happy when we are server-rendering since we may have a <script> here
+  // containing the hydrated server-side staticContext (from StaticRouterProvider).
+  // useId relies on the component tree structure to generate deterministic id's
+  // so we need to ensure it remains the same on the client even though
+  // we don't need the <script> tag
+  return (
+    <>
+      <DataRouterContext.Provider value={dataRouterContext}>
+        <DataRouterStateContext.Provider
+          value={{
+            ...state,
+            revalidation: revalidating ? "loading" : state.revalidation,
+          }}
+        >
+          <FetchersContext.Provider value={fetcherData.current}>
+            {/* <ViewTransitionContext.Provider value={vtContext}> */}
+            <Router
+              basename={basename}
+              location={state.location}
+              navigationType={state.historyAction}
+              navigator={navigator}
+            >
+              <MemoizedDataRoutes
+                routes={router.routes}
+                future={router.future}
+                state={state}
+                unstable_onError={unstable_onError}
+              />
+            </Router>
+            {/* </ViewTransitionContext.Provider> */}
+          </FetchersContext.Provider>
+        </DataRouterStateContext.Provider>
+      </DataRouterContext.Provider>
+      {null}
+    </>
+  );
+}
+
 /**
  * Render the UI for the given {@link DataRouter}. This component should
  * typically be at the top of an app's element tree.

packages/react-router/lib/rsc/browser.tsx

@@ -1,7 +1,7 @@
 import * as React from "react";
 import * as ReactDOM from "react-dom";
 
-import { RouterProvider } from "../components";
+import { UNSTABLE_TransitionEnabledRouterProvider as RouterProvider } from "../components";
 import {
   RSCRouterContext,
   type DataRouteMatch,
@@ -153,8 +153,7 @@ export function createCallServer({
     }
 
     if (payload.rerender) {
-      React.startTransition(
-        // @ts-expect-error - We have old react types that don't know this can be async
+      (globalVar.__reactRouterDataRouter as any).__startRevalidation(
         async () => {
           const rerender = await payload.rerender;
           if (!rerender) return;
@@ -176,41 +175,38 @@ export function createCallServer({
               return;
             }
 
-            let lastMatch: RSCRouteManifest | undefined;
-            for (const match of rerender.matches) {
-              globalVar.__reactRouterDataRouter.patchRoutes(
-                lastMatch?.id ?? null,
-                [createRouteFromServerManifest(match)],
-                true,
-              );
-              lastMatch = match;
-            }
-            (
-              window as WindowWithRouterGlobals
-            ).__reactRouterDataRouter._internalSetStateDoNotUseOrYouWillBreakYourApp(
-              {},
-            );
+            (globalVar.__reactRouterDataRouter as any).__startRevalidation(
+              () => {
+                let lastMatch: RSCRouteManifest | undefined;
+                for (const match of rerender.matches) {
+                  globalVar.__reactRouterDataRouter.patchRoutes(
+                    lastMatch?.id ?? null,
+                    [createRouteFromServerManifest(match)],
+                    true,
+                  );
+                  lastMatch = match;
+                }
 
-            React.startTransition(() => {
-              (
-                window as WindowWithRouterGlobals
-              ).__reactRouterDataRouter._internalSetStateDoNotUseOrYouWillBreakYourApp(
-                {
-                  loaderData: Object.assign(
-                    {},
-                    globalVar.__reactRouterDataRouter.state.loaderData,
-                    rerender.loaderData,
-                  ),
-                  errors: rerender.errors
-                    ? Object.assign(
-                        {},
-                        globalVar.__reactRouterDataRouter.state.errors,
-                        rerender.errors,
-                      )
-                    : null,
-                },
-              );
-            });
+                (
+                  window as WindowWithRouterGlobals
+                ).__reactRouterDataRouter._internalSetStateDoNotUseOrYouWillBreakYourApp(
+                  {
+                    loaderData: Object.assign(
+                      {},
+                      globalVar.__reactRouterDataRouter.state.loaderData,
+                      rerender.loaderData,
+                    ),
+                    errors: rerender.errors
+                      ? Object.assign(
+                          {},
+                          globalVar.__reactRouterDataRouter.state.errors,
+                          rerender.errors,
+                        )
+                      : null,
+                  },
+                );
+              },
+            );
           }
         },
       );