From 5659a4632b2187c1b481203e66b2a79e6e76ff27 Mon Sep 17 00:00:00 2001
From: Matt Brophy <matt@brophy.org>
Date: Mon, 12 Dec 2022 18:05:43 -0500
Subject: [PATCH] ScrollRestoration on RR 6.4

---
 packages/remix-react/scroll-restoration.tsx | 146 +++++++-------------
 1 file changed, 48 insertions(+), 98 deletions(-)

diff --git a/packages/remix-react/scroll-restoration.tsx b/packages/remix-react/scroll-restoration.tsx
index 59ce0dd5c91..b14edb6a004 100644
--- a/packages/remix-react/scroll-restoration.tsx
+++ b/packages/remix-react/scroll-restoration.tsx
@@ -1,19 +1,15 @@
 import * as React from "react";
-import { useLocation } from "react-router-dom";
+import type { ScrollRestorationProps as ScrollRestorationPropsRR } from "react-router-dom";
+import {
+  useLocation,
+  UNSAFE_useScrollRestoration as useScrollRestoration,
+} from "react-router-dom";
 
-import { useBeforeUnload, useTransition } from "./components";
 import type { ScriptProps } from "./components";
+import { useMatches } from "./components";
 
 let STORAGE_KEY = "positions";
-
-let positions: { [key: string]: number } = {};
-
-if (typeof document !== "undefined") {
-  let sessionPositions = sessionStorage.getItem(STORAGE_KEY);
-  if (sessionPositions) {
-    positions = JSON.parse(sessionPositions);
-  }
-}
+let hydrated = false;
 
 /**
  * This component will emulate the browser's scroll restoration on location
@@ -21,29 +17,46 @@ if (typeof document !== "undefined") {
  *
  * @see https://remix.run/api/remix#scrollrestoration
  */
-export function ScrollRestoration(props: ScriptProps) {
-  useScrollRestoration();
-
-  // wait for the browser to restore it on its own
-  React.useEffect(() => {
-    window.history.scrollRestoration = "manual";
-  }, []);
-
-  // let the browser restore on it's own for refresh
-  useBeforeUnload(
-    React.useCallback(() => {
-      window.history.scrollRestoration = "auto";
-    }, [])
+export function ScrollRestoration({
+  getKey,
+  ...props
+}: ScriptProps & {
+  getKey: ScrollRestorationPropsRR["getKey"];
+}) {
+  let location = useLocation();
+  let matches = useMatches();
+
+  useScrollRestoration({
+    getKey,
+    storageKey: STORAGE_KEY,
+    skip: !hydrated,
+  });
+
+  // In order to support `getKey`, we need to compute a "key" here so we can
+  // hydrate that up so that SSR scroll restoration isn't waiting on React to
+  // hydrate. *However*, our key on the server is not the same as our key on
+  // the client!  So if the user's getKey implementation returns the SSR
+  // location key, then let's ignore it and let our inline <script> below pick
+  // up the client side history state key
+  let key = React.useMemo(
+    () => {
+      if (!getKey) return null;
+      let userKey = getKey(location, matches);
+      return userKey !== location.key ? userKey : null;
+    },
+    // Nah, we only need this the first time for the SSR render
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    []
   );
 
-  let restoreScroll = ((STORAGE_KEY: string) => {
+  let restoreScroll = ((STORAGE_KEY: string, restoreKey: string) => {
     if (!window.history.state || !window.history.state.key) {
       let key = Math.random().toString(32).slice(2);
       window.history.replaceState({ key }, "");
     }
     try {
       let positions = JSON.parse(sessionStorage.getItem(STORAGE_KEY) || "{}");
-      let storedY = positions[window.history.state.key];
+      let storedY = positions[restoreKey || window.history.state.key];
       if (typeof storedY === "number") {
         window.scrollTo(0, storedY);
       }
@@ -53,84 +66,21 @@ export function ScrollRestoration(props: ScriptProps) {
     }
   }).toString();
 
+  // Use this to skip restoration on the initial load since we'll do it in
+  // the inline <script> below
+  if (!hydrated) {
+    hydrated = true;
+  }
+
   return (
     <script
       {...props}
       suppressHydrationWarning
       dangerouslySetInnerHTML={{
-        __html: `(${restoreScroll})(${JSON.stringify(STORAGE_KEY)})`,
+        __html: `(${restoreScroll})(${JSON.stringify(
+          STORAGE_KEY
+        )}, ${JSON.stringify(key)})`,
       }}
     />
   );
 }
-
-let hydrated = false;
-
-function useScrollRestoration() {
-  let location = useLocation();
-  let transition = useTransition();
-
-  let wasSubmissionRef = React.useRef(false);
-
-  React.useEffect(() => {
-    if (transition.submission) {
-      wasSubmissionRef.current = true;
-    }
-  }, [transition]);
-
-  React.useEffect(() => {
-    if (transition.location) {
-      positions[location.key] = window.scrollY;
-    }
-  }, [transition, location]);
-
-  useBeforeUnload(
-    React.useCallback(() => {
-      sessionStorage.setItem(STORAGE_KEY, JSON.stringify(positions));
-    }, [])
-  );
-
-  if (typeof document !== "undefined") {
-    // eslint-disable-next-line
-    React.useLayoutEffect(() => {
-      // don't do anything on hydration, the component already did this with an
-      // inline script.
-      if (!hydrated) {
-        hydrated = true;
-        return;
-      }
-
-      let y = positions[location.key];
-
-      // been here before, scroll to it
-      if (y != undefined) {
-        window.scrollTo(0, y);
-        return;
-      }
-
-      // try to scroll to the hash
-      if (location.hash) {
-        let el = document.getElementById(location.hash.slice(1));
-        if (el) {
-          el.scrollIntoView();
-          return;
-        }
-      }
-
-      // don't do anything on submissions
-      if (wasSubmissionRef.current === true) {
-        wasSubmissionRef.current = false;
-        return;
-      }
-
-      // otherwise go to the top on new locations
-      window.scrollTo(0, 0);
-    }, [location]);
-  }
-
-  React.useEffect(() => {
-    if (transition.submission) {
-      wasSubmissionRef.current = true;
-    }
-  }, [transition]);
-}