use-complex-state.js

import React from 'react';

export const useComplexState = (initialState) => {
    const initialStateRef = React.useRef(null);
    if (initialStateRef.current === null) {
        // Only get initialState once. This is a sort of memoization of initalState functions.
        initialStateRef.current = typeof initialState === 'function' ? initialState() : initialState;
    }

    // useReducer returns [state, dispatch].
    // Dispatch here effectively becomes setState, as we aren't passing in / switching on actions.
    return React.useReducer((prevState, dispatchArgs) => {
        const updatedState = typeof dispatchArgs === 'function' ? dispatchArgs(prevState) : dispatchArgs;

        // Shallow compare states. Avoid updating for equal states. Avoids unecessary re-renders.
        const sameValues = Object.entries(updatedState).every(([stateKey, stateValue]) => (
            Object.is(stateValue, prevState[stateKey])
        ));

        if (sameValues) return prevState;

        return {
            ...prevState,
            ...updatedState,
        };
    }, initialStateRef.current);
};