Practical application of Zustand, from theory to practice

import { useEffect, useRef } from 'react'
import { useClickAway } from 'react-use'
import useModal from '../hooks/useModal'

export default function Modal() {
  // стан модалки
  const modal = useModal()
  // посилання на модалку
  const modalRef = useRef(null)
  // посилання на вміст модалки
  const modalContentRef = useRef(null)

  useEffect(() => {
    if (!modalRef.current) return

    // показуємо/приховуємо модалку в залежності від значення індикатора `isOpen`
    // `showModal` и `close` - це нативні методи, що надаються HTML-елементом `dialog`
    if (modal.isOpen) {
      modalRef.current.showModal()
    } else {
      modalRef.current.close()
    }
  }, [modal.isOpen])

  // приховуємо модалку при натисканні за межами її вмісту
  useClickAway(modalContentRef, modal.close)

  if (!modal.isOpen) return null

  return (
    <dialog
      style={{
        padding: 0,
      }}
      ref={modalRef}
    >
      <div
        style={{
          padding: '1rem',
          display: 'flex',
          alignItems: 'center',
          gap: '1rem',
        }}
        ref={modalContentRef}
      >
        <div>Modal content</div>
        <button onClick={modal.close}>X</button>
      </div>
    </dialog>
  )
}

Define minimal styles in the index.css file:

body {
  margin: 0;
}

#root {
  min-height: 100vh;
  display: grid;
  place-content: center;
}

dialog::backdrop {
  background-color: rgba(0, 0, 0, 0.4);
}

Finally, let’s render the modal in the App.jsx file:

import Modal from './components/Modal'
import useModal from './hooks/useModal'

function App() {
  const modal = useModal()

  return (
    <>
      <button onClick={modal.open}>Open modal</button>
      <Modal />
    </>
  )
}

export default App

It was easy, wasn’t it? And all thanks to the magic of the create function.

The source code of the zustand can be found here. Since we will only be looking at the core functionality provided by this package, we are interested in 2 files – vanilla.ts and react.ts.

The code contained in the vanilla.ts file is an implementation of the publisher/subscriber, pub/sub pattern.

We create a file zustand/vanilla.js with the following content:

const createStoreImpl = (createState) => {
  // стан
  let state
  // обробники
  const listeners = new Set()

  // функція оновлення стану
  const setState = (partial, replace) => {
    // наступний стан
    const nextState = typeof partial === 'function' ? partial(state) : partial
    // якщо стан змінився
    if (!Object.is(nextState, state)) {
      // попередній/поточний стан
      const previousState = state
      // оновлюємо стан за допомогою `nextState` (якщо `replace === true` або значенням `nextState` є примітив)
      // або нового об'єкта, що об'єднує 'state' і 'nextState`
      state =
        replace ?? typeof nextState !== 'object'
          ? nextState
          : Object.assign({}, state, nextState)
      // запускаємо обробники
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  // функція вилучення стану
  const getState = () => state

  // функція передплати
  // `listener` - обробник `onStoreChange`
  // див. код хука `useSyncExternalStoreWithSelector` - про це трохи пізніше
  const subscribe = (listener) => {
    // додаємо/реєструємо обробник
    listeners.add(listener)
    // повертаємо функцію відписки
    return () => listeners.delete(listener)
  }

  // функція видалення всіх оброблювачів
  const destroy = () => {
    listeners.clear()
  }

  const api = { setState, getState, subscribe, destroy }
  // ініціалізуємо стан
  state = createState(setState, getState, api)
  // повертаємо методи
  return api
}

// в залежності від того, чи передається функція ініціалізації стану,
// повертаємо або `api`, або функцію `createStoreImpl`
export const createStore = (createState) =>
  createState ? createStoreImpl(createState) : createStoreImpl

I think everything is clear here. Let’s move on.

The code contained in the react.ts file represents the integration or implementation of the pub/sub React fiber in question.

We create a file zustand/react.js with the following content:

// `CommonJS`
import useSyncExternalStoreExports from 'use-sync-external-store/shim/with-selector'
import { createStore } from './vanilla.js'

const { useSyncExternalStoreWithSelector } = useSyncExternalStoreExports

export function useStore(api, selector = api.getState, equalityFn) {
  // отримуємо частину (зріз) стану
  const slice = useSyncExternalStoreWithSelector(
    api.subscribe,
    api.getState,
    api.getServerState || api.getState,
    selector,
    equalityFn,
  )
  // і повертаємо його
  return slice
}

const createImpl = (createState) => {
  // отримуємо методи, що повертаються функцією `createStore`
  const api =
    typeof createState === 'function' ? createStore(createState) : createState

 // визначаємо хук, що викликає хук `useStore` з переданою
 // функцією-селектором (`selector`) для отримання частини стану і
 // функцією порівняння (`equalityFn`) для визначення необхідності повторного рендерингу
  const useBoundStore = (selector, equalityFn) =>
    useStore(api, selector, equalityFn)

  // це потрібно для того, щоб мати можливість
  // викликати хук за межами компонента -
  // `useModal.getState()`
  Object.assign(useBoundStore, api)

  return useBoundStore
}

// можна отримати або хук `useBoundStore`, або функцію `createImpl`
export const create = (createState) =>
  createState ? createImpl(createState) : createImpl

Try replacing import { create } from ‘zustand’ with import { create } from ‘../zustand/react’ and useModal.js to make sure nothing has changed in terms of functionality. This is where the magic begins.

The useSyncExternalStoreWithSelector hook is an advanced version of the useSyncExternalStore hook (useSyncExternalStore and its variants are for some reason in a separate package). The difference between the two is that useSyncExternalStoreWithSelector takes 2 additional parameters:

• selector – a selector function for obtaining a part of the state (by default, the entire state is returned);

• equalityFn – a function for comparing the current and new states, which is used to determine the need for re-rendering.

Call useModal with a selector:

const isModalOpen = useModal((state) => state.isOpen)

Call useModal with selector and comparison function:

import { shallow } from 'zustand/shallow'

const { open, close } = useModal(({ open, close }) => ({ open, close }), shallow)

What is the useSyncExternalStore hook for?

Typically, React components read data from props, state, and context. However, sometimes a component may need to read time-varying data from storage outside of React. Such storage can be:

• a third-party state management library (such as zustand) that stores state outside of React;

• a browser API that exposes the mutated value and events to subscribe to its changes.

useSyncExternalStore accepts 2 mandatory and 1 optional parameter:

• subscribe(required parameter) – a function that accepts a callback parameter and a subscription to the repository. callback is called for any change to the repository. This causes the component to be re-rendered. subscribe must return the unsubscribe function from the repository;

• getSnapshot(mandatory parameter) – a function that returns a snapshot (snapshot) of the state from the storage consumed by the component. If the state has not changed, repeated calls to getSnapshot should return the same values. If the new state is different from the current state, React re-renders the component;

• getServerSnapshot(optional parameter) – a function that returns an initial state snapshot from the repository. It is used only in the process of server rendering of content and its hydration on the client.

useSyncExternalStore returns a snapshot of the storage for use in React’s rendering logic (loop).

Thus, useSyncExternalStore allows you to subscribe to changes to state residing in an external store in a way that is compatible with React’s competitive capabilities. The render cycle involves, among other things, calling the same sequence of hooks used by the component for initial and repeated renders. The same sequence of calling (and number) of hooks is ensured by the rules of using hooks. This is logical: calling hooks in a different sequence or in a smaller/larger number will lead to an inconsistency in the state of the component.

useSyncExternalStore makes our pub/sub (external store) part of the hook system that forms the component’s final state.

The hook code in question can be found here (mountSyncExternalStore followed by updateSyncExternalStore functions).

The “bare” mountSyncExternalStore looks like this:

function mountSyncExternalStore(
  subscribe,
  getSnapshot,
  getServerSnapshot,
) {
  // волокно
  const fiber = currentlyRenderingFiber
  // поточний/виконуваний хук
  const hook = mountWorkInProgressHook()

  // наступний знімок стану
  let nextSnapshot
  const isHydrating = getIsHydrating()
  if (isHydrating) {
    nextSnapshot = getServerSnapshot()
  } else {
    nextSnapshot = getSnapshot()
  }

  // читаємо поточний знімок зі сховища на кожному рендерингу
  // це порушує звичайні правила React і працює тільки завдяки тому,
  // що оновлення сховища завжди є синхронними
  hook.memoizedState = nextSnapshot
  const inst = {
    value: nextSnapshot,
    getSnapshot,
  }
  hook.queue = inst

 // тут плануються ефекти для підписки на сховище і 
 // для оновлення полів екземпляра (`inst`), 
 // які оновлюються при будь-якій зміні `subscribe`, `getSnapshot` або значення 
 // ці нутрощі нас не цікавлять

  return nextSnapshot
}

The differences between updateSyncExternalStore and mountSyncExternalStore are as follows:

// попередній знімок
const prevSnapshot = (currentHook || hook).memoizedState;
// чи змінилося стан?
const snapshotChanged = !is(prevSnapshot, nextSnapshot);
// якщо змінилося
if (snapshotChanged) {
  hook.memoizedState = nextSnapshot;
  markWorkInProgressReceivedUpdate();
}
const inst = hook.queue;

As a bonus, catch a slightly modified shallow function that allows you to compare objects in depth, which can find many uses:

function equal<T>(objA: T, objB: T): boolean {
  if (Object.is(objA, objB)) {
    return true
  }

  if (
    typeof objA !== 'object' ||
    objA === null ||
    typeof objB !== 'object' ||
    objB === null
  ) {
    return false
  }

  if (
    (Array.isArray(objA) && !Array.isArray(objB)) ||
    (Array.isArray(objB) && !Array.isArray(objA))
  ) {
    return false
  }

  if (objA instanceof Map && objB instanceof Map) {
    if (objA.size !== objB.size) return false

    for (const [key, value] of objA) {
      if (!Object.is(value, objB.get(key))) {
        return false
      }
    }
    return true
  }

  if (objA instanceof Set && objB instanceof Set) {
    if (objA.size !== objB.size) return false

    for (const value of objA) {
      if (!objB.has(value)) {
        return false
      }
    }
    return true
  }

  if (objA instanceof Date && objB instanceof Date) {
    return Object.is(objA.getTime(), objA.getTime())
  }

  const keysA = Object.keys(objA)
  if (keysA.length !== Object.keys(objB).length) {
    return false
  }

  return keysA.every((key) => equal(objA[key as keyof T], objB[key as keyof T]))
}

I hope you learned something new and spent your time in vain.

The article is made for educational purposes