use-element-effect.ts

import React from "react";

/**
 * Hook to run effects related to a dom element.
 * Unlike useEffect, this hooks returns a ref callback.
 * The effect will be run every time the element changes.
 */
const useElementEffect = <T extends HTMLElement>(
  callback: (el: T) => (() => void) | undefined | void,
  deps: React.DependencyList,
): React.RefCallback<T> => {
  // eslint-disable-next-line react-hooks/exhaustive-deps
  const memoedCallback = React.useCallback(callback, deps);

  const unmountRef = React.useRef<(() => void) | undefined | void>(undefined);
  return React.useCallback(
    (node) => {
      if (node) {
        unmountRef.current?.();
        unmountRef.current = memoedCallback(node);
      } else {
        unmountRef.current?.();
        unmountRef.current = undefined;
      }
    },
    [memoedCallback],
  );
};