React Hook: useComponentDidUpdate()

usage.tsx

const useComponentDidUpdate = useGetSnapshotBeforeUpdate(
  (_prevProps, prevState) => {
    if (!messagesWrapperRef.current || !messages.data || !prevState) return undefined;

    if (messagesWrapperRef.current.scrollHeight > (prevState.scrollHeight || 0)) {
      return prevState.scrollHeight;
    }

    return undefined;
  },
  undefined,
  { scrollHeight: messagesWrapperRef.current?.scrollHeight },
);

useComponentDidUpdate((_prevProps, _prevState, snapshot) => {
  if (snapshot && messagesWrapperRef.current) {
    messagesWrapperRef.current.scrollTop = messagesWrapperRef.current.scrollHeight - snapshot;
  }
});

use-get-snapshot-before-update.hook.ts

import { useEffect, useLayoutEffect, useRef } from 'react';

const usePrevPropsAndState = <Props, State>(props: Props, state: State) => {
  const prevPropsAndStateRef = useRef<{ props: Props | null; state: State | null }>({
    props: null,
    state: null,
  });
  const prevProps = prevPropsAndStateRef.current.props;
  const prevState = prevPropsAndStateRef.current.state;

  useEffect(() => {
    prevPropsAndStateRef.current = { props, state };
  });

  return { prevProps, prevState };
};

export const useGetSnapshotBeforeUpdate = <Props = any, State = any, Snapshot = any>(
  getSnapshotBeforeUpdateCb: (prevProps: Props | null, prevState: State | null) => Snapshot,
  props: Props,
  state: State,
) => {
  // get prev props and state
  const { prevProps, prevState } = usePrevPropsAndState<Props, State>(props, state);
  const snapshot = useRef<Snapshot | null>(null);

  // getSnapshotBeforeUpdate - not run on mount + run on every update
  const componentJustMounted = useRef(true);
  useLayoutEffect(() => {
    if (!componentJustMounted.current) {
      snapshot.current = getSnapshotBeforeUpdateCb(prevProps, prevState);
    }
    componentJustMounted.current = false;
  });

  const useComponentDidUpdate = (
    componentDidUpdateCb: (
      prevProps: Props | null,
      prevState: State | null,
      snapshot: Snapshot | null,
    ) => void,
  ) => {
    useEffect(() => {
      if (!componentJustMounted.current) {
        componentDidUpdateCb(prevProps, prevState, snapshot.current);
      }
    });
  };

  return useComponentDidUpdate;
};