Jotai

React 原始而灵活的状态管理

核心内部

本指南对于那些希望理解 Jotai 核心实现的人非常有益。它并不是核心实现的完整示例，而是一个简化版本。它受到了 Daishi Kato（@dai_shi）的一系列推文的启发。

第一版

让我们从一个简单的例子开始。原子就是一个将返回配置对象的函数。我们使用 WeakMap 来映射原子及其状态。 WeakMap 不会在内存中保留其键，所以如果一个原子被垃圾收集，它的状态也会被垃圾收集。这有助于避免内存泄漏。

import { useState, useEffect } from 'react'

// atom 函数返回一个包含初始值的配置对象
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 中，我们可以创建派生原子。派生原子是依赖于其他原子的原子。

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。这被称为依赖跟踪。

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 也是用来读取原子值的，但它不会被跟踪。
    // write 函数中的 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
}

// 如果原子是基本类型，我们返回它的值。
// 如果原子是派生的，我们读取父原子的值
// 并将当前原子添加到父原子的依赖集合中（递归）。
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]
}

这是一个使用我们的派生原子实现的例子。派生计数器示例

参考推文：支持派生原子

JotaiJotai

核心内部

第一版

第二版

Jotai