精讀《zustand 原始碼》

黃子毅發表於2022-01-24
原始碼

zustand 是一個非常時髦的狀態管理庫，也是 2021 年 Star 增長最快的 React 狀態管理庫。它的理念非常函式式，API 設計的很優雅，值得學習。

概述

首先介紹 zustand 的使用方法。

建立 store

通過 create 函式建立 store，回撥可拿到 get set 就類似 Redux 的 getState 與 setState，可以獲取 store 瞬時值與修改 store。返回一個 hook 可以在 React 元件中訪問 store。

import create from 'zustand'

const useStore = create((set, get) => ({
  bears: 0,
  increasePopulation: () => set(state => ({ bears: state.bears + 1 })),
  removeAllBears: () => set({ bears: 0 })
}))

上面例子是全域性唯一的 store，也可以通過 createContext 方式建立多例項 store，結合 Provider 使用：

import create from 'zustand'
import createContext from 'zustand/context'

const { Provider, useStore } = createContext()

const createStore = () => create(...)

const App = () => (
  <Provider createStore={createStore}>
    ...
  </Provider>
)

訪問 store

通過 useStore 在元件中訪問 store。與 redux 不同的是，無論普通資料還是函式都可以存在 store 裡，且函式也通過 selector 語法獲取。因為函式引用不可變，所以實際上下面第二個例子不會引發重渲染：

function BearCounter() {
  const bears = useStore(state => state.bears)
  return <h1>{bears} around here ...</h1>
}

function Controls() {
  const increasePopulation = useStore(state => state.increasePopulation)
  return <button onClick={increasePopulation}>one up</button>
}

如果嫌訪問變數需要呼叫多次 useStore 麻煩，可以自定義 compare 函式返回一個物件：

const { nuts, honey } = useStore(state => ({ nuts: state.nuts, honey: state.honey }), shallow)

細粒度 memo

利用 useCallback 甚至可以跳過普通 compare，而僅關心外部 id 值的變化，如：

const fruit = useStore(useCallback(state => state.fruits[id], [id]))

原理是 id 變化時，useCallback 返回值才會變化，而 useCallback 返回值如果不變，useStore 的 compare 函式引用對比就會為 true，非常巧妙。

set 合併與覆蓋

set 函式第二個引數預設為 false，即合併值而非覆蓋整個 store，所以可以利用這個特性清空 store：

const useStore = create(set => ({
  salmon: 1,
  tuna: 2,
  deleteEverything: () => set({ }, true), // clears the entire store, actions included
}))

非同步

所有函式都支援非同步，因為修改 store 並不依賴返回值，而是呼叫 set，所以是否非同步對資料流框架來說都一樣。

監聽指定變數

還是用英文比較表意，即 subscribeWithSelector，這個中介軟體可以讓我們把 selector 用在 subscribe 函式上，相比於 redux 傳統的 subscribe，就可以有針對性的監聽了：

mport { subscribeWithSelector } from 'zustand/middleware'
const useStore = create(subscribeWithSelector(() => ({ paw: true, snout: true, fur: true })))

// Listening to selected changes, in this case when "paw" changes
const unsub2 = useStore.subscribe(state => state.paw, console.log)
// Subscribe also exposes the previous value
const unsub3 = useStore.subscribe(state => state.paw, (paw, previousPaw) => console.log(paw, previousPaw))
// Subscribe also supports an optional equality function
const unsub4 = useStore.subscribe(state => [state.paw, state.fur], console.log, { equalityFn: shallow })
// Subscribe and fire immediately
const unsub5 = useStore.subscribe(state => state.paw, console.log, { fireImmediately: true })

後面還有一些結合中介軟體、immer、localstorage、redux like、devtools、combime store 就不細說了，都是一些細節場景。值得一提的是，所有特性都是正交的。

精讀

其實大部分使用特性都在利用 React 語法，所以可以說 50% 的特性屬於 React 通用特性，只是寫在了 zustand 文件裡，看上去像是 zustand 的特性，所以這個庫真的挺會借力的。

建立 store 例項

任何資料流管理工具，都有一個最核心的 store 例項。對 zustand 來說，便是定義在 vanilla.ts 檔案的 createStore 了。

createStore 返回一個類似 redux store 的資料管理例項，擁有四個非常常見的 API：

export type StoreApi<T extends State> = {
  setState: SetState<T>
  getState: GetState<T>
  subscribe: Subscribe<T>
  destroy: Destroy
}

首先 getState 的實現：

const getState: GetState<TState> = () => state

就是這麼簡單粗暴。再看 state，就是一個普通物件：

let state: TState

這就是資料流簡單的一面，沒有魔法，資料儲存用一個普通物件，僅此而已。

接著看 setState，它做了兩件事，修改 state 並執行 listenser：

const setState: SetState<TState> = (partial, replace) => {
  const nextState = typeof partial === 'function' ? partial(state) : partial
  if (nextState !== state) {
    const previousState = state
    state = replace ? (nextState as TState) : Object.assign({}, state, nextState)
    listeners.forEach((listener) => listener(state, previousState))
  }
}

修改 state 也非常簡單，唯一重要的是 listener(state, previousState)，那麼這些 listeners 是什麼時候註冊和宣告的呢？其實 listeners 就是一個 Set 物件：

const listeners: Set<StateListener<TState>> = new Set()

註冊和銷燬時機分別是 subscribe 與 destroy 函式呼叫時，這個實現很簡單、高效。對應程式碼就不貼了，很顯然，subscribe 時註冊的監聽函式會作為 listener 新增到 listeners 佇列中，當發生 setState 時便會被呼叫。

最後我們看 createStore 的定義與結尾：

function createStore(createState) {
  let state: TState
  const setState = /** ... */
  const getState = /** ... */
  /** ... */
  const api = { setState, getState, subscribe, destroy }
  state = createState(setState, getState, api)
  return api
}

雖然這個 state 是個簡單的物件，但回顧使用文件，我們可以在 create 建立 store 利用 callback 對 state 賦值，那個時候的 set、get、api 就是上面程式碼倒數第二行傳入的：

import { create } from 'zustand'

const useStore = create((set, get) => ({
  bears: 0,
  increasePopulation: () => set(state => ({ bears: state.bears + 1 })),
  removeAllBears: () => set({ bears: 0 })
}))

至此，初始化 store 的所有 API 的來龍去脈就梳理清楚了，邏輯簡單清晰。

create 函式的實現

上面我們說清楚瞭如何建立 store 例項，但這個例項是底層 API，使用文件介紹的 create 函式在 react.ts 檔案定義，並呼叫了 createStore 建立框架無關資料流。之所 create 定義在 react.ts，是因為返回的 useStore 是一個 Hooks，所以本身具有 React 環境特性，因此得名。

該函式第一行就呼叫 createStore 建立基礎 store，因為對框架來說是內部 API，所以命名也叫 api：

const api: CustomStoreApi = typeof createState === 'function' ? createStore(createState) : createState

const useStore: any = <StateSlice>(
  selector: StateSelector<TState, StateSlice> = api.getState as any,
  equalityFn: EqualityChecker<StateSlice> = Object.is
) => /** ... */

接下來所有程式碼都在建立 useStore 這個函式，我們看下其內部實現：

簡單來說就是利用 subscribe 監聽變化，並在需要的時候強制重新整理當前元件，並傳入最新的 state 給到 useStore。所以第一步當然是建立 forceUpdate 函式:

const [, forceUpdate] = useReducer((c) => c + 1, 0) as [never, () => void]

然後通過呼叫 API 拿到 state 並傳給 selector，並呼叫 equalityFn（這個函式可以被定製）判斷狀態是否發生了變化：

const state = api.getState()
newStateSlice = selector(state)
hasNewStateSlice = !equalityFn(
  currentSliceRef.current as StateSlice,
  newStateSlice
)

如果狀態變化了，就更新 currentSliceRef.current：

useIsomorphicLayoutEffect(() => {
  if (hasNewStateSlice) {
    currentSliceRef.current = newStateSlice as StateSlice
  }
  stateRef.current = state
  selectorRef.current = selector
  equalityFnRef.current = equalityFn
  erroredRef.current = false
})
useIsomorphicLayoutEffect 是同構框架常用 API 套路，在前端環境是 useLayoutEffect，在 node 環境是 useEffect：

說明一下 currentSliceRef 與 newStateSlice 的功能。我們看 useStore 最後的返回值：

const sliceToReturn = hasNewStateSlice
  ? (newStateSlice as StateSlice)
  : currentSliceRef.current
useDebugValue(sliceToReturn)
return sliceToReturn

發現邏輯是這樣的：如果 state 變化了，則返回新的 state，否則返回舊的，這樣可以保證 compare 函式判斷相等時，返回物件的引用完全相同，這個是不可變資料的核心實現。另外我們也可以學習到閱讀原始碼的技巧，即要經常跳讀。

那麼如何在 selector 變化時更新 store 呢？中間還有一段核心程式碼，呼叫了 subscribe，相信你已經猜到了，下面是核心程式碼片段：

useIsomorphicLayoutEffect(() => {
  const listener = () => {
    try {
      const nextState = api.getState()
      const nextStateSlice = selectorRef.current(nextState)
      if (!equalityFnRef.current(currentSliceRef.current as StateSlice, nextStateSlice)) {
        stateRef.current = nextState
        currentSliceRef.current = nextStateSlice
        forceUpdate()
      }
    } catch (error) {
      erroredRef.current = true
      forceUpdate()
    }
  }
  const unsubscribe = api.subscribe(listener)
  if (api.getState() !== stateBeforeSubscriptionRef.current) {
    listener() // state has changed before subscription
  }
  return unsubscribe
}, [])

這段程式碼要先從 api.subscribe(listener) 看，這使得任何 setState 都會觸發 listener 的執行，而 listener 利用 api.getState() 拿到最新 state，並拿到上一次的 compare 函式 equalityFnRef 執行一下判斷值前後是否發生了改變，如果改變則更新 currentSliceRef 並進行一次強制重新整理（呼叫 forceUpdate）。

context 的實現

注意到 context 語法，可以建立多個互不干擾的 store 例項：

import create from 'zustand'
import createContext from 'zustand/context'

const { Provider, useStore } = createContext()

const createStore = () => create(...)

const App = () => (
  <Provider createStore={createStore}>
    ...
  </Provider>
)

首先我們知道 create 建立的 store 是例項間互不干擾的，問題是 create 返回的 useStore 只有一個例項，也沒有 <Provider> 宣告作用域，那麼如何構造上面的 API 呢？

首先 Provider 儲存了 create 返回的 useStore：

const storeRef = useRef<TUseBoundStore>()
storeRef.current = createStore()

那麼 useStore 本身其實並不實現資料流功能，而是將 <Provider> 提供的 storeRef 拿到並返回：

const useStore: UseContextStore<TState> = <StateSlice>(
  selector?: StateSelector<TState, StateSlice>,
  equalityFn = Object.is
) => {
  const useProviderStore = useContext(ZustandContext)
  return useProviderStore(
    selector as StateSelector<TState, StateSlice>,
    equalityFn
  )
}

所以核心邏輯還是是現在 create 函式裡，context.ts 只是利用 ReactContext 將 useStore “注入” 到元件，且利用 ReactContext 特性，這個注入可以存在多個例項，且不會相互影響。

中介軟體

中介軟體其實不需要怎麼實現。比如看這個 redux 中介軟體的例子：

import { redux } from 'zustand/middleware'
const useStore = create(redux(reducer, initialState))

可以將 zustand 用法改變為 reducer，實際上是利用了函式式理念，redux 函式本身可以拿到 set, get, api，如果想保持 API 不變，則原樣返回 callback 就行了，如果想改變用法，則返回特定的結構，就是這麼簡單。

為了加深理解，我們看看 redux 中介軟體原始碼：

export const redux = ( reducer, initial ) => ( set, get, api ) => {
  api.dispatch = action => {
    set(state => reducer(state, action), false, action)
    return action
  }
  api.dispatchFromDevtools = true
  return { dispatch: (...a) => api.dispatch(...a), ...initial }
}

將 set, get, api 封裝為 redux API：dispatch 本質就是呼叫 set。

總結

zustand 是一個實現精巧的 React 資料流管理工具，自身框架無關的分層合理，中介軟體實現巧妙，值得學習。

討論地址是：精讀《zustand 原始碼》· Issue #392 · dt-fe/weekly

如果你想參與討論，請 點選這裡，每週都有新的主題，週末或週一釋出。前端精讀 - 幫你篩選靠譜的內容。

關注 前端精讀微信公眾號

<img width=200 src="https://img.alicdn.com/tfs/TB165W0MCzqK1RjSZFLXXcn2XXa-258-258.jpg">

版權宣告：自由轉載-非商用-非衍生-保持署名（創意共享 3.0 許可證）

相關文章