useScrollRestoration for Next.js - V2

app-shell.tsx

import { useScrollRestoration } from "./use-scroll-restoration";

type AppShellProps = {
  content: React.ReactNode;
  topBar?: React.ReactNode;
  bottomBar?: React.ReactNode;
};

export const AppShell = ({ content, topBar, bottomBar }: AppShellProps) => {
  useScrollRestoration();

  return (
    <>
      <style jsx global>{`
    html, body, #__next {
      width: 100%; height: 100%;
    }`}
      </style>
      <div className={`AppShell w-full h-full grid grid-cols-1 grid-rows-[auto_minmax(0,1fr)_auto]`}>
        <header className="AppShell__topBar">{topBar}</header>
        <main className="AppShell__content h-full overflow-auto">{content}</main>
        <div className="AppShell__bottomBar">{bottomBar}</div>
      </div>
    </>
  );
};

use-scroll-restoration.ts

import Router, { type NextRouter } from "next/router";
import { useEffect } from "react";
import { LocalStorage } from "./utils.local-storage";

type PageUrl = string;
type ScrollPosition = { x: number, y: number; };
type PagesScrollPosition = Record<PageUrl, ScrollPosition>;
const _pagesScrollPosition = {
  getStorage: () => new LocalStorage<PagesScrollPosition>('app-shell-pagesScrollPosition'),
  getDomNode: () => window.document.getElementsByClassName('AppShell__content')[0],
  persistScrollPositionSnapshot: (url: PageUrl) => {
    const node = _pagesScrollPosition.getDomNode();
    if (!node) return;
    const storage = _pagesScrollPosition.getStorage();
    const scrollPosition: ScrollPosition = {
      x: node.scrollLeft,
      y: node.scrollTop,
    };
    storage.save({
      ...storage.retrieve() ?? {},
      [url]: scrollPosition,
    });
  },
  restoreScrollPositionSnapshot: (url: PageUrl) => {
    const storage = _pagesScrollPosition.getStorage();
    const pagesScrollPosition = storage.retrieve();
    if (!pagesScrollPosition) return;
    if (!pagesScrollPosition[url]) return;
    const node = _pagesScrollPosition.getDomNode();
    if (!node) return;
    const { x, y } = pagesScrollPosition[url];
    setTimeout(() => {
      node.scrollTo({ left: x, top: y, behavior: 'auto' });
    }, 50);
  },
  scrollToTop: () => {
    const node = _pagesScrollPosition.getDomNode();
    if (!node) return;
    setTimeout(() => {
      node.scrollTo({ left: 0, top: 0, behavior: 'auto' });
    }, 50);
  }
};

type LastNavigationEvent = { type: 'BACK_OR_FORWARD' | 'REGULAR_NAVIGATION'; };
const _lastNavigationEvent = {
  getStorage: () => new LocalStorage<LastNavigationEvent>('app-shell-lastNavigationEvent'),
  saveEventType: (type: LastNavigationEvent['type']) => {
    const storage = _lastNavigationEvent.getStorage();
    storage.save({ type });
  },
  getLastEvent: (): LastNavigationEvent => {
    const storage = _lastNavigationEvent.getStorage();
    return storage.retrieve() ?? { type: 'REGULAR_NAVIGATION' };
  },
  forget: () => {
    const storage = _lastNavigationEvent.getStorage();
    storage.delete();
  }
};

export const useScrollRestoration = () => {
  useEffect(() => {

    // Router.events.on handlers
    type NextEventHandler = Parameters<NextRouter['events']['on']>[1];
    // const beforeHistoryChange: NextEventHandler = (newUrl, options) => {
    // };
    // const hashChangeComplete: NextEventHandler = (newUrl, options) => {
    // };
    // const hashChangeStart: NextEventHandler = (newUrl, options) => {
    // };
    // const routeChangeError: NextEventHandler = (newUrl, options) => {
    // };
    const routeChangeStart: NextEventHandler = (newUrl, options) => {
      // Note:
      // this handler is invoked BEFORE transitioning to new page/route, 
      // (<Link> click or router.push() ) start the transition

      const oldUrl = Router.asPath;
      _pagesScrollPosition.persistScrollPositionSnapshot(oldUrl);
    };
    const routeChangeComplete: NextEventHandler = (newUrl, options) => {
      // Note:
      // this handler is invoked AFTER transitioning to new page/route, 
      if (_lastNavigationEvent.getLastEvent()?.type === 'BACK_OR_FORWARD') {
        _pagesScrollPosition.restoreScrollPositionSnapshot(newUrl);
      }
      if (_lastNavigationEvent.getLastEvent()?.type === 'REGULAR_NAVIGATION') {
        _pagesScrollPosition.scrollToTop();
      }
      _lastNavigationEvent.forget();
    };

    // Router.beforePopState
    const beforePopState: Parameters<NextRouter['beforePopState']>[0] = (state) => {
      // Note:
      // this handler is invoked BEFORE transitioning to new page/route 
      // but only when going "back" (or "forward") in history, usually via browser buttons
      // <Link> click or router.push() DOES NOT invoke this handler

      _lastNavigationEvent.saveEventType('BACK_OR_FORWARD');
      return true;
    };

    // Window.on('beforeunload')
    const onWindowBeforeUnload = (event: BeforeUnloadEvent) => {
      // Note:
      // This handler is invoked when the browser tab is going to be deactivated and closed
      const oldUrl = Router.asPath;
      _pagesScrollPosition.persistScrollPositionSnapshot(oldUrl);
      delete event['returnValue'];
    };

    // Router.events.on('beforeHistoryChange', beforeHistoryChange);
    // Router.events.on('hashChangeComplete', hashChangeComplete);
    // Router.events.on('hashChangeStart', hashChangeStart);
    // Router.events.on('routeChangeError', routeChangeError);
    Router.events.on('routeChangeStart', routeChangeStart);
    Router.events.on('routeChangeComplete', routeChangeComplete);
    Router.beforePopState(beforePopState);
    window.addEventListener('beforeunload', onWindowBeforeUnload);

    return () => {
      // Router.events.off('beforeHistoryChange', beforeHistoryChange);
      // Router.events.off('hashChangeComplete', hashChangeComplete);
      // Router.events.off('hashChangeStart', hashChangeStart);
      // Router.events.off('routeChangeError', routeChangeError);
      Router.events.off('routeChangeStart', routeChangeStart);
      Router.events.off('routeChangeComplete', routeChangeComplete);
      Router.beforePopState(() => true);
      window.removeEventListener('beforeunload', onWindowBeforeUnload);
    };
  }, []);
};

utils.local-storage.ts

export class LocalStorage<Value> {
  key: string;
  constructor(key: string) {
    this.key = key;
  }
  serialize(x: Value) {
    return JSON.stringify(x);
  }
  deserialize(x: string): Value {
    return JSON.parse(x);
  }
  save(newValue: Value) {
    window.localStorage.setItem(this.key, this.serialize(newValue));
  }
  retrieve() {
    const rawValue = window.localStorage.getItem(this.key);
    if (!rawValue) return null;
    return this.deserialize(rawValue);
  }
  delete() {
    window.localStorage.removeItem(this.key);
  }
}