jotai/utils.mjs

👻 Next gen state management that will spook you

import { useMemo, useCallback } from 'react';
import { atom, useAtom } from 'jotai';

function useUpdateAtom(anAtom) {
  const writeOnlyAtom = useMemo(() => atom(null, (_get, set, update) => set(anAtom, update)), [anAtom]);
  writeOnlyAtom.scope = anAtom.scope;
  return useAtom(writeOnlyAtom)[1];
}

function useAtomValue(anAtom) {
  return useAtom(anAtom)[0];
}

const RESET = Symbol();
function atomWithReset(initialValue) {
  const anAtom = atom(initialValue, (get, set, update) => {
    if (update === RESET) {
      set(anAtom, initialValue);
    } else {
      set(anAtom, typeof update === 'function' ? update(get(anAtom)) : update);
    }
  });
  return anAtom;
}

function useResetAtom(anAtom) {
  const writeOnlyAtom = useMemo(() => atom(null, (_get, set, _update) => set(anAtom, RESET)), [anAtom]);
  writeOnlyAtom.scope = anAtom.scope;
  return useAtom(writeOnlyAtom)[1];
}

function useReducerAtom(anAtom, reducer) {
  const [state, setState] = useAtom(anAtom);
  const dispatch = useCallback(action => {
    setState(prev => reducer(prev, action));
  }, [setState, reducer]);
  return [state, dispatch];
}

function atomWithReducer(initialValue, reducer) {
  const anAtom = atom(initialValue, (get, set, action) => set(anAtom, reducer(get(anAtom), action)));
  return anAtom;
}

function atomFamily(initializeRead, initializeWrite, areEqual) {
  let shouldRemove = null;
  const atoms = new Map();

  const createAtom = param => {
    let item;

    if (areEqual === undefined) {
      item = atoms.get(param);
    } else {
      // Custom comparator, iterate over all elements
      for (let [key, value] of atoms) {
        if (areEqual(key, param)) {
          item = value;
          break;
        }
      }
    }

    if (item !== undefined) {
      if (shouldRemove != null && shouldRemove(item[1], param)) {
        atoms.delete(param);
      } else {
        return item[0];
      }
    }

    const newAtom = atom(initializeRead(param), initializeWrite && initializeWrite(param));
    atoms.set(param, [newAtom, Date.now()]);
    return newAtom;
  };

  createAtom.remove = param => {
    if (areEqual === undefined) {
      atoms.delete(param);
    } else {
      for (let [key] of atoms) {
        if (areEqual(key, param)) {
          atoms.delete(key);
          break;
        }
      }
    }
  };

  createAtom.setShouldRemove = fn => {
    shouldRemove = fn;
    if (!shouldRemove) return;

    for (let [key, value] of atoms) {
      if (shouldRemove(value[1], key)) {
        atoms.delete(key);
      }
    }
  };

  return createAtom;
}

function useSelector(anAtom, selector, equalityFn = Object.is) {
  const sliceAtom = useMemo(() => {
    let initialized = false;
    let prevSlice;
    const derivedAtom = atom(get => {
      const slice = selector(get(anAtom));

      if (initialized && equalityFn(prevSlice, slice)) {
        return prevSlice;
      }

      initialized = true;
      prevSlice = slice; // self contained mutation?

      return slice;
    });
    derivedAtom.scope = anAtom.scope;
    return derivedAtom;
  }, [anAtom, selector, equalityFn]);
  return useAtom(sliceAtom)[0];
}

function useAtomCallback(callback, scope) {
  const anAtom = useMemo(() => atom(null, (get, set, [arg, resolve, reject]) => {
    try {
      resolve(callback(get, set, arg));
    } catch (e) {
      reject(e);
    }
  }), [callback]);
  anAtom.scope = scope;
  const [, invoke] = useAtom(anAtom);
  return useCallback(arg => new Promise((resolve, reject) => {
    invoke([arg, resolve, reject]);
  }), [invoke]);
}

const deepFreeze = obj => {
  if (typeof obj !== 'object' || obj === null) return;
  Object.freeze(obj);
  const propNames = Object.getOwnPropertyNames(obj);

  for (const name of propNames) {
    const value = obj[name];
    deepFreeze(value);
  }

  return obj;
};

function freezeAtom(anAtom) {
  const frozenAtom = atom(get => deepFreeze(get(anAtom)), (_get, set, arg) => set(anAtom, arg));
  frozenAtom.scope = anAtom.scope;
  return frozenAtom;
}

const atomFrozen = (read, write) => {
  const anAtom = atom(read, write);
  const origRead = anAtom.read;

  anAtom.read = get => deepFreeze(origRead(get));

  return anAtom;
};

const atomFrozenInDev = typeof process === 'object' && process.env.NODE_ENV === 'development' ? atomFrozen : atom;

export { RESET, atomFamily, atomFrozenInDev, atomWithReducer, atomWithReset, freezeAtom, useAtomCallback, useAtomValue, useReducerAtom, useResetAtom, useSelector, useUpdateAtom };