What is Atomic State Management - Create One Yourself

Where does the state live?

The answer is not straightforward. React developers typically employ two strategies for structuring the application state: component state (using useState) and global store (using Redux). The state can either be closely linked to the component or stored in the Redux store, which means it is closely tied to the source and cannot be created independently.

Have you ever found yourself in a situation where you wanted to utilize the useState hook but also pass a reference to your state object? This is where the atomic state management model comes in.

Atoms

Atomic State Management involves utilizing Atoms as a central repository for state management. It can be seen as an upgraded version of the useState hook, allowing for state sharing between components. This approach combines the benefits of both component state and global store patterns. Atoms are purposefully designed to hold a singular value.

It’s short in writing and easy for sharing between components, as demonstrated in the example below.

// Example from jotai.org
const animeAtom = atom(animeAtom);

const Header = () => {
  const [anime, setAnime] = useAtom(animeAtom)
  ...
}

As you can see in the above example Atomic State Management model reduces boilerplate code compared to approaches like flux pattern and is very similar to React's useState hook.

TL;DR Use atomic state management techniques to achieve better flexibility in organizing application state management.

Build your own from scratch.

Before we proceed you can check the project on github. This implementation is for learning purposes, for production use check Jotai or Recoil.

Atom Creators / Factory implementation

let atomId = 0;

function atomFactory(payload) {
  const atom = {};

  const subscribers = new Map();

  let subscriberIds = 0;

  const key = atomId++;

  // This function returns the current value
  atom.get = function () {
    return atom.value;
  };

  // sets value and notify to subscribers
  atom.set = function (value) {
    atom.value = value;
    notify(value);
  };

  // notifier function to notify value
  function notify(value) {
    subscribers.forEach((subscriber) => {
      subscriber(value);
    });
  }

  // subscribe to changes; returns unsubscribe fn
  atom.subscribe = function (fn, initialId) {
    const id = initialId ?? (subscriberIds += 1);

    subscribers.set(id, fn);

    return () => void subscribers.delete(id);
  };

  // actual atom value
  atom.value = payload;

  return atom;
}

export { atomFactory as atom };

It is a very basic implementation of atom factory it returns an atom object.

// atom returned by factory fn

{
    get: () => void
    set: (value: any) => void
    subscribe: () => (() => void)
}

useAtom hook implementation

export function useAtom(atom) {
  const [state, setState] = useState(atom.get());

  useEffect(() => {
    // subscribe on mount and sets local state with new value (used for sync atom to reacts state)
    const unSubscribe = atom.subscribe(setState);

    // unsubscribe on unmount
    return () => unSubscribe();
  }, [atom]);

  // just setter function.
  const setAtomValue = useCallback((value) => atom.set(value), [atom]);

  return [state, setAtomValue];
}

uhhmmm.... it's good but we need a little bit of refactoring, we need useAtomValue / useAtomSetter hooks like Jotai to optimize rerenders.

useAtomValue and useAtomSetter Implementation

Here we are breaking useAtom hooks into two parts.

// useAtomValue
export function useAtomValue(atom) {
  const [state, setState] = useState(atom.get());

  useEffect(() => {
    const unSubscribe = atom.subscribe(setState);
    return () => unSubscribe();
  }, [atom]);

  return state;
}

// useAtomSetter
export function useAtomSetter(atom) {
  return useCallback((value) => atom.set(value), [atom]);
}

Refactored useAtom Hook

export function useAtom(atom) {
  return [useAtomValue(atom), useAtomSetter(atom)];
}

Usage

It's the same as Jotai

// Example from jotai.org
const animeAtom = atom('bleach');

const Header = () => {
  const [anime, setAnime] = useAtom(animeAtom)
  ...
}

Derived Atom Implementation.

// refactored atom factory fn
function atomFactory(payload) {
  const atom = {};

  const subscribers = new Map();

  let subscriberIds = 0;

  const key = atomId++;

  // getAtom function used to subscribe to another atom (for derived state)
  atom.getAtom = function (prevAtom) {
    prevAtom.subscribe(() => {
      if (payload instanceof Function) {
        atom.value = payload(atom.getAtom);
        notify(atom.value);
      }
    }, `atom_${key}`);

    return prevAtom.get();
  };

  atom.get = function () {
    return atom.value;
  };

  atom.set = function (value) {
    atom.value = value;
    notify(value);
  };

  function notify(value) {
    subscribers.forEach((subscriber) => {
      subscriber(value);
    });
  }

  atom.subscribe = function (fn, initialId) {
    const id = initialId ?? (subscriberIds += 1);

    subscribers.set(id, fn);

    return () => void subscribers.delete(id);
  };

  // check if the payload is a function (derived atom) or normal atom
  if (payload instanceof Function) {
    atom.value = payload(atom.getAtom);
  } else {
    atom.value = payload;
  }

  return atom;
}

export { atomFactory as atom };

useAtom will remain the same.

Derived atom example

import { atom, useAtom, useAtomValue } from './lib';

const priceAtom = createAtom(15);

const discountAtom = createAtom(10);

const discountedPriceAtom = createAtom((get) => {
    return (get(priceAtom) / 100) * get(discountAtom);
});

const Component = () => {
  const [price, setPrice] = useAtom(priceAtom);

  const discountedPrice = useAtomValue(discountedPriceAtom);
  ...
}

BONUS: atomWithLocalStorage Plugin

import { atom } from "./lib";

export function atomWithLocalStorage(key, payload) {
  //Create new atom
  const newAtom = atom(payload);

  // check value exists in localstorage or not
  const prevVal = JSON.parse(localStorage.getItem(key) || "null");

  if (prevVal) {
    // if the value exists in localstorage sets to atom
    newAtom.set(prevVal.data);
  }

  // subscribe to changes and set value in localstorage
  newAtom.subscribe((val) =>
    localStorage.setItem(key, JSON.stringify({ data: val }))
  );

  return newAtom;
}