Jotai State Management

Jotai provides a bottom-up atomic model that serves as a more efficient and flexible alternative to the top-down useContext + useState pattern.

Key Advantages over Context:

• Avoids Provider Hell: Instead of nesting many context providers at the root, you use atoms.
• Dynamic State: Adding or removing state doesn't require re-mounting component subtrees with new providers.
• Optimized Re-renders: Unlike standard Context, where any change to the provider value re-renders all consumers, Jotai optimizes renders based on specific atom dependencies, avoiding the need for manual memoization or selector functions.

Usage Comparison

In React Context, sharing multiple states requires creating multiple contexts and wrapping the tree in multiple providers. In Jotai, you simply define atoms and use the useAtom hook.

import { Provider, atom, useAtom } from 'jotai'

const atom1 = atom(0)
const atom2 = atom(0)

// Optional: Use Provider to scope state. 
// If omitted, Jotai uses a default global provider (Provider-less mode).
const Parent = ({ children }) => {
  return <Provider>{children}</Provider>
}

const Component1 = () => {
  const [state, setState] = useAtom(atom1)
}

const Component2 = () => {
  const [state, setState] = useAtom(atom2)
}

The atomWithListeners pattern allows you to create an atom and a corresponding hook that can listen to state changes. This is particularly useful for components that need to perform side effects (like updating other state or triggering external APIs) when an atom changes, but do not need to re-render themselves every time that atom's value updates.

atomWithListeners returns a tuple containing:

1. An atom: A standard Jotai atom that holds the value.
2. A hook (useListener): A hook that accepts a callback. This callback is executed every time the atom's value is set. The callback receives the get function, set function, the newVal, and the prevVal.

The hook also returns a cleanup function (via useEffect) to automatically remove the listener when the component unmounts.

import { useEffect } from 'react'
import {
  atom,
  useAtom,
  useSetAtom,
  Getter,
  Setter,
  SetStateAction,
} from 'jotai'

type Callback<Value> = (
  get: Getter,
  set: Setter,
  newVal: Value,
  prevVal: Value,
) => void

export function atomWithListeners<Value>(initialValue: Value) {
  const baseAtom = atom(initialValue)
  const listenersAtom = atom<Callback<Value>[]>([])
  const anAtom = atom(
    (get) => get(baseAtom),
    (get, set, arg: SetStateAction<Value>) => {
      const prevVal = get(baseAtom)
      set(baseAtom, arg)
      const newVal = get(baseAtom)
      get(listenersAtom).forEach((callback) => {
        callback(get, set, newVal, prevVal)
      })
    },
  )
  const useListener = (callback: Callback<Value>) => {
    const setListeners = useSetAtom(listenersAtom)
    useEffect(() => {
      setListeners((prev) => [...prev, callback])
      return () =>
        setListeners((prev) => {
          const index = prev.indexOf(callback)
          return [...prev.slice(0, index), ...prev.slice(index + 1)]
        })
    }, [setListeners, callback])
  }
  return [anAtom, useListener] as const
}

The atomWithBroadcast pattern allows you to share state between different browsing contexts (tabs, windows, frames, iframes, or workers) on the same origin using the BroadcastChannel API. This is an alternative to atomWithStorage when you want to avoid using localStorage.

Key Characteristics

• Communication: Uses BroadcastChannel to enable communication between contexts.
• Initialization Limitation: Unlike atomWithStorage, BroadcastChannel does not support receiving messages during the initialization phase. If you need to persist state across sessions (not just across tabs in a single session), you should combine this with local storage.
• Usage: Once created, the atom behaves like a standard Jotai atom, but updates made in one tab will be broadcast to all other tabs listening to the same key.

import { atom, useAtom } from 'jotai'

// Implementation of the atomWithBroadcast pattern
export function atomWithBroadcast<Value>(key: string, initialValue: Value) {
  const baseAtom = atom(initialValue)
  const listeners = new Set<(event: MessageEvent<any>) => void>()
  const channel = new BroadcastChannel(key)

  channel.onmessage = (event) => {
    listeners.forEach((l) => l(event))
  }

  const broadcastAtom = atom(
    (get) => get(baseAtom),
    (get, set, update: { isEvent: boolean; value: SetStateAction<Value> }) => {
      set(baseAtom, update.value)

      if (!update.isEvent) {
        channel.postMessage(get(baseAtom))
      }
    },
  )

  broadcastAtom.onMount = (setAtom) => {
    const listener = (event: MessageEvent<any>) => {
      setAtom({ isEvent: true, value: event.data })
    }

    listeners.add(listener)

    return () => {
      listeners.delete(listener)
    }
  }

  const returnedAtom = atom(
    (get) => get(broadcastAtom),
    (_get, set, update: SetStateAction<Value>) => {
      set(broadcastAtom, { isEvent: false, value: update })
    },
  )

  return returnedAtom
}

// Usage Example
const broadAtom = atomWithBroadcast('count', 0)

const ListOfThings = () => {
  const [count, setCount] = useAtom(broadAtom)

  return (
    <div>
      {count}
      <button onClick={() => setCount(count + 1)}>+1</button>
    </div>
  )
}

In its simplest form, an atom is a configuration object that holds an initial value. The useAtom hook manages the state of that atom by using a WeakMap to associate the atom with its current value and a set of listeners. This ensures that when an atom is garbage collected, its state is also removed, preventing memory leaks. Components subscribe to atom changes via useEffect, and updating an atom via the returned setter notifies all subscribed listeners.

import { useState, useEffect }

// atom function returns a config object which contains initial value
export const atom = (initialValue) => ({ init: initialValue })

// we need to keep track of the state of the atom.
// we are using weakmap to avoid memory leaks
const atomStateMap = new WeakMap()
const getAtomState = (atom) => {
  let atomState = atomStateMap.get(atom)
  if (!atomState) {
    atomState = { value: atom.init, listeners: new Set() }
    atomStateMap.set(atom, atomState)
  }
  return atomState
}

// useAtom hook returns a tuple of the current value
// and a function to update the atom's value
export const useAtom = (atom) => {
  const atomState = getAtomState(atom)
  const [value, setValue] = useState(atomState.value)
  useEffect(() => {
    const callback = () => setValue(atomState.value)

    // same atom can be used at multiple components, so we need to
    // keep listening for atom's state change till component is unmounted.
    atomState.listeners.add(callback)
    callback()
    return () => atomState.listeners.delete(callback)
  }, [atomState])

  const setAtom = (nextValue) => {
    atomState.value = nextValue

    // let all the subscribed components know that the atom's state has changed
    atomState.listeners.forEach((l) => l())
  }

  return [value, setAtom]
}

The atomWithToggle pattern creates a WritableAtom that manages a boolean state. It simplifies state management by providing a single setter function that can either toggle the current boolean value or force it to a specific state (true or false).

This avoids the boilerplate of creating a separate atom just to handle the toggle logic for an existing boolean atom.

import { useAtom } from 'jotai'
import { atomWithToggle } from 'XXX' // Replace 'XXX' with the actual module path

// Initialize with a specific value
const isActiveAtom = atomWithToggle(true)

const Toggle = () => {
  const [isActive, toggle] = useAtom(isActiveAtom)

  return (
    <>
      {/* Toggles the current state */}
      <button onClick={() => toggle()}>
        isActive: {isActive ? 'yes' : 'no'}
      </button>

      {/* Forces state to true */}
      <button onClick={() => toggle(true)}>force true</button>

      {/* Forces state to false */}
      <button onClick={() => toggle(false)}>force false</button>
    </>
  )
}

A Bunja can depend on another Bunja using bunja.use. This allows managing hierarchical lifetimes (e.g., a 'Page' Bunja that lives longer than a 'Modal' Bunja).

When a child Bunja is used via useBunja, it automatically ensures the parent Bunja is also instantiated. When the child Bunja is unmounted, its instance is destroyed. If the parent has no other active users, it is also destroyed.

const pageBunja = bunja(() => {
  const pageStateAtom = atom({})
  return { pageStateAtom }
})

const modalBunja = bunja(() => {
  // Accessing the parent Bunja
  const { pageStateAtom } = bunja.use(pageBunja)
  
  const modalStateAtom = atom((get) => ({
    ...get(pageStateAtom),
    modal: 'state',
  }))

  bunja.effect(() => {
    console.log('modal opened')
    return () => console.log('modal closed')
  })

  return { modalStateAtom }
})

function Modal() {
  const { modalStateAtom } = useBunja(modalBunja)
  const modalState = useAtomValue(modalStateAtom)
  // ...
}

When composing atoms, follow these architectural guidelines:

1. Prefer wide graphs over deep graphs: An atom that derives from several other atoms is better than a long chain of atoms where each derives from the previous one. Deep chains can hit the JavaScript call stack limit and cause overflows during reads.
2. Atoms as functions: Think of atoms as building blocks. Composing atoms is analogous to composing functions.
3. Handling non-serializable data: If you need to store a function inside an atom, wrap it in an object. Jotai expects the getter function of a derived atom to be a pure getter; passing a raw function as the atom value might cause confusion in the derivation logic.

// To store a function, wrap it in an object
const doublerAtom = atom({ callback: (n) => n * 2 })

// Usage in a component
// const [doubler] = useAtom(doublerAtom)
// const doubledValue = doubler.callback(50) // 100