use-composed-ref.ts

import { useMemo } from "react";

export function useComposedRef<TElement>(handlers: RefHandler<TElement>[]) {
  // We actually want the handlers to be the dependency list.
  // eslint-disable-next-line react-hooks/exhaustive-deps
  return useMemo(() => composeRefs(handlers), handlers);
}

function composeRefs<TElement>(
  handlers: RefHandler<TElement>[],
): React.MutableRefObject<null | TElement> {
  let element: null | TElement = null;

  return {
    get current() {
      return element;
    },

    set current(value) {
      setRefs(value, handlers);
      element = value;
    },
  };
}

function setRefs<TElement>(
  element: null | TElement,
  handlers: RefHandler<TElement>[],
) {
  handlers.forEach((handler) => {
    // Ingore falsy handlers to allow syntaxes like:
    //
    // ```js
    // const child = React.Children.only(this.props.children)
    // return React.cloneElement(child, {
    //   ref: ref => setRefs(ref, [this.myRef, child.ref])
    // })
    // ```
    if (handler) {
      switch (typeof handler) {
        case "function": {
          handler(element);
          break;
        }

        case "object": {
          handler.current = element;
          break;
        }

        default:
          throw new Error(
            `Only refs of type function and React.createRef() are supported. Got ${typeof handler}.`,
          );
      }
    }
  });
}

type RefHandler<TElement> =
  | React.RefCallback<null | TElement>
  | React.MutableRefObject<null | TElement>
  | undefined
  | null;