源码学习——手写 zustand

项目搭建

cra 创建 react 项目

npx create-react-app my-zustand

运行

npm start

安装 zustand

npm i zustand

使用 zustand

// App.js
import { create } from 'zustand';

const useUserStore = create((set) => ({
  firstName: '',
  lastName: '',
  updateFirstName: (firstName) => set(() => ({ firstName })),
  updateLastName: (lastName) => set(() => ({ lastName })),
}));

function App() {
  const { firstName, updateFirstName } = useUserStore((state) => state);
  return (
    <div
      style={{
        height: '100vh',
        display: 'flex',
        justifyContent: 'center',
        alignItems: 'center',
      }}
    >
      <div style={{ marginRight: '10px' }}>
        <div>firstName: {firstName}</div>
        <input onChange={(e) => updateFirstName(e.currentTarget.value)} />
      </div>

      <LastName />
    </div>
  );
}

function LastName() {
  const { lastName, updateLastName } = useUserStore((state) => state);
  return (
    <div>
      <div>lastName: {lastName}</div>
      <input onChange={(e) => updateLastName(e.currentTarget.value)} />
    </div>
  );
}

export default App;

实现 zustand

核心原理：

1. 基于闭包保存全局状态
2. 基于发布订阅模式实现响应式

// my-zustand.js
const createStore = (createState) => {
  let state;
  const listeners = new Set();

  const setState = (partial) => {
    const nextState = typeof partial === 'function' ? partial(state) : partial;

    if (!Object.is(nextState, state)) {
      const previousState = state;
      state = Object.assign({}, state, nextState);
      listeners.forEach((listener) => listener(state, previousState));
    }
  };

  const getState = () => state;

  const subscribe = (listener) => {
    listeners.add(listener);
    return () => listeners.delete(listener);
  };

  const destory = () => {
    listeners.clear();
  };

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

  state = createState(setState, getState, api);

  return api;
};

状态变了如何触发页面渲染？使用useState

// 模拟实现，有漏洞，仅用于表达思想
const useStore = (api, selector) => {
  const [, render] = useState(0);
  useEffect(() => {
    api.subscribe((state, previousState) => {
      const newState = selector(state);
      const oldState = selector(previousState);
      if (newState !== oldState) {
        render(Math.random());
      }
    });
  }, []);

  return selector(api.getState());
};

export const create = (createState) => {
  const api = createStore(createState);
  return (selector) => useStore(api, selector);
};

替换 create 函数后运行，功能正常

// import { create } from 'zustand'; ---
import { create } from './my-zustand'; // +++

使用 useSyncExternalStore 优化

事实上，react 提供了一个 hook，用来订阅外部 store，store 变化以后会触发 rerender

重构 useStore

const useStore = (api, selector) =>
  useSyncExternalStore(api.subscribe, () => selector(api.getState()));

运行后功能正常

最终源码

// my-zustand.js
import { useSyncExternalStore } from 'react';
const createStore = (createState) => {
  let state;
  const listeners = new Set();

  const setState = (partial) => {
    const nextState = typeof partial === 'function' ? partial(state) : partial;

    if (!Object.is(nextState, state)) {
      const previousState = state;
      state = Object.assign({}, state, nextState);
      listeners.forEach((listener) => listener(state, previousState));
    }
  };

  const getState = () => state;

  const subscribe = (listener) => {
    listeners.add(listener);
    return () => listeners.delete(listener);
  };

  const destory = () => {
    listeners.clear();
  };

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

  state = createState(setState, getState, api);

  return api;
};

const useStore = (api, selector) =>
  useSyncExternalStore(api.subscribe, () => selector(api.getState()));

export const create = (createState) => {
  const api = createStore(createState);
  return (selector) => useStore(api, selector);
};