useCachedRects.ts

import React, { useState, useEffect, useCallback } from 'react';

const flipMapKeys = (mapObj: Map<unknown, unknown> | null) => {
  if (!mapObj) return new Map();
  return new Map([...mapObj].map(([key, value]) => [value, key]));
};

/**
 * useCachedRects hook
 */
export function useCachedRects<Item>(
  itemNodeMap: React.RefObject<Map<Item, Element | null>>,
) {
  const [rects, setRects] = useState<Map<Item, DOMRect>>(new Map());

  const updateRects = useCallback(
    (nodes: (Element | null)[]) => {
      const nodeItemMap = flipMapKeys(itemNodeMap?.current);
      const rects: Array<[Item, DOMRect]> = [];

      for (const node of nodes) {
        if (nodeItemMap?.has(node)) {
          const item = nodeItemMap?.get(node);
          const rect = node?.getBoundingClientRect();
          rect && rects.push([item, rect]);
        }
      }

      setRects((prev) => {
        const newRects = [...prev, ...rects].filter(
          ([item]) => itemNodeMap.current?.get(item), // ignore null nodes
        );
        return new Map(newRects);
      });
    },
    [itemNodeMap],
  );

  // Resize
  // ---------
  const handleResize = useCallback(() => {
    requestAnimationFrame(() => {
      const nodes = Array.from(itemNodeMap.current?.values() || []);
      updateRects(nodes);
    });
  }, [itemNodeMap, updateRects]);

  useEffect(() => {
    const nodeMap = itemNodeMap?.current;
    if (!nodeMap || nodeMap.size === 0) return;

    const observer = new ResizeObserver(handleResize);
    [...nodeMap.values()].forEach((child: Element | null) => {
      child && observer.observe(child);
    });

    return () => observer.disconnect();
  }, [itemNodeMap, handleResize]);

  // Scroll
  // ---------
  useEffect(() => {
    const nodeMap = itemNodeMap?.current;
    if (!nodeMap || nodeMap.size === 0) return;

    let rafId: number;

    const handleScroll = () => {
      cancelAnimationFrame(rafId);
      rafId = requestAnimationFrame(() => updateRects([...nodeMap.values()]));
    };
    window.addEventListener('scroll', handleScroll);

    return () => window.removeEventListener('scroll', handleScroll);
  }, [itemNodeMap, updateRects]);

  return rects;
}
  
  
// USAGE
// ---------------------
const rectsCache = useCachedRects<Item>(childNodes);
const rect = rectsCache.get(item);