70 行代码实现 Zustand 核心功能

前端目前主流的开发技术栈如 React、Vue 等都是状态数据驱动 UI 更新（即 UI = f(state)），所以状态管理是项目开发的重要一环。

React 和 Vue 除了自带的状态管理 API，同时还有一些功能强大的状态管理库可供选择。Vue 常见的状态管理库有 Vuex 和 Pinia，React 状态管理相对更多，有 redux、mobox、zustand、jotai 等等。

在 React 中，redux 还是最热门的状态管理库，相信你肯定在 React 开发中有使用过它。其他的状态库，都有各自的设计理念，在某些场景和开发规范，它们可能更适合你的项目。

本文将介绍 zustand 的核心实现，zustand 库和 redux 类似，都参考了 flux 设计理念，它一些特点如下：

易于上手，学习成本低
轻量级设计，gzip 压缩后仅 1KB
TypeScript 友好，有助于提升代码质量和开发体验
强大的可扩展性，通过中间件可以实现日志，数据持久化等能力
zustand 在设计上注重性能，采用高效的更新机制减少不必要的渲染，同时支持状态分片。

基于上述特点，zustand 还是比较受欢迎的，你可以看到 zustand 的使用量是排在前头的。

Zustand 的使用

zustand 的使用起来很简单, 使用 create 创建一个 useStore，可以把状态值和更新状态函数都保存在 state 中，随后在组件中调用即可。

jsx 复制代码

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>
  )
}

同时 zustand 核心代码也可以在普通 JS 中调用，把上述功能用普通 JS 实现就如下：

html 复制代码

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

js 复制代码

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()
  }
}

Zustand 的实现

Vanilla 版本

zustand 的核心实现非常简洁，我们先实现一个普通版本的 zustand，因为 react hook 版本也需要使用到它。从上面 zustand 使用案例代码可以看出，state 状态值不能直接修改，要通过 setState 来触发修改，这个和 redux 一致，对于通知状态变化则使用了发布订阅模式。

核心实现大概如下：

js 复制代码

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

React Hook

接着基于普通版本实现 React Hook 版本。在实现前，我们先了解一个 React 自带的 Hook - useSyncExternalStore。

useSyncExternalStore 的作用是让你可以订阅外部的状态源，当外部状态源发生变化时，React 会触发重选渲染。

js 复制代码

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

useSyncExternalStore 调用第一个参数 subscribe 订阅数据源变化，当数据源变化了，就触发重新渲染，并调用 getSnapshot 返回最新的状态值。

有了这个 Hook 的支持，我们就可以轻松实现 React Hook 版本 zustand。

js 复制代码

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

至此我们已经完成了 zustand 核心功能的代码编写。

拓展

useSyncExternalStoreWithSelector

React Hook 版本的 zustand 中使用到了 useSyncExternalStoreWithSelector，这个 Hook 是基于 useSyncExternalStore 实现的，可以简单了解下它的实现，简化版源码如下（去除了服务端渲染等内容）：

js 复制代码

// 相比于 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
}

总结

可以看到 zustand 核心代码还是很简洁的。通过实现核心代码，我们可以更好地理解和使用 zustand。有兴趣的同学可以继续了解下 zustand 插件相关的内容。