이 가이드는 Jotai의 핵심 구현을 이해하고자 하는 분들에게 유용합니다. 이 글은 완전한 핵심 구현 예제가 아니라 단순화된 버전을 다룹니다. Daishi Kato(@dai_shi)의 트윗 모음을 참고하여 작성되었습니다.

첫 번째 버전

간단한 예제부터 시작해 보겠습니다. 아톰은 단순히 설정 객체를 반환하는 함수입니다. 우리는 아톰과 그 상태를 매핑하기 위해 WeakMap을 사용합니다. WeakMap은 키를 메모리에 유지하지 않기 때문에, 아톰이 가비지 컬렉션되면 그 상태도 함께 가비지 컬렉션됩니다. 이는 메모리 누수를 방지하는 데 도움이 됩니다.

import { useState, useEffect } from 'react'

// 아톰 함수는 초기값을 포함한 설정 객체를 반환합니다.
export const atom = (initialValue) => ({ init: initialValue })

// 아톰의 상태를 추적해야 합니다.
// 메모리 누수를 방지하기 위해 WeakMap을 사용합니다.
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 훅은 현재 값과 아톰의 값을 업데이트하는 함수를 튜플로 반환합니다.
export const useAtom = (atom) => {
  const atomState = getAtomState(atom)
  const [value, setValue] = useState(atomState.value)
  useEffect(() => {
    const callback = () => setValue(atomState.value)

    // 동일한 아톰이 여러 컴포넌트에서 사용될 수 있으므로,
    // 컴포넌트가 언마운트될 때까지 아톰의 상태 변경을 계속해서 듣습니다.
    atomState.listeners.add(callback)
    callback()
    return () => atomState.listeners.delete(callback)
  }, [atomState])

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

    // 구독된 모든 컴포넌트에게 아톰의 상태가 변경되었음을 알립니다.
    atomState.listeners.forEach((l) => l())
  }

  return [value, setAtom]
}

여기 우리가 간단히 구현한 아톰을 사용한 예제가 있습니다. 카운터 예제

참고 트윗: Jotai의 내부 구조를 파헤치다

두 번째 버전

잠깐만요! 더 나은 방법이 있습니다. Jotai에서는 파생된 아톰(derived atom)을 만들 수 있습니다. 파생된 아톰은 다른 아톰에 의존하는 아톰입니다.

const priceAtom = atom(10)
const readOnlyAtom = atom((get) => get(priceAtom) * 2)
const writeOnlyAtom = atom(
  null, // 첫 번째 인자로 `null`을 전달하는 것이 관례입니다.
  (get, set, args) => {
    set(priceAtom, get(priceAtom) - args)
  },
)
const readWriteAtom = atom(
  (get) => get(priceAtom) * 2,
  (get, set, newPrice) => {
    set(priceAtom, newPrice / 2)
    // 동시에 여러 아톰을 설정할 수 있습니다.
  },
)

모든 의존성을 추적하기 위해 아톰의 상태에 하나의 속성을 더 추가해야 합니다. 아톰 X가 아톰 Y에 의존한다고 가정해 봅시다. 그러면 아톰 Y를 업데이트할 때 아톰 X도 업데이트해야 합니다. 이를 의존성 추적(dependency tracking)이라고 합니다.

const atomState = {
  value: atom.init,
  listeners: new Set(),
  dependents: new Set(),
}

이제 아톰의 상태를 읽고 쓰는 함수를 만들어야 합니다. 이 함수는 의존하는 아톰의 상태를 업데이트할 수 있어야 합니다.

import { useState, useEffect } from 'react'

export const atom = (read, write) => {
  if (typeof read === 'function') {
    return { read, write }
  }
  const config = {
    init: read,

    // read 함수에서 get은 아톰 값을 읽는 데 사용됩니다.
    // 이는 반응적이며 의존성을 추적합니다.
    read: (get) => get(config),

    // write 함수에서 get은 아톰 값을 읽는 데 사용되지만 추적되지 않습니다.
    // set은 아톰 값을 쓰는 데 사용되며 대상 아톰의 write 함수를 호출합니다.
    write:
      write ||
      ((get, set, arg) => {
        if (typeof arg === 'function') {
          set(config, arg(get(config)))
        } else {
          set(config, arg)
        }
      }),
  }
  return config
}

// 위와 동일하지만 상태에 하나의 추가 속성이 있습니다: dependents
const atomStateMap = new WeakMap()
const getAtomState = (atom) => {
  let atomState = atomStateMap.get(atom)
  if (!atomState) {
    atomState = {
      value: atom.init,
      listeners: new Set(),
      dependents: new Set(),
    }
    atomStateMap.set(atom, atomState)
  }
  return atomState
}

// 아톰이 기본형(primitive)이면 값을 반환합니다.
// 아톰이 파생형(derived)이면 부모 아톰의 값을 읽고
// 현재 아톰을 부모의 의존성 집합에 추가합니다 (재귀적으로).
const readAtom = (atom) => {
  const atomState = getAtomState(atom)
  const get = (a) => {
    if (a === atom) {
      return atomState.value
    }
    const aState = getAtomState(a)
    aState.dependents.add(atom) // XXX 추가만 합니다.
    return readAtom(a) // XXX 캐싱하지 않습니다.
  }
  const value = atom.read(get)
  atomState.value = value
  return value
}

// atomState가 수정되면 모든 의존하는 아톰에게 알려야 합니다 (재귀적으로).
// 이제 이 아톰에 의존하는 모든 컴포넌트의 콜백을 실행합니다.
const notify = (atom) => {
  const atomState = getAtomState(atom)
  atomState.dependents.forEach((d) => {
    if (d !== atom) notify(d)
  })
  atomState.listeners.forEach((l) => l())
}

// writeAtom은 필요한 매개변수와 함께 atom.write를 호출하고 notify 함수를 트리거합니다.
const writeAtom = (atom, value) => {
  const atomState = getAtomState(atom)

  // 'a'는 atomStateMap의 어떤 아톰입니다.
  const get = (a) => {
    const aState = getAtomState(a)
    return aState.value
  }

  // 'a'가 atom과 같으면 값을 업데이트하고 아톰에게 알린 후 반환합니다.
  // 그렇지 않으면 'a'에 대해 writeAtom을 호출합니다 (재귀적으로).
  const set = (a, v) => {
    if (a === atom) {
      atomState.value = v
      notify(atom)
      return
    }
    writeAtom(a, v)
  }

  atom.write(get, set, value)
}

export const useAtom = (atom) => {
  const [value, setValue] = useState()
  useEffect(() => {
    const callback = () => setValue(readAtom(atom))
    const atomState = getAtomState(atom)
    atomState.listeners.add(callback)
    callback()
    return () => atomState.listeners.delete(callback)
  }, [atom])
  const setAtom = (nextValue) => {
    writeAtom(atom, nextValue)
  }
  return [value, setAtom]
}

여기 파생된 아톰 구현을 사용한 예제가 있습니다. 파생된 카운터 예제

참고 트윗: 파생된 아톰 지원