React hook wrapping ResizeObserver

useMeasurements.tsx

import React, { useReducer, useRef, useCallback, useLayoutEffect } from 'react'
import _ from 'lodash'
import ResizeObserver from 'resize-observer-polyfill'
import { assertNever } from 'lib/misc/Assert'

type DOMRectKeys = keyof DOMRectReadOnly

interface State<T extends DOMRectKeys> {
  keys: T[]
  value: Pick<DOMRectReadOnly, T> | undefined
}

type Action = { type: 'update'; rect: DOMRectReadOnly } | { type: 'reset' }

function reducer<T extends DOMRectKeys>(state: State<T>, action: Action): State<T> {
  const update = (next: Pick<DOMRectReadOnly, T> | undefined) =>
    _.isEqual(state.value, next) ? state : { ...state, value: next }

  switch (action.type) {
    case 'update':
      return update(_.pick(action.rect, state.keys))
    case 'reset':
      return update(undefined)
    default:
      assertNever(action)
  }
}

type UseMeasurementsReturn<T extends DOMRectKeys> = {
  measurements: Pick<DOMRectReadOnly, T> | undefined
  refSetter: (node: null | HTMLElement) => void
}

export function useMeasurements<T extends DOMRectKeys>(keys: T[]): UseMeasurementsReturn<T> {
  const [state, dispatch] = useReducer(reducer, { keys, value: undefined })

  const nodeRef = useRef<HTMLElement | null>(null)
  const prevNodesRef = useRef<HTMLElement[]>([])
  const roRef = useRef<ResizeObserver | null>(null)

  const measure = useCallback(() => {
    if (nodeRef.current) {
      const rect = nodeRef.current.getBoundingClientRect()

      dispatch({ type: 'update', rect })
    } else {
      dispatch({ type: 'reset' })
    }
  }, [])

  useLayoutEffect(() => {
    measure()

    if (!roRef.current) {
      roRef.current = new ResizeObserver(measure)
      if (nodeRef.current) {
        roRef.current.observe(nodeRef.current)
      }
    }

    return () => {
      roRef.current?.disconnect()
      roRef.current = null
      prevNodesRef.current = []
    }
  }, [measure])

  const refSetter = useCallback((node: null | HTMLElement) => {
    if (node) {
      if (roRef.current) {
        roRef.current.observe(node)
      }
    } else {
      if (nodeRef.current) {
        prevNodesRef.current.push(nodeRef.current)
      }

      window.setTimeout(() => {
        prevNodesRef.current.forEach(n => {
          if (n !== nodeRef.current && roRef.current) {
            roRef.current.unobserve(n)
          }
        })
        prevNodesRef.current = []
      }, 0)
    }

    nodeRef.current = node
  }, [])

  return { measurements: state.value, refSetter }
}

interface UseMeasurementsProps<T extends DOMRectKeys> {
  keys: T[]
  children: (props: UseMeasurementsReturn<T>) => React.ReactElement
}

export const UseMeasurements = <T extends DOMRectKeys>({ keys, children }: UseMeasurementsProps<T>) =>
  children(useMeasurements(keys))