手写 Zustand：从零实现 React 轻量级状态管理库

为什么选择 Zustand?

在 React 开发中,组件间通信一直是个令人头疼的问题。当组件层级复杂时,通过 props 层层传递状态不仅代码冗余,维护成本也直线上升。这时候就需要一个中央状态管理库来解决这个痛点。

相比老牌的 Redux,Zustand 的优势非常明显:

• 极简 API:没有繁琐的 reducer、action、dispatch 概念
• 零样板代码:不需要包裹 Provider,直接创建 store 即用
• 性能优秀:基于订阅机制实现精准更新,避免无效渲染
• 体积小巧:核心代码仅 1KB 左右

正因如此,Zustand 在 GitHub 上已经收获了 4 万+ Star,成为近年来最受欢迎的 React 状态管理方案之一。

核心原理拆解

要手写 Zustand,首先需要理解其三大核心机制:

1. 状态存储与管理

Zustand 采用闭包方式存储状态,通过 createStore 创建一个独立的状态容器:

javascript

const createStore = (createState) => {
  let state;  // 闭包变量,存储状态
  const getState = () => state;
  // ... 其他方法
}

这种设计让状态完全脱离 React 组件树,既可以在组件内使用,也可以在组件外直接操作。

2. 订阅发布模式

这是 Zustand 的灵魂所在。当状态改变时,如何通知所有使用该状态的组件更新?答案是观察者模式:

• 发布者(Store) :维护一个订阅者列表 listeners
• 订阅者(组件) :通过 subscribe 注册监听函数
• 状态更新时:遍历执行所有订阅者的回调

javascript

const listeners = new Set();

const subscribe = (listener) => {
  listeners.add(listener);
  return () => listeners.delete(listener);  // 返回取消订阅函数
}

const setState = (partial, replace = false) => {
  // 更新状态后通知所有订阅者
  listeners.forEach(listener => listener(state, previousState));
}

3. 选择器(Selector)与精准更新

这是 Zustand 性能优化的关键。通过 selector 函数,组件可以只订阅自己关心的状态切片:

javascript

const count = useCounterStore((state) => state.count);

当 state.text 改变时,只订阅 count 的组件不会重新渲染。实现原理是在订阅回调中比较 selector 返回值:

javascript

api.subscribe((state, previousState) => {
  const newObj = selector(state);
  const oldObj = selector(previousState);
  if (newObj !== oldObj) {
    forceRender(Math.random());  // 仅当关心的状态变化才强制更新
  }
})

完整实现详解

第一步:创建 Store

createStore 函数负责初始化状态并返回操作 API:

javascript

const createStore = (createState) => {
  let state;
  const listeners = new Set();
  
  const getState = () => state;
  
  const setState = (partial, replace = false) => {
    const nextState = typeof partial === 'function' 
      ? partial(state) 
      : partial;
    
    if (!Object.is(nextState, state)) {
      const previousState = state;
      if (!replace) {
        // 默认浅合并,保留未修改的字段
        state = Object.assign({}, state, nextState);
      } else {
        // replace 模式直接替换整个 state
        state = nextState;
      }
      listeners.forEach(listener => listener(state, previousState));
    }
  }
  
  const subscribe = (listener) => {
    listeners.add(listener);
    return () => listeners.delete(listener);
  }
  
  const api = { setState, getState, subscribe };
  state = createState(setState, getState, api);
  return api;
}

关键细节:

• Object.is() 判断状态是否真正改变,避免无效更新
• setState 支持传入函数,方便基于旧状态计算新值
• subscribe 返回取消订阅函数,符合 React useEffect 清理机制

第二步:实现 Hook 适配层

useStore 将订阅机制桥接到 React 组件:

javascript

const useStore = (api, selector) => {
  const [, forceRender] = useState(0);
  
  useEffect(() => {
    const unsubscribe = api.subscribe((state, previousState) => {
      const newObj = selector(state);
      const oldObj = selector(previousState);
      if (newObj !== oldObj) {
        forceRender(Math.random());  // 强制重渲染
      }
    });
    return unsubscribe;  // 组件卸载时自动取消订阅
  }, []);
  
  return selector(api.getState());
}

这里用了一个巧妙的技巧:通过修改 state 触发组件更新,而不是直接操作 DOM。

第三步:暴露便捷的 create API

javascript

export const create = (createState) => {
  const api = createStore(createState);
  
  const useBoundStore = (selector) => useStore(api, selector);
  
  // 将 API 方法挂载到 Hook 上,支持在组件外调用
  Object.assign(useBoundStore, api);
  
  return useBoundStore;
}

Object.assign 这一步很关键,它让我们可以:

• 组件内:通过 useCounterStore(selector) 使用
• 组件外:通过 useCounterStore.setState() 直接操作状态

实战验证

基于上面的实现,我们创建一个计数器和文本编辑器共存的案例:

javascript

const useCounterStore = create((set) => ({
  count: 0,
  text: '初始文本',
  increment: () => set((state) => ({ count: state.count + 1 })),
  updateText: (newText) => set({ text: newText }),
}));

CountDisplay 组件只订阅 count:

javascript

const CountDisplay = () => {
  console.log('CountDisplay 渲染了');
  const count = useCounterStore((state) => state.count);
  const increment = useCounterStore((state) => state.increment);
  
  return (
    <div>
      <p>当前计数: {count}</p>
      <button onClick={increment}>增加</button>
    </div>
  );
};

TextDisplay 组件只订阅 text:

javascript

const TextDisplay = () => {
  console.log('TextDisplay 渲染了');
  const text = useCounterStore((state) => state.text);
  const updateText = useCounterStore((state) => state.updateText);
  
  return (
    <div>
      <p>当前文本: {text}</p>
      <input value={text} onChange={(e) => updateText(e.target.value)} />
    </div>
  );
};

验证结果:

• 修改文本时,控制台只打印 TextDisplay 渲染了
• 点击计数按钮时,控制台只打印 CountDisplay 渲染了

这证明了精准更新机制生效!没有使用的组件不会重新渲染,性能得到保障。

进阶:API 直接调用

得益于 Object.assign,我们可以在任何地方直接操作状态:

javascript

const handleBatchUpdate = () => {
  useCounterStore.setState((prev) => ({ 
    count: prev.count + 10, 
    text: '批量修改完成!' 
  }));
  
  // 同步读取最新状态(不触发渲染)
  console.log(useCounterStore.getState());
};

这在处理异步逻辑或非 React 环境(如 WebSocket 回调)时非常有用。

源码阅读的价值

通过手写 Zustand,我们收获了什么?

1. 设计模式的实战应用

• 观察者模式:订阅发布机制
• 闭包:状态隔离与持久化
• 高阶函数:create 返回定制化 Hook

2. React 性能优化技巧

• 通过 selector 避免无效渲染
• Object.is() 精准判断状态变化
• useEffect 清理函数自动取消订阅

3. 框架设计思路

为什么 Zustand 这么简单?因为它:

• 没有引入中间件、异步处理等复杂概念
• 直接利用 JS 闭包和 React Hooks,没有额外抽象
• API 设计符合直觉,学习成本极低

总结

Zustand 的核心只有 200 行代码,却解决了 React 状态管理的本质问题。通过手写实现,我们深刻理解了:

• 状态管理 = 存储 + 订阅 + 通知
• 性能优化 = 精准订阅 + 浅比较
• 好的 API = 隐藏复杂度 + 暴露灵活性

当你下次在项目中使用 Zustand 时,不妨打开 DevTools 观察组件的渲染次数,你会发现这个 1KB 的小库,背后有着极其精妙的设计哲学。