我们一起聊聊70 行代码实现 Zustand 核心功能

import { create } from 'zustand'

const useStore = create((set) => ({
  count: 1,
  // 通过 set 方法更新状态值，set 支持传入函数和状态对象值
  inc: () => set((state) => ({ count: state.count + 1 })),
}))

function Counter() {
  const count = useStore((state) => state.count)
  const inc = useStore((state) => state.inc)
  return (
    <div>
      <p>{`Count: ${count}`}</p>
      <button onClick={inc}>+1</button>
    </div>
  )
}

<div>
  <p>Count: <span id="value"></span></p>
  <button id="btn">+1</button>
</div>

import { createStore } from 'zustand@4.5.4/vanilla'

const store = createStore((set) => ({
  count: 1,
  inc: () => set((state) => ({ count: state.count + 1 })),
}))

const { getState, setState, subscribe, getInitialState } = store

window.onload = () => {
  const value = document.querySelector('#value')
  value.innerHTML = getInitialState().count // 设置 store 中 count 值
  // 使用 subscribe 订阅状态变化，并更新数值
  subscribe((state) => {
    value.innerHTML = state.count
  })

  const btn = document.querySelector('#btn')
  btn.onclick = () => {
    // 触发更新
    getState().inc()
  }
}

const create = (createState) => {
  let state
  let initialState
  const listeners = new Set()

  const setState = (partial, replace) => {
    // 判断是否为函数，为函数就调用，并传入当前状态值
    const nextState = typeof partial === 'function'
      ? partial(state)
      : state
    
    // 对比状态值是否有变化
    if (!Object.is(nextState, state)) {
      const previousState = state
      // 如果是替换整个状态值，或者状态值为基础值或 null，则直接赋值，不然使用 Object.aasign 合并状态值
      state = replace ?? (typeof nextState !== 'object' || nextState === null)
        ? nextState
        : Object.assign({}, state, nextState)
      
      // 触发订阅函数
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  const getState = () => state
  
  const getInitialState = () => initialState

  const subscribe = (listener) => {
    listeners.add(listener)
    // 返回一个取消订阅的方法
    return () => {
      listeners.delete(listener)
    }
  }

  // 清空订阅
  const destory = () => listeners.clear()

  const api = {
    setState,
    getState,
    getInitialState,
    subscribe,
    destory
  }

  // 调用 createState，createState 参数为 set、get 和 api 对象，函数返回状态初始值
  initialState = (state = createState(setState, getState, api))

  return api
}

export default create

const snapshot = useSyncExternalStore(subscribe, getSnapshot, getServerSnapshot?)

import createImpl from './vanilla'
import useSyncExternalStoreExports from 'use-sync-external-store/shim/with-selector'

const { useSyncExternalStoreWithSelector } = useSyncExternalStoreExports

const create = (createState) => {
  // 使用普通版本 zustand 创建一个支持发布订阅的数据源
  const api = createImpl(createState)

  // zustand 版本 Hook，参数为状态值选择器和判断状态是否变化函数
  const useBearStore = (selector, equiltyFn?) => 
    // 和 useSyncExternalStore 类似，不过支持传入 selector，获取部分数据
    useSyncExternalStoreWithSelector(
      api.subcribe,
      api.getState,
      api.getInitialState,
      selector,
      equiltyFn,
    )

  // 把 api 合并到 Hook 对象上
  Object.assign(useBearStore, api)

  return useBearStore
}

export default create

// 相比于 useSyncExternalStore ，多了 selector 和 isEqual 参数
function useSyncExternalStoreWithSelector(
  subscribe,
  getSnapshot,
  getServerSnapshot,
  selector,
  isEqual?,
) {
  const [getSelection, getServerSelection] = useMemo(() => {
    let memoizedSnapshot; // 缓存的整个状态值
    let memoizedSelection: Selection; // 缓存的使用 selector 选中的部分状态值
  
    const memoizedSelector = (nextSnapshot: Snapshot) => {
      const prevSnapshot = memoizedSnapshot
      const prevSelection = memoizedSelection

      // 如果整体状态值相等，直接返回缓存的 selector 选中的状态值。
      if (is(prevSnapshot, nextSnapshot)) {
        return prevSelection;
      }

      // 使用 selector 函数获取最新的状态值
      const nextSelection = selector(nextSnapshot);

      // 有传入判断状态是否相等函数，相等的话就返回上次的 selector 选中值。
      if (isEqual !== undefined && isEqual(prevSelection, nextSelection)) {
        // 记录最新的整体状态值
        memoizedSnapshot = nextSnapshot;
        return prevSelection;
      }

      // 记录最新一次的更新值
      memoizedSnapshot = nextSnapshot;
      memoizedSelection = nextSelection;
      // 返回 selector 函数获取最新的状态值
      return nextSelection;
    };
      
    const getSnapshotWithSelector = () => memoizedSelector(getSnapshot());
    return [getSnapshotWithSelector, () => {}];
  }, [getSnapshot, getServerSnapshot, selector, isEqual]);
  
  // 调用 useSyncExternalStore 方法，第二参数不是整体获取整个状态值，而是 selector 的状态值
  const value = useSyncExternalStore(
    subscribe,
    getSelection,
    getServerSelection,
  );
  return value
}