从0～1实现一个简易的zustand

了解zustand的使用方式

ts 复制代码

import { Button, Avatar, Badge, Space, Image } from 'antd'
import { create } from 'zustand'


type Store = {
  count: number
  inc: () => void
}


const useStore = create<Store>()((set) => ({
  count: 1,
  inc: () => set((state) => ({ count: state.count + 1 })),
}))


function Counter() {
  const { count, inc } = useStore()


  return (
    <Space size="large">
      <Badge count={count}>
        <Image src="https://docs.pmnd.rs/zustand.ico" style={{ width: '60px', height: '60px' }} preview={false} />
      </Badge>
      <Button type="primary" onClick={inc}>one up</Button>
    </Space>
  )
}


export default Counter

通过官方用例可以发现，通过create函数创建出一个useStore并导出，在组件中，通过useStore的方式引入所需要的状态（当然这里也可以按需引入自己需要的状态，避免不必要的re-render）。

列出zustand主要实现的方法

ts 复制代码

const createStore = (createState) => {
    let state;  // store内部状态存储于state上
    const getState = () => state; 
    const setState = () => {}; // setState就是create接收函数的入参
    const subscribe = () => {}; // 每次订阅时将subscribe加入到listeners，subscribe的作用是触发组件重新渲染
    const api = { getState, setState, subscribe };
    state = createState(setState); // state的初始值就是createState的调用结果
    return api;
}

const useStore = (api, selector, equalityFn) => {};

export const create = (createState) => {
  const api = createStore(createState); // 拿到store，包含了全部操作store的方法
  const useBoundStore = (selector, equalityFn) =>
    useStore(api, selector, equalityFn);
  return useBoundStore;
};

createStore用来创建Store，接收createState参数，这里的createState就是create接收函数的入参，getStae获取状态，setState设置状态，subscribe，组件使用这个状态时会订阅这个Store，当状态改变时会重新渲染。
useBoundStore 接收 selector（从完整的状态中选取部分状态），equalityFn（用来对比选取状态是否发生变化，从而决定是否重新渲染）。
useStore 借助 useSyncExternalStoreWithSelector 完成订阅、状态选取、re-render 优化，返回选取的状态。
create 完成上述函数的组合。

useStore实现

在实现useStore前，我们需要对useStore第一个参数声明一下类型。

api对象中包含getState、setState、subscribe。

getState返回状态，状态是由用户自定义的，这里可以用泛型T表示

ts 复制代码

type GetState<T> = ()=>T

setState是设置状态，设置状态需要分三种情况，一种是整体的状态更新，一种是部分状态更新，还有一种是函数式更新（传入一个函数）

ts 复制代码

type SetState<T> = (
  partial: T | Partial<T> | ((state: T) => T | Partial<T>),
) => void

subscribe，这里因为 subscribe 会作为参数传入到 useSyncExternalStoreWithSelector 中，因此我们直接用 useSyncExternalStoreWithSelector 定义的类型就好了：

ts 复制代码

type Subscribe = Parameters<typeof useSyncExternalStoreWithSelector>[0]

最终效果：

ts 复制代码

type StoreApi<T> = {
  setState: SetState<T>
  getState: GetState<T>
  subscribe: Subscribe
}

useStore内部会调用useSyncExternalStoreWithSelector方法做re-render处理并且接收五个参数（subscribe、getState、setState、selector、equalityFn），因此很容易写出useStore的函数源码,

ts 复制代码

const useStore = <State, StateSlice>(
  api: StoreApi<State>,
  selector: (state: State) => StateSlice = api.getState as any,
  equalityFn: (a: StateSlice, b: StateSlice) => boolean
) => {
  const slice = useSyncExternalStoreWithSelector(
    api.subscribe,
    api.getState,
    api.getState,
    selector,
    equalityFn
  );
  return slice;
};

selector 可选传入，不传入默认返回全部状态，传入这里的 slice 值为调用 selector 的返回结果。

subscribe实现

接下来我们实现 subscribe 函数，当在组件中获取状态时需要对 Store 进行订阅，这样当 Store 内部状态发生变化时才能够通知组件完成 re-render。订阅函数需要接收一个函数参数（调用这个函数来完成组件的 re-render），并保存这个函数（这里用了 Set 结构来保存），最终需要返回一个函数，当组件卸载时会被调用，用来取消订阅

ts 复制代码

type StateCreator<T> = (setState: SetState<T>) => T

const createStore = <T>(createState: StateCreator<T>): StoreApi<T> => {
  const listeners = new Set<() => void>();
  let state: T;
  const setState = () => {};
  const getState = () => state;
  const subscribe: Subscribe = (subscribe) => {
    listeners.add(subscribe);
    return () => listeners.delete(subscribe);
  };
  const api = { setState, getState, subscribe };
  state = createState(setState);
  return api;
};

我们定义了一个 listeners 的 Set 结构，并将 subscribe 接收的参数保存到 listeners 中，这样当 Store 状态发生变化（也就是调用 setState）时，依次遍历 listeners 保存的所有函数来 re-render 所有订阅该 Store 的组件即可。

这里通知re-render的逻辑在useSyncExternalStoreWithSelector中，组件通过useStore方法拿到状态时，useSyncExternalStoreWithSelector方法就会在组件中运行useLayoutEffect和useEffect方法监听状态变化，如果发现Store改变就会re-render。

setState实现

ts 复制代码

const setState: SetState<T> = (partial) => {
  const nextState =
    typeof partial === "function"
      ? (partial as (state: T) => T)(state)
      : partial;
  if (!Object.is(nextState, state)) {
    state =
      typeof nextState !== "object" || nextState === null
        ? (nextState as T)
        : Object.assign({}, state, nextState);
    listeners.forEach((listener) => listener());
  }
};

setState接收三种情况的传参，所以我们需要对传入的参数做一下判断，如果是函数，则需要调用拿到具体的状态值，最后需要通过Object.is判断前后状态是否一致，不一致的情况就需要进行状态更新。