useReactiveState.ts

import type {Dispatch, SetStateAction} from "react";
import {useCallback, useState} from "react";

export const UNINITIALIZED = Symbol("UNINITIALIZED");

/**
 * `useState` but allows mutation during retrieval as well.
 *
 * `useReactiveState(value, (x) => x)` is equivalent to just `useState(value)`.
 * `getValue` will be called each render (except the first) with the current
 * value and whatever is returned will be set as the new state (but won't cause
 * a re-render).
 *
 * @example Typical pattern
 * const didColorChange = useDidChange(props.color);
 * const color = useReactiveState(props.color, function(current): string {
 *     if (didColorChange) {
 *         return props.color;
 *     } else {
 *         return current;
 *     }
 * });
 */
function useReactiveState<T>(
    initialValue: T,
    getValue: (current: T) => T,
): [value: T, setValue: Dispatch<SetStateAction<T>>] {
    return useReactiveStateLazy<T>(function (current) {
        if (current === UNINITIALIZED) {
            return initialValue;
        } else {
            return getValue(current);
        }
    });
}

/**
 * `useState` but allows mutation during retrieval as well.
 *
 * Same as `useReactiveState` but instead of passing an initial value
 * separately, `getValue` is called with `UNINITIALIZED` to get the initial
 * value.
 *
 * @see useReactiveState
 */
export function useReactiveStateLazy<T>(
    getValue: (current: T | typeof UNINITIALIZED) => T,
): [value: T, setValue: Dispatch<SetStateAction<T>>] {
    const [value, setValue] = useState<{current: T | typeof UNINITIALIZED}>({
        current: UNINITIALIZED,
    });

    value.current = getValue(value.current);

    return [
        value.current,
        useCallback(function (valueOrTransformer) {
            if (valueOrTransformer instanceof Function) {
                setValue(function (oldValue) {
                    return {
                        current: valueOrTransformer(
                            oldValue.current as Exclude<
                                typeof oldValue.current,
                                typeof UNINITIALIZED
                            >,
                        ),
                    };
                });
            } else {
                setValue({current: valueOrTransformer});
            }
        }, []),
    ];
}

export default useReactiveState;