Finish migration and keep previous compatibility

packages/gatsby-react-router-scroll/package.json

@@ -33,9 +33,9 @@
     "directory": "packages/gatsby-react-router-scroll"
   },
   "scripts": {
-    "build": "babel src --out-dir . --ignore \"**/__tests__\"",
+    "build": "babel src --out-dir . --ignore \"**/__tests__\" --extensions \".ts,.tsx\"",
     "prepare": "cross-env NODE_ENV=production npm run build",
-    "watch": "babel -w src --out-dir . --ignore \"**/__tests__\""
+    "watch": "babel -w src --out-dir . --ignore \"**/__tests__\" --extensions \".ts,.tsx\""
   },
   "engines": {
     "node": ">=10.13.0"

packages/gatsby-react-router-scroll/src/index.ts

@@ -0,0 +1,5 @@
+import { ScrollHandler as ScrollContext } from "./scroll-handler"
+import { ScrollContainer } from "./scroll-container"
+import { useScrollRestoration } from "./use-scroll-restoration"
+
+export { ScrollContext, ScrollContainer, useScrollRestoration }

packages/gatsby-react-router-scroll/src/scroll-container.tsx

@@ -0,0 +1,64 @@
+import React from "react"
+import ReactDOM from "react-dom"
+import PropTypes from "prop-types"
+import { ScrollContext } from "./scroll-handler"
+import { SessionStorage } from "./session-storage"
+import { Location } from "@reach/router"
+import { Location as HLocation } from "history"
+
+const propTypes = {
+  scrollKey: PropTypes.string.isRequired,
+  shouldUpdateScroll: PropTypes.func,
+  children: PropTypes.element.isRequired,
+}
+
+type Props = {
+  scrollKey: string
+  shouldUpdateScroll?: Function
+  children: React.ReactNode
+}
+
+type PropsWithContextAndLocation = Props & {
+  context: SessionStorage
+  location: HLocation
+}
+
+class ScrollContainerImplementation extends React.Component<
+  PropsWithContextAndLocation
+> {
+  componentDidMount() {
+    const node = ReactDOM.findDOMNode(this) as Element
+    const { location, scrollKey } = this.props
+
+    if (!node) return
+
+    node.addEventListener("scroll", () => {
+      this.props.context.save(location, scrollKey, node.scrollTop)
+    })
+
+    const position = this.props.context.read(location, scrollKey)
+    node.scrollTo(0, position)
+  }
+
+  render() {
+    return this.props.children
+  }
+}
+
+export const ScrollContainer = (props: Props) => (
+  <Location>
+    {({ location }) => (
+      <ScrollContext.Consumer>
+        {context => (
+          <ScrollContainerImplementation
+            {...props}
+            context={context}
+            location={location}
+          />
+        )}
+      </ScrollContext.Consumer>
+    )}
+  </Location>
+)
+
+ScrollContainer.propTypes = propTypes

packages/gatsby-react-router-scroll/src/ScrollBehaviorContext.js → packages/gatsby-react-router-scroll/src/scroll-handler.tsx

@@ -1,14 +1,30 @@
 import React from "react"
+import { LocationContext } from "@reach/router"
 import PropTypes from "prop-types"
-import SessionStorage from "./StateStorage"
-
-export default class ScrollContext extends React.Component {
-  constructor(props, context) {
-    super(props, context)
-
-    this._stateStorage = new SessionStorage()
+import { SessionStorage } from "./session-storage"
+
+export const ScrollContext = React.createContext<SessionStorage>(
+  new SessionStorage()
+)
+ScrollContext.displayName = `GatsbyScrollContext`
+
+type ShouldUpdateScrollFn = (
+  prevRouterProps: LocationContext | undefined,
+  routerProps: LocationContext
+) => boolean
+type ShouldUpdateScroll = undefined | ShouldUpdateScrollFn
+
+export class ScrollHandler extends React.Component<
+  LocationContext & { shouldUpdateScroll: ShouldUpdateScroll }
+> {
+  static propTypes = {
+    shouldUpdateScroll: PropTypes.func,
+    children: PropTypes.element.isRequired,
+    location: PropTypes.object.isRequired,
   }
 
+  _stateStorage: SessionStorage = new SessionStorage()
+
   scrollListener = () => {
     const { key } = this.props.location
 
@@ -23,17 +39,17 @@ export default class ScrollContext extends React.Component {
       this.props.location.key
     )
     if (scrollPosition) {
-      this.windowScroll(scrollPosition)
+      this.windowScroll(scrollPosition, undefined)
     } else if (this.props.location.hash) {
-      this.scrollToHash(decodeURI(this.props.location.hash))
+      this.scrollToHash(decodeURI(this.props.location.hash), undefined)
     }
   }
 
   componentWillUnmount() {
     window.removeEventListener(`scroll`, this.scrollListener)
   }
 
-  componentDidUpdate(prevProps) {
+  componentDidUpdate(prevProps: LocationContext): void {
     const { hash } = this.props.location
 
     const scrollPosition = this._stateStorage.read(
@@ -47,21 +63,30 @@ export default class ScrollContext extends React.Component {
     }
   }
 
-  windowScroll = (position, prevProps) => {
+  windowScroll = (
+    position: number,
+    prevProps: LocationContext | undefined
+  ): void => {
     if (this.shouldUpdateScroll(prevProps, this.props)) {
       window.scroll(0, position)
     }
   }
 
-  scrollToHash = (hash, prevProps) => {
+  scrollToHash = (
+    hash: string,
+    prevProps: LocationContext | undefined
+  ): void => {
     const node = document.querySelector(hash)
 
     if (node && this.shouldUpdateScroll(prevProps, this.props)) {
       node.scrollIntoView()
     }
   }
 
-  shouldUpdateScroll = (prevRouterProps, routerProps) => {
+  shouldUpdateScroll = (
+    prevRouterProps: LocationContext | undefined,
+    routerProps: LocationContext
+  ): boolean => {
     const { shouldUpdateScroll } = this.props
     if (!shouldUpdateScroll) {
       return true
@@ -72,12 +97,10 @@ export default class ScrollContext extends React.Component {
   }
 
   render() {
-    return this.props.children
+    return (
+      <ScrollContext.Provider value={this._stateStorage}>
+        {this.props.children}
+      </ScrollContext.Provider>
+    )
   }
 }
-
-ScrollContext.propTypes = {
-  shouldUpdateScroll: PropTypes.func,
-  children: PropTypes.element.isRequired,
-  location: PropTypes.object.isRequired,
-}

packages/gatsby-react-router-scroll/src/StateStorage.js → packages/gatsby-react-router-scroll/src/session-storage.ts

@@ -1,13 +1,14 @@
+import { Location } from "history"
 const STATE_KEY_PREFIX = `@@scroll|`
 const GATSBY_ROUTER_SCROLL_STATE = `___GATSBY_REACT_ROUTER_SCROLL`
 
-export default class SessionStorage {
-  read(location, key) {
+export class SessionStorage {
+  read(location: Location, key: string): number {
     const stateKey = this.getStateKey(location, key)
 
     try {
       const value = window.sessionStorage.getItem(stateKey)
-      return JSON.parse(value)
+      return value ? JSON.parse(value) : 0
     } catch (e) {
       if (process.env.NODE_ENV !== `production`) {
         console.warn(
@@ -23,11 +24,11 @@ export default class SessionStorage {
         return window[GATSBY_ROUTER_SCROLL_STATE][stateKey]
       }
 
-      return {}
+      return 0
     }
   }
 
-  save(location, key, value) {
+  save(location: Location, key: string, value: number) {
     const stateKey = this.getStateKey(location, key)
     const storedValue = JSON.stringify(value)
 
@@ -49,7 +50,7 @@ export default class SessionStorage {
     }
   }
 
-  getStateKey(location, key) {
+  getStateKey(location: Location, key: string) {
     const locationKey = location.key || location.pathname
     const stateKeyBase = `${STATE_KEY_PREFIX}${locationKey}`
     return key === null || typeof key === `undefined`

packages/gatsby-react-router-scroll/src/use-scroll-restoration.ts

@@ -0,0 +1,21 @@
+import { ScrollContext } from "./scroll-handler"
+import { useRef, useContext, useLayoutEffect } from "react"
+import { useLocation } from "@reach/router"
+
+export function useScrollRestoration(identifier: string) {
+  const location = useLocation()
+  const state = useContext(ScrollContext)
+  const ref = useRef<HTMLElement>()
+
+  useLayoutEffect(() => {
+    const position = state.read(location, identifier)
+    ref.current!.scrollTo(0, position)
+  }, [])
+
+  return {
+    ref,
+    onScroll() {
+      state.save(location, identifier, ref.current!.scrollTop)
+    },
+  }
+}