Zustand原理与源码分析

Zustand的基本用法

让我们先看一个简单的例子，展示Zustand的基本用法：

import create from 'zustand'

const useBlogStore = create<BlogState>((set) => ({
	age: 11,
	blog: {
		comments: {
			count: 11,
			list: [1, 2, 3],
		},
		count: 12,
	},
	setAge: (age) => set({ age }),
	setBlogCommentsCount: (count) =>
		set({ blog: { ...state.blog, comments: { ...state.blog.comments, count } } }),
}))

// 使用store的React组件
function BlogComponent() {
	const age = useBlogStore((state) => state.age)
	const blog = useBlogStore((state) => state.blog)
	const setAge = useBlogStore((state) => state.setAge)
	const setBlogCommentsCount = useBlogStore((state) => state.setBlogCommentsCount)

	return (
		<div>
			<div>Age: {age}</div>
			<button onClick={() => setAge(age + 1)}>增加年龄</button>

			<div>评论数: {blog.comments.count}</div>
			<div>文章数: {blog.count}</div>
			<button onClick={() => setBlogCommentsCount(blog.comments + 1)}>增加评论</button>
		</div>
	)
}

这个例子展示了两种类型的状态：

基本类型状态：age 是一个数字类型
复合类型状态：blog 是一个包含 comments 和 count 的对象

核心源码

Zustand源码主要包括两部分：

zustand store的实现，主要通过订阅发布模式实现
react中如何使用zustand store，主要通过一个自定义hook实现

// zustand/react.ts react和zustand的结合
export function useStore<TState, StateSlice>(
  api: ReadonlyStoreApi<TState>,
  selector: (state: TState) => StateSlice = identity as any,
) {
  const slice = React.useSyncExternalStore(
    api.subscribe,
    () => selector(api.getState()),
    () => selector(api.getInitialState()),
  )
  React.useDebugValue(slice)
  return slice
}


const createImpl = <T>(createState: StateCreator<T, [], []>) => {
  const api = createStore(createState)

  const useBoundStore: any = (selector?: any) => useStore(api, selector)

  Object.assign(useBoundStore, api)

  return useBoundStore
}

export const create = (<T>(createState: StateCreator<T, [], []> | undefined) =>
  createState ? createImpl(createState) : createImpl) as Create

// zustand/valinna.ts zustand核心代码，store的实现
const createStoreImpl: CreateStoreImpl = (createState) => {
  type TState = ReturnType<typeof createState>
  type Listener = (state: TState, prevState: TState) => void
  let state: TState
  const listeners: Set<Listener> = new Set()

  const setState: StoreApi<TState>['setState'] = (partial, replace) => {
    // TODO: Remove type assertion once https://github.com/microsoft/TypeScript/issues/37663 is resolved
    // https://github.com/microsoft/TypeScript/issues/37663#issuecomment-759728342
    const nextState =
      typeof partial === 'function'
        ? (partial as (state: TState) => TState)(state)
        : partial
    if (!Object.is(nextState, state)) {
      const previousState = state
      state =
        (replace ?? (typeof nextState !== 'object' || nextState === null))
          ? (nextState as TState)
          : Object.assign({}, state, nextState)
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  const getState: StoreApi<TState>['getState'] = () => state

  const getInitialState: StoreApi<TState>['getInitialState'] = () =>
    initialState

  const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener)
    // Unsubscribe
    return () => listeners.delete(listener)
  }

  const api = { setState, getState, getInitialState, subscribe }
  const initialState = (state = createState(setState, getState, api))
  return api as any
}

export const createStore = ((createState) =>
  createState ? createStoreImpl(createState) : createStoreImpl) as CreateStore

上述代码引用的就是zustand最核心的代码了。 valinna.ts文件中定义了整个 store 的实现，react.ts文件中定义了如何在React组件中使用 store。

createStoreImpl 函数

Zustand的核心是 createStoreImpl 函数，它创建了一个状态存储。让我们看看它的关键部分：结合上述代码，我们可以看到，createStoreImpl 函数主要做了以下几件事：

创建一个状态存储 state,一个观察者列表 listeners，一个初始化数据存储initialState。
创建一个状态更新函数 setState
创建一个状态获取函数 getState
创建一个状态订阅函数 subscribe
创建一个初始值获取函数 getInitialState
返回一个api对象，包含 setState、getState、getInitialState、subscribe 四个函数

调用getState会获取最新的状态，调用getInitialState获取初始值，调用subscribe 会在观察者中增加listen回调。比较重要的是setState调用，重点看这段代码：

const setState: StoreApi<TState>['setState'] = (partial, replace) => {
    // TODO: Remove type assertion once https://github.com/microsoft/TypeScript/issues/37663 is resolved
    // https://github.com/microsoft/TypeScript/issues/37663#issuecomment-759728342
    const nextState =
      typeof partial === 'function'
        ? (partial as (state: TState) => TState)(state)
        : partial
    if (!Object.is(nextState, state)) {
      const previousState = state
      state =
        (replace ?? (typeof nextState !== 'object' || nextState === null))
          ? (nextState as TState)
          : Object.assign({}, state, nextState)
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

setState 函数是Zustand的核心，它有几个重要的特点：

支持函数式更新和对象式更新 partial
使用浅比较nextState和state来判断状态是否发生变化
传入了replace参数或者nextState是非对象，则使用nextState替换整个state。

一般不要使用replace这个参数，或让nextState为基本类型或者null，这样会替换整个state，状态管理可能变得更加混乱。

默认使用浅拷贝来合并状态。

如果state中包含深层对象，更新深层对象时，需要手动合并整个对象，否则会丢失其他属性。
更新深层对象的某一个属性时，也会导致整个对象被替换。
组件中深层state对象更新后，所有依赖该state的子属性（非基本类型）的组件都会触发渲染

const useStore = create<BlogState>((set) => ({
	age: 11,
	blog: {
		comments: {
			count: 11,
			list: [1, 2, 3],
		},
		count: 12,
	},
	setAge: (age) => set({ age }),
	setBlogCommentsCount: (count) =>
		set({ blog: { ...state.blog, comments: { ...state.blog.comments, count } } }),
}))

如这里的代码，某个组件A依赖了blog的comments的list属性，当组件B调用setBlogCommentsCount更新了blog的comments的count属性时，组件A会重新渲染。

所以设计state时，需要考虑组件的依赖关系，避免不必要的渲染。 尽量让state扁平化，避免深层对象。

更新了状态后，会调用观察者列表listeners中的所有观察者，通知他们状态发生了变化

Store绑定到React的useStore Hook 实现

const createImpl = <T>(createState: StateCreator<T, [], []>) => {
  const api = createStore(createState)

  const useBoundStore: any = (selector?: any) => useStore(api, selector)

  Object.assign(useBoundStore, api)

  return useBoundStore
}

上述代码就是zustand/react的核心代码，即如何将store的数据和React组件关联。

useBoundStore 函数就是组件中const age = useBlogStore((state) => state.age)的useBlogStore函数，它接收一个可选的selector参数，返回一个store的子状态。

`useStore`的实现：

export function useStore<TState, StateSlice>(
  api: ReadonlyStoreApi<TState>,
  selector: (state: TState) => StateSlice = identity as any,
) {
  const slice = React.useSyncExternalStore(
    api.subscribe,
    () => selector(api.getState()),
    () => selector(api.getInitialState()),
  )
  React.useDebugValue(slice)
  return slice
}

Zustand的React绑定是通过useSyncExternalStore实现的。

useSyncExternalStore 是React 18引入的新Hook，能让你订阅外部Store。

让我们看看useSyncExternalStore的简化实现：

function useSyncExternalStore<T>(
	subscribe: (callback: () => void) => () => void,
	getSnapshot: () => T,
	getServerSnapshot?: () => T
): T {
	// 判断是否在服务端渲染
	const isServer = typeof window === 'undefined'

	// 获取当前快照
	const value = isServer ? getServerSnapshot?.() ?? getSnapshot() : getSnapshot()

	// 使用useState存储快照
	const [state, setState] = useState(value)

	// 使用useEffect订阅外部store的变化
	useEffect(() => {
		// 在服务端不进行订阅
		if (isServer) return

		// 订阅store的变化
		const unsubscribe = subscribe(() => {
			// 当store变化时，获取新的快照并更新state
			const nextValue = getSnapshot()
			setState(nextValue)
		})

		// 清理订阅
		return () => unsubscribe()
	}, [subscribe, getSnapshot, isServer])

	return state
}

结合useStore和useSyncExternalStore的实现，我们可以知道：

useStore会调用api.subscribe、api.getState、api.getInitialState函数，从而实现状态的订阅和获取。
每在一个组件中调用useStore，都会创建一个订阅，订阅store的变化。
当store发生变化时，会调用api.subscribe中的回调函数，获取最新的state状态，然后更新组件。

总结

Zustand通过简单的API和实现，提供了一个强大的状态管理解决方案。通过阅读源码，可以得到一些最佳实践：

尽量让state扁平化，避免深层对象。
调用set时，尽量不要传入replace参数，或让nextState为基本类型或者null，这样会替换整个state，状态管理可能变得更加混乱。
每个组件中调用useStore时，都会创建一个订阅，订阅store的变化。过多调用useStore，会导致订阅过多，影响性能。
每个组件尽量订阅的值为基本类型，避免订阅深层对象。可以考虑拆分订阅的state为多个值。而不是订阅一个复合对象，然后解构获取。

// good
const commentsCount = useBlogStore((state) => state.blog.comments.count)
const blogCount = useBlogStore((state) => state.blog.count)

// bad
const { count, comments: {count: commentsCount} } = useBlogStore(state => state.blog)

Zustand的基本用法

核心源码

createStoreImpl 函数

Store绑定到React的useStore Hook 实现

useStore的实现：

总结

`useStore`的实现：