gistfile1.ts

import React, {
  useCallback,
  useEffect,
  useMemo,
  useRef,
  useState,
} from 'react';
import useConstant from 'use-constant';
import produce from 'immer';
import { useIsMountedFn } from './useIsMounted';

export type SyncSetState<S> = (stateUpdate: React.SetStateAction<S>) => void;
export type AsyncSetState<S> = (
  stateUpdate: React.SetStateAction<S>,
) => Promise<S>;
export type SyncStateProducer<S> = (stateProducer: (draft: S) => void) => void;
export type AsyncStateProducer<S> = (
  stateProducer: (draft: S) => void,
) => Promise<S>;

export type AwesomeState<S> = {
  initialState: S;
  getState: () => S;
  setState: SyncSetState<S>;
  setStateAsync: AsyncSetState<S>;
  produceState: SyncStateProducer<S>;
  produceStateAsync: AsyncStateProducer<S>;
};
export type AwesomeStateReturn<S> = [S, AwesomeState<S>];

export type UseAwesomeStateInitializer<S> = S | (() => S);

const useAsyncSetState = <S>(
  state: S,
  setState: SyncSetState<S>,
): AsyncSetState<S> => {
  // hold resolution function for all setState calls still unresolved
  const resolvers = useRef<((state: S) => void)[]>([]);

  // ensure resolvers are called once state updates have been applied
  useEffect(() => {
    resolvers.current.forEach(resolve => resolve(state));
    resolvers.current = [];
  }, [state, setState]);

  // make setState return a promise
  return useCallback(
    (stateUpdate: React.SetStateAction<S>) => {
      return new Promise<S>(resolve => {
        resolvers.current.push(resolve);
        setState(stateBefore => {
          const stateAfter =
            stateUpdate instanceof Function
              ? stateUpdate(stateBefore)
              : stateBefore;
          // If state does not change, we must resolve the promise because react won't re-render and effect will not resolve
          if (stateAfter === stateBefore) {
            resolve(stateAfter);
          }
          return stateAfter;
        });
      });
    },
    [setState],
  );
};

const useInitialState = <S>(
  initialStateArg: UseAwesomeStateInitializer<S>,
): S => {
  return useConstant(() => {
    return initialStateArg instanceof Function
      ? initialStateArg()
      : initialStateArg;
  });
};

// TODO enhance with logging configuration to help debugging
const useLoggingState = <S>(
  initialStateArg: UseAwesomeStateInitializer<S>,
): [S, SyncSetState<S>] => {
  const isMounted = useIsMountedFn();
  const initialState = useInitialState(initialStateArg);
  const [state, setState] = useState(initialState);
  //useEffect(() => log('state', state), [state]);
  const loggingSetState = useCallback(
    (stateUpdate: React.SetStateAction<S>) => {
      if (!isMounted()) {
        console.debug('setState while mounted: ignoring'); // TODO make this configurable
        return;
      }
      setState(stateBefore => {
        const stateAfter =
          stateUpdate instanceof Function
            ? stateUpdate(stateBefore)
            : stateUpdate;
        // log('setState (fn)', { stateAfter, stateBefore });
        return stateAfter;
      });
    },
    [setState],
  );

  return [state, loggingSetState];
};

const useGetState = <S>(state: S): (() => S) => {
  const stateRef = useRef(state);
  useEffect(() => {
    stateRef.current = state;
  });
  return useCallback(() => stateRef.current, [stateRef]);
};

const useStateProducer = <S>(
  setState: SyncSetState<S>,
): SyncStateProducer<S> => {
  return useCallback(
    producer => {
      return setState(state => produce(state, producer));
    },
    [setState],
  );
};

const useStateProducerAsync = <S>(
  setState: AsyncSetState<S>,
): AsyncStateProducer<S> => {
  return useCallback(
    producer => {
      return setState(state => produce(state, producer));
    },
    [setState],
  );
};

const useAwesomeState = <S>(
  initialStateArg: UseAwesomeStateInitializer<S>,
): AwesomeStateReturn<S> => {
  const initialState = useInitialState(initialStateArg);
  const [state, setState] = useLoggingState(initialState);
  const getState = useGetState(state);
  const setStateAsync = useAsyncSetState(state, setState);
  const produceState = useStateProducer(setState);
  const produceStateAsync = useStateProducerAsync(setStateAsync);

  /*
  useEffect(() => console.debug("initialState"),[initialState]);
  useEffect(() => console.debug("getState"),[getState]);
  useEffect(() => console.debug("setState"),[setState]);
  useEffect(() => console.debug("setStateAsync"),[setStateAsync]);
  useEffect(() => console.debug("produceState"),[produceState]);
  useEffect(() => console.debug("produceStateAsync"),[produceStateAsync]);
  */

  const api: AwesomeState<S> = useMemo(() => {
    return {
      initialState,
      getState,
      setState,
      setStateAsync,
      produceState,
      produceStateAsync,
    };
  }, [
    // All the api methods must should be stable!
    initialState,
    getState,
    setState,
    setStateAsync,
    produceState,
    produceStateAsync,
  ]);
  return useMemo(() => [state, api], [state, api]);
};

export default useAwesomeState;