useAnimatedHeight.ts

import * as React from 'react'
import { SpringValue, easings, useSpring } from 'react-spring'

/**
 * Hook that animates height when args.animationKey changes
 *
 * Ex:
 *   const animatedBlock = useAnimatedHeight({
 *     animationKey: key,
 *   })
 *
 *   return (
 *     <animated.div style={animatedBlock.style}>
 *       <div ref={animatedBlock.measuredInnerRef}>
 *         {content}
 *       </div>
 *     </animated.div>
 *   )
 */
export function useAnimatedHeight<AnimationKey>(args: { animationKey: AnimationKey }): {
  style: {
    height: SpringValue<string | number>
  }
  measuredInnerRef: React.RefObject<HTMLDivElement>
} {
  const measuredInnerRef = React.useRef<HTMLDivElement>(null)
  const [prevAnimationKey, setPrevAnimationKey] = React.useState(args.animationKey)
  React.useLayoutEffect(() => {
    setPrevAnimationKey(args.animationKey)
  }, [args.animationKey])
  const [prevHeight, setPrevHeight] = React.useState<number | null>(null)
  const [animateHeight, setAnimateHeight] = React.useState<number | null>(null)

  const animatedStyle = useSpring({
    height: animateHeight ?? 'auto',
    onRest: React.useCallback(() => {
      setAnimateHeight(null)
    }, []),
    config: {
      duration: 400,
      easing: easings.easeOutQuart,
      clamp: true,
    },
  })

  if (args.animationKey !== prevAnimationKey && prevHeight === null) {
    /**
     * This block runs right before we render, only if the
     * prevHeight value is not set.
     */
    const currentAnimatedValue = animatedStyle.height.get()
    const prevHeightOrCurrentAnimation =
      typeof currentAnimatedValue === 'number'
        ? currentAnimatedValue
        : measuredInnerRef.current?.getBoundingClientRect().height ?? 0

    animatedStyle.height.set(prevHeightOrCurrentAnimation)
    setPrevHeight(prevHeightOrCurrentAnimation)
  } else if (prevHeight !== null && args.animationKey === prevAnimationKey) {
    /**
     * This this block runs right after a render, only
     * if the prevHeight field is set.
     */
    setAnimateHeight(measuredInnerRef.current?.getBoundingClientRect().height ?? null)
    setPrevHeight(null)
  }

  return {
    style: animatedStyle,
    measuredInnerRef: measuredInnerRef,
  }
}