特定のdiv要素に対してScrollRestorationを設定できるようにした

概要

Tanstack react-routerを使用したプロジェクトで、特定のdiv要素のスクロール位置を制御させたかったのですが、Scroll Restorationというライブラリではwindow全体に対するスクロール制御しかできなかったため、仕方なくScroll Restorationのコードをローカルにコピーし、特定のdiv要素に対し設定できるようにカスタマイズしたお話です。

ライブラリのコードをTypeScriptに変換

まず、ローカルのnode_modules内にインストールしてあるライブラリ本体のコードを見に行きました。

@tanstack/react-router/dist/esm/scroll-restoration.js

import * as React from "react";
import { useRouter } from "./useRouter.js";
import { functionalUpdate } from "./utils.js";
const useLayoutEffect = typeof window !== "undefined" ? React.useLayoutEffect : React.useEffect;
const windowKey = "window";
const delimiter = "___";
let weakScrolledElements = /* @__PURE__ */ new WeakSet();
const sessionsStorage = typeof window !== "undefined" && window.sessionStorage;
const cache = sessionsStorage ? (() => {
  const storageKey = "tsr-scroll-restoration-v2";
  const state = JSON.parse(
    window.sessionStorage.getItem(storageKey) || "null"
  ) || { cached: {}, next: {} };
  return {
    state,
    set: (updater) => {
      cache.state = functionalUpdate(updater, cache.state);
      window.sessionStorage.setItem(storageKey, JSON.stringify(cache.state));
    }
  };
})() : void 0;
const defaultGetKey = (location) => location.state.key;
function useScrollRestoration(options) {
  const router = useRouter();
  useLayoutEffect(() => {
    const getKey = (options == null ? void 0 : options.getKey) || defaultGetKey;
    const { history } = window;
    history.scrollRestoration = "manual";
    const onScroll = (event) => {
      if (weakScrolledElements.has(event.target)) return;
      weakScrolledElements.add(event.target);
      let elementSelector = "";
      if (event.target === document || event.target === window) {
        elementSelector = windowKey;
      } else {
        const attrId = event.target.getAttribute(
          "data-scroll-restoration-id"
        );
        if (attrId) {
          elementSelector = `[data-scroll-restoration-id="${attrId}"]`;
        } else {
          elementSelector = getCssSelector(event.target);
        }
      }
      if (!cache.state.next[elementSelector]) {
        cache.set((c) => ({
          ...c,
          next: {
            ...c.next,
            [elementSelector]: {
              scrollX: NaN,
              scrollY: NaN
            }
          }
        }));
      }
    };
    if (typeof document !== "undefined") {
      document.addEventListener("scroll", onScroll, true);
    }
    const unsubOnBeforeLoad = router.subscribe("onBeforeLoad", (event) => {
      if (event.pathChanged) {
        const restoreKey = getKey(event.fromLocation);
        for (const elementSelector in cache.state.next) {
          const entry = cache.state.next[elementSelector];
          if (elementSelector === windowKey) {
            entry.scrollX = window.scrollX || 0;
            entry.scrollY = window.scrollY || 0;
          } else if (elementSelector) {
            const element = document.querySelector(elementSelector);
            entry.scrollX = (element == null ? void 0 : element.scrollLeft) || 0;
            entry.scrollY = (element == null ? void 0 : element.scrollTop) || 0;
          }
          cache.set((c) => {
            const next = { ...c.next };
            delete next[elementSelector];
            return {
              ...c,
              next,
              cached: {
                ...c.cached,
                [[restoreKey, elementSelector].join(delimiter)]: entry
              }
            };
          });
        }
      }
    });
    const unsubOnResolved = router.subscribe("onResolved", (event) => {
      if (event.pathChanged) {
        if (!router.resetNextScroll) {
          return;
        }
        router.resetNextScroll = true;
        const restoreKey = getKey(event.toLocation);
        let windowRestored = false;
        for (const cacheKey in cache.state.cached) {
          const entry = cache.state.cached[cacheKey];
          const [key, elementSelector] = cacheKey.split(delimiter);
          if (key === restoreKey) {
            if (elementSelector === windowKey) {
              windowRestored = true;
              window.scrollTo(entry.scrollX, entry.scrollY);
            } else if (elementSelector) {
              const element = document.querySelector(elementSelector);
              if (element) {
                element.scrollLeft = entry.scrollX;
                element.scrollTop = entry.scrollY;
              }
            }
          }
        }
        if (!windowRestored) {
          window.scrollTo(0, 0);
        }
        cache.set((c) => ({ ...c, next: {} }));
        weakScrolledElements = /* @__PURE__ */ new WeakSet();
      }
    });
    return () => {
      document.removeEventListener("scroll", onScroll);
      unsubOnBeforeLoad();
      unsubOnResolved();
    };
  }, [options == null ? void 0 : options.getKey, router]);
}
function ScrollRestoration(props) {
  useScrollRestoration(props);
  return null;
}
function useElementScrollRestoration(options) {
  var _a;
  const router = useRouter();
  const getKey = options.getKey || defaultGetKey;
  let elementSelector = "";
  if (options.id) {
    elementSelector = `[data-scroll-restoration-id="${options.id}"]`;
  } else {
    const element = (_a = options.getElement) == null ? void 0 : _a.call(options);
    if (!element) {
      return;
    }
    elementSelector = getCssSelector(element);
  }
  const restoreKey = getKey(router.latestLocation);
  const cacheKey = [restoreKey, elementSelector].join(delimiter);
  return cache.state.cached[cacheKey];
}
function getCssSelector(el) {
  const path = [];
  let parent;
  while (parent = el.parentNode) {
    path.unshift(
      `${el.tagName}:nth-child(${[].indexOf.call(parent.children, el) + 1})`
    );
    el = parent;
  }
  return `${path.join(" > ")}`.toLowerCase();
}
export {
  ScrollRestoration,
  useElementScrollRestoration,
  useScrollRestoration
};
//# sourceMappingURL=scroll-restoration.js.map

そして、そのコードをプロジェクトに合わせTypeScriptに書き換えました。
必要に応じてコード自体も添削しています。

customUseScrollRestoration

import { functionalUpdate, useRouter } from "@tanstack/react-router";
import { useEffect, useRef } from "react";

interface ScrollPosition {
  scrollX: number;
  scrollY: number;
}

export interface CacheState {
  cached: Record<string, Record<string, ScrollPosition>>;
  next: Record<string, ScrollPosition>;
}

const nextKey = "key";

let weakScrolledElements = new WeakSet();
const sessionsStorage = typeof window !== "undefined" && window.sessionStorage;
const cache = sessionsStorage
  ? (() => {
      const storageKey = "tsr-scroll-restoration-v2";
      const state = JSON.parse(
        window.sessionStorage.getItem(storageKey) ?? "null",
      ) ?? { cached: {}, next: {} };

      return {
        state,
        set: (updater: (c: CacheState) => CacheState) => {
          if (cache) {
            cache.state = functionalUpdate(updater, cache.state);
          }
          window.sessionStorage.setItem(
            storageKey,
            JSON.stringify(cache?.state),
          );
        },
      };
    })()
  : void 0;

export function useScrollRestoration() {
  const router = useRouter();
  const scrollableRef = useRef<HTMLDivElement>(null);

  useEffect(() => {
    const { history } = window;
    history.scrollRestoration = "manual";

    const onScroll = () => {
      const element = scrollableRef.current;
      if (!element || weakScrolledElements.has(element)) return;

      weakScrolledElements.add(element);

      if (!cache?.state.next[nextKey]) {
        cache?.set((c: CacheState) => ({
          ...c,
          next: {
            ...c.next,
            [nextKey]: {
              scrollX: element.scrollLeft ?? 0,
              scrollY: element.scrollTop ?? 0,
            },
          },
        }));
      }
    };

    const element = scrollableRef.current;
    if (element) {
      element.addEventListener("scroll", onScroll, true);
    }

    const unsubOnBeforeLoad = router.subscribe("onBeforeLoad", (event) => {
      if (!scrollableRef.current || !event.pathChanged) return;

      const entry = cache?.state.next[nextKey];
      if (!entry) return;

      entry.scrollX = scrollableRef.current.scrollLeft;
      entry.scrollY = scrollableRef.current.scrollTop;

      cache?.set((c: CacheState): CacheState => {
        const next = { ...c.next };
        if (!scrollableRef.current) return { cached: {}, next: {} };
        delete next[nextKey];
        return {
          ...c,
          next,
          cached: {
            ...c.cached,
            [`${event.fromLocation.state.key}`]: entry,
          },
        };
      });
    });

    const unsubOnResolved = router.subscribe("onResolved", async (event) => {
      await waitForScrollableRef(scrollableRef);

      if (
        !event.pathChanged ||
        !router.resetNextScroll ||
        !scrollableRef.current
      )
        return;

      const entry = cache?.state.cached[`${event.toLocation.state.key}`];

      scrollableRef.current.scrollLeft = entry ? entry.scrollX : 0;
      scrollableRef.current.scrollTop = entry ? entry.scrollY : 0;

      cache?.set((c: CacheState) => ({ ...c, next: {} }));
      weakScrolledElements = new WeakSet();
    });

    return () => {
      if (element) {
        element.removeEventListener("scroll", onScroll);
      }
      unsubOnBeforeLoad();
      unsubOnResolved();
    };
  }, [router]);

  // このrefをスクロールさせたいdivにセットする
  return { scrollableRef };
}

// refの要素が現れるまで待機させる
async function waitForScrollableRef(
  ref: React.RefObject<HTMLDivElement>,
): Promise<HTMLDivElement | null> {
  while (!ref.current) {
    await new Promise((resolve) => setTimeout(() => resolve(ref.current), 100));
  }
  return ref.current;
}

感想

もしかしたらライブラリの機能だけでなんとかできたのかもしれないです。
こうやってオープンソースのコードをカスタマイズするのは初めてだったのでいい経験になりました。