From 4cff152a3891d4fe67edfc6457693ec323fe2fd4 Mon Sep 17 00:00:00 2001
From: Matt Brophy <matt@brophy.org>
Date: Fri, 13 Oct 2023 15:29:14 -0400
Subject: [PATCH] POC of TransitionProvider approach

---
 examples/view-transitions/src/main.tsx        | 13 ++-
 .../__tests__/data-browser-router-test.tsx    |  7 +-
 .../__tests__/exports-test.tsx                |  8 +-
 packages/react-router-dom/index.tsx           | 97 ++++---------------
 packages/react-router-dom/server.tsx          | 33 ++++---
 .../__tests__/exports-test.tsx                |  1 +
 packages/react-router/index.ts                |  2 +
 packages/react-router/lib/components.tsx      | 23 ++++-
 packages/react-router/lib/context.ts          |  9 ++
 9 files changed, 83 insertions(+), 110 deletions(-)

diff --git a/examples/view-transitions/src/main.tsx b/examples/view-transitions/src/main.tsx
index bfe23d1fc0..16f1f2103a 100644
--- a/examples/view-transitions/src/main.tsx
+++ b/examples/view-transitions/src/main.tsx
@@ -10,6 +10,7 @@ import {
   NavLink,
   Outlet,
   RouterProvider,
+  TransitionProvider,
   unstable_useViewTransitionState,
   useActionData,
   useLoaderData,
@@ -246,13 +247,21 @@ function NavImage({ src, idx }: { src: string; idx: number }) {
 const rootElement = document.getElementById("root") as HTMLElement;
 ReactDOMClient.createRoot(rootElement).render(
   <React.StrictMode>
-    <RouterProvider
+    <TransitionProvider
       router={router}
       future={{
         // Wrap all state updates in React.startTransition()
         v7_startTransition: true,
       }}
-    />
+    >
+      <RouterProvider
+        router={router}
+        future={{
+          // Wrap all state updates in React.startTransition()
+          v7_startTransition: true,
+        }}
+      />
+    </TransitionProvider>
   </React.StrictMode>
 );
 
diff --git a/packages/react-router-dom/__tests__/data-browser-router-test.tsx b/packages/react-router-dom/__tests__/data-browser-router-test.tsx
index aca4d565f9..ff24b27756 100644
--- a/packages/react-router-dom/__tests__/data-browser-router-test.tsx
+++ b/packages/react-router-dom/__tests__/data-browser-router-test.tsx
@@ -34,6 +34,7 @@ import {
   useMatches,
   useSearchParams,
   createRoutesFromElements,
+  TransitionProvider,
 } from "react-router-dom";
 
 import createDeferred from "../../router/__tests__/utils/createDeferred";
@@ -6006,7 +6007,11 @@ function testDomRouter(
           ],
           { window: testWindow }
         );
-        render(<RouterProvider router={router} />);
+        render(
+          <TransitionProvider router={router}>
+            <RouterProvider router={router} />
+          </TransitionProvider>
+        );
 
         expect(screen.getByText("Home")).toBeDefined();
         fireEvent.click(screen.getByText("/a"));
diff --git a/packages/react-router-dom/__tests__/exports-test.tsx b/packages/react-router-dom/__tests__/exports-test.tsx
index dcf2e22ba7..2d94b97884 100644
--- a/packages/react-router-dom/__tests__/exports-test.tsx
+++ b/packages/react-router-dom/__tests__/exports-test.tsx
@@ -4,21 +4,15 @@ import * as ReactRouterDOM from "react-router-dom";
 let nonReExportedKeys = new Set([
   "UNSAFE_mapRouteProperties",
   "UNSAFE_useRoutesImpl",
+  "UNSAFE_DataRouterSubscriberContext",
 ]);
 
-let modifiedExports = new Set(["RouterProvider"]);
-
 describe("react-router-dom", () => {
   for (let key in ReactRouter) {
     if (nonReExportedKeys.has(key)) {
       it(`does not re-export ${key} from react-router`, () => {
         expect(ReactRouterDOM[key]).toBe(undefined);
       });
-    } else if (modifiedExports.has(key)) {
-      it(`re-exports a different version of ${key}`, () => {
-        expect(ReactRouterDOM[key]).toBeDefined();
-        expect(ReactRouterDOM[key]).not.toBe(ReactRouter[key]);
-      });
     } else {
       it(`re-exports ${key} from react-router`, () => {
         expect(ReactRouterDOM[key]).toBe(ReactRouter[key]);
diff --git a/packages/react-router-dom/index.tsx b/packages/react-router-dom/index.tsx
index 8c62c56893..68d958d5e9 100644
--- a/packages/react-router-dom/index.tsx
+++ b/packages/react-router-dom/index.tsx
@@ -4,15 +4,12 @@
  */
 import * as React from "react";
 import type {
-  DataRouteObject,
   FutureConfig,
   Location,
   NavigateOptions,
   NavigationType,
-  Navigator,
   RelativeRoutingType,
   RouteObject,
-  RouterProviderProps,
   To,
 } from "react-router";
 import {
@@ -31,7 +28,7 @@ import {
   UNSAFE_RouteContext as RouteContext,
   UNSAFE_mapRouteProperties as mapRouteProperties,
   UNSAFE_useRouteId as useRouteId,
-  UNSAFE_useRoutesImpl as useRoutesImpl,
+  UNSAFE_DataRouterSubscriberContext as DataRouterSubscriberContext,
 } from "react-router";
 import type {
   BrowserHistory,
@@ -148,6 +145,7 @@ export {
   Outlet,
   Route,
   Router,
+  RouterProvider,
   Routes,
   createMemoryRouter,
   createPath,
@@ -419,14 +417,20 @@ class Deferred<T> {
   }
 }
 
+interface TransitionProviderProps {
+  router: RemixRouter;
+  future?: Partial<FutureConfig>;
+  children: React.ReactNode | React.ReactNode[];
+}
+
 /**
- * Given a Remix Router instance, render the appropriate UI
+ * Enable support for View Transitions in a RouterProvider
  */
-export function RouterProvider({
-  fallbackElement,
+export function TransitionProvider({
   router,
   future,
-}: RouterProviderProps): React.ReactElement {
+  children,
+}: TransitionProviderProps): React.ReactElement {
   let [state, setStateImpl] = React.useState(router.state);
   let [pendingState, setPendingState] = React.useState<RouterState>();
   let [vtContext, setVtContext] = React.useState<ViewTransitionContextObject>({
@@ -487,10 +491,6 @@ export function RouterProvider({
     [optInStartTransition, transition, renderDfd, router.window]
   );
 
-  // Need to use a layout effect here so we are subscribed early enough to
-  // pick up on any render-driven redirects/navigations (useEffect/<Navigate>)
-  React.useLayoutEffect(() => router.subscribe(setState), [router, setState]);
-
   // When we start a view transition, create a Deferred we can use for the
   // eventual "completed" render
   React.useEffect(() => {
@@ -546,78 +546,15 @@ export function RouterProvider({
     }
   }, [vtContext.isTransitioning, interruption]);
 
-  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,
-    }),
-    [router, navigator, basename]
-  );
-
-  // 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}>
-          <ViewTransitionContext.Provider value={vtContext}>
-            <Router
-              basename={basename}
-              location={state.location}
-              navigationType={state.historyAction}
-              navigator={navigator}
-            >
-              {state.initialized ? (
-                <DataRoutes routes={router.routes} state={state} />
-              ) : (
-                fallbackElement
-              )}
-            </Router>
-          </ViewTransitionContext.Provider>
-        </DataRouterStateContext.Provider>
-      </DataRouterContext.Provider>
-      {null}
-    </>
+    <DataRouterSubscriberContext.Provider value={[state, setState]}>
+      <ViewTransitionContext.Provider value={vtContext}>
+        {children}
+      </ViewTransitionContext.Provider>
+    </DataRouterSubscriberContext.Provider>
   );
 }
 
-function DataRoutes({
-  routes,
-  state,
-}: {
-  routes: DataRouteObject[];
-  state: RouterState;
-}): React.ReactElement | null {
-  return useRoutesImpl(routes, undefined, state);
-}
-
 export interface BrowserRouterProps {
   basename?: string;
   children?: React.ReactNode;
diff --git a/packages/react-router-dom/server.tsx b/packages/react-router-dom/server.tsx
index 7946218731..8c4390e772 100644
--- a/packages/react-router-dom/server.tsx
+++ b/packages/react-router-dom/server.tsx
@@ -19,6 +19,7 @@ import {
   UNSAFE_convertRoutesToDataRoutes as convertRoutesToDataRoutes,
 } from "@remix-run/router";
 import {
+  UNSAFE_DataRouterSubscriberContext as DataRouterSubscriberContext,
   UNSAFE_mapRouteProperties as mapRouteProperties,
   UNSAFE_useRoutesImpl as useRoutesImpl,
 } from "react-router";
@@ -130,21 +131,23 @@ export function StaticRouterProvider({
 
   return (
     <>
-      <DataRouterContext.Provider value={dataRouterContext}>
-        <DataRouterStateContext.Provider value={state}>
-          <ViewTransitionContext.Provider value={{ isTransitioning: false }}>
-            <Router
-              basename={dataRouterContext.basename}
-              location={state.location}
-              navigationType={state.historyAction}
-              navigator={dataRouterContext.navigator}
-              static={dataRouterContext.static}
-            >
-              <DataRoutes routes={router.routes} state={state} />
-            </Router>
-          </ViewTransitionContext.Provider>
-        </DataRouterStateContext.Provider>
-      </DataRouterContext.Provider>
+      <DataRouterSubscriberContext.Provider value={null}>
+        <ViewTransitionContext.Provider value={{ isTransitioning: false }}>
+          <DataRouterContext.Provider value={dataRouterContext}>
+            <DataRouterStateContext.Provider value={state}>
+              <Router
+                basename={dataRouterContext.basename}
+                location={state.location}
+                navigationType={state.historyAction}
+                navigator={dataRouterContext.navigator}
+                static={dataRouterContext.static}
+              >
+                <DataRoutes routes={router.routes} state={state} />
+              </Router>
+            </DataRouterStateContext.Provider>
+          </DataRouterContext.Provider>
+        </ViewTransitionContext.Provider>
+      </DataRouterSubscriberContext.Provider>
       {hydrateScript ? (
         <script
           suppressHydrationWarning
diff --git a/packages/react-router-native/__tests__/exports-test.tsx b/packages/react-router-native/__tests__/exports-test.tsx
index c13a20c8e1..ca5eade57f 100644
--- a/packages/react-router-native/__tests__/exports-test.tsx
+++ b/packages/react-router-native/__tests__/exports-test.tsx
@@ -4,6 +4,7 @@ import * as ReactRouterNative from "react-router-native";
 let nonReExportedKeys = new Set([
   "UNSAFE_mapRouteProperties",
   "UNSAFE_useRoutesImpl",
+  "UNSAFE_DataRouterSubscriberContext",
   "UNSAFE_ViewTransitionContext",
 ]);
 
diff --git a/packages/react-router/index.ts b/packages/react-router/index.ts
index 4ac70a3422..470171a40b 100644
--- a/packages/react-router/index.ts
+++ b/packages/react-router/index.ts
@@ -86,6 +86,7 @@ import type {
 import {
   DataRouterContext,
   DataRouterStateContext,
+  DataRouterSubscriberContext,
   LocationContext,
   NavigationContext,
   RouteContext,
@@ -305,6 +306,7 @@ export function createMemoryRouter(
 export {
   DataRouterContext as UNSAFE_DataRouterContext,
   DataRouterStateContext as UNSAFE_DataRouterStateContext,
+  DataRouterSubscriberContext as UNSAFE_DataRouterSubscriberContext,
   LocationContext as UNSAFE_LocationContext,
   NavigationContext as UNSAFE_NavigationContext,
   RouteContext as UNSAFE_RouteContext,
diff --git a/packages/react-router/lib/components.tsx b/packages/react-router/lib/components.tsx
index 4ec7dbc856..96f13aaed9 100644
--- a/packages/react-router/lib/components.tsx
+++ b/packages/react-router/lib/components.tsx
@@ -35,6 +35,7 @@ import {
   AwaitContext,
   DataRouterContext,
   DataRouterStateContext,
+  DataRouterSubscriberContext,
   LocationContext,
   NavigationContext,
   RouteContext,
@@ -92,23 +93,35 @@ export function RouterProvider({
   router,
   future,
 }: RouterProviderProps): React.ReactElement {
-  let [state, setStateImpl] = React.useState(router.state);
+  let ctx = React.useContext(DataRouterSubscriberContext);
+  let ctxState = ctx ? ctx[0] : null;
+  let ctxSetState = ctx ? ctx[1] : null;
+  let [localState, localSetStateImpl] = React.useState(router.state);
   let { v7_startTransition } = future || {};
 
   let setState = React.useCallback<RouterSubscriber>(
     (newState: RouterState) => {
       if (v7_startTransition && startTransitionImpl) {
-        startTransitionImpl(() => setStateImpl(newState));
+        startTransitionImpl(() => localSetStateImpl(newState));
       } else {
-        setStateImpl(newState);
+        localSetStateImpl(newState);
       }
     },
-    [setStateImpl, v7_startTransition]
+    [localSetStateImpl, v7_startTransition]
   );
 
+  // If we're inside a DataRouterSubscriberContext that needs fine-grained
+  // control of the state updates (for startViewTransition), then prefer those
+  // over our own state/setStateImpl
+  let state = ctxState || localState;
+  let subscriber = ctxSetState || setState;
+
   // Need to use a layout effect here so we are subscribed early enough to
   // pick up on any render-driven redirects/navigations (useEffect/<Navigate>)
-  React.useLayoutEffect(() => router.subscribe(setState), [router, setState]);
+  React.useLayoutEffect(
+    () => router.subscribe(subscriber),
+    [router, subscriber]
+  );
 
   let navigator = React.useMemo((): Navigator => {
     return {
diff --git a/packages/react-router/lib/context.ts b/packages/react-router/lib/context.ts
index 212dc7160a..2e0da12399 100644
--- a/packages/react-router/lib/context.ts
+++ b/packages/react-router/lib/context.ts
@@ -9,6 +9,8 @@ import type {
   Action as NavigationType,
   RelativeRoutingType,
   Router,
+  RouterState,
+  RouterSubscriber,
   StaticHandlerContext,
   To,
   TrackedPromise,
@@ -84,6 +86,13 @@ if (__DEV__) {
   DataRouterStateContext.displayName = "DataRouterState";
 }
 
+export const DataRouterSubscriberContext = React.createContext<
+  [RouterState, RouterSubscriber] | null
+>(null);
+if (__DEV__) {
+  DataRouterSubscriberContext.displayName = "DataRouterSubscriber";
+}
+
 export const AwaitContext = React.createContext<TrackedPromise | null>(null);
 if (__DEV__) {
   AwaitContext.displayName = "Await";