Intensive reading "zustand source code"

zustand is a very trendy state management library and the fastest growing React state management library for Star in 2021. Its concept is very functional, the API design is very elegant, and it is worth learning.

Overview

First, we will introduce how to use zustand

create store

Create a store through the create function, and the callback can get get set is similar to Redux's getState and setState . It can obtain the instantaneous value of the store and modify the store. Returns a hook to access the store in React components.

import create from 'zustand'

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

The above example is a globally unique store. You can also createContext and use it in conjunction with Provider:

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

const { Provider, useStore } = createContext()

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

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

visit store

Access the store in the component via useStore Unlike redux, both ordinary data and functions can be stored in the store, and functions are also obtained through selector syntax. Because function references are immutable, the second example below does not actually cause a re-render:

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 multiple times to access variables, you can customize the compare function to return an object:

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

fine-grained memo

Using useCallback you can even skip the normal compare and only care about the change of the external id value, such as:

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

The principle is that when the id changes, the useCallback will change, and useCallback useStore does not change, the comparison of the compare function of true will be 061ee0129a872f, which is very clever.

set merge and overwrite

set second parameter of the false , that is, the value is merged instead of overwriting the entire store, so you can use this feature to clear the store:

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

asynchronous

All functions support asynchrony, because modifying the store does not depend on the return value, but calls set , so whether it is asynchronous or not is the same for the data flow framework.

Listen to the specified variable

Or use English to compare the meaning, that is, subscribeWithSelector . This middleware allows us to use the selector in the subscribe function. Compared with the traditional redux subscribe, we can monitor it in a targeted manner:

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 })

There are also some combination middleware, immer, localstorage, redux like, devtools, combime store, which are all detailed scenarios. It is worth mentioning that all features are orthogonal.

intensive reading

In fact, most of the features are using React syntax, so it can be said that 50% of the features belong to the common features of React, but they are written in the zustand document, which looks like the features of zustand, so this library is really good at leveraging of.

Create a store instance

Any data flow management tool has a core store instance. For zustand, it is defined in the vanilla.ts file createStore up.

createStore returns a redux store-like data management instance with four very common APIs:

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

First, the implementation of getState

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

It's that simple and rude. Look at state again, it is an ordinary object:

let state: TState

That's the simple side of dataflow, no magic, just a plain object for data storage, that's all.

Then look at setState , it does two things, modify state and execute 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))
  }
}

Modifying state also very simple, the only important thing is listener(state, previousState) , so when were these listeners registered and declared? In fact, listeners is a Set object:

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

When the registration and destruction timings are subscribe and destroy function calls respectively, this implementation is very simple and efficient. The corresponding code is not posted, it is clear, subscribe registered when listening function as listener added to listeners queue occurs when setState when it will be called.

Finally, let's look at the definition and ending of createStore

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

Although this state is a simple object, but looking back at the documentation, we can create use the callback to assign a value to the state. At that time, set , get , and api were passed in the penultimate line of the above code:

import { create } from 'zustand'

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

So far, the ins and outs of all the APIs for initializing the store have been sorted out, and the logic is simple and clear.

Implementation of the create function

We have explained how to create a store instance above, but this instance is the underlying API, which is defined in the react.ts create function described in the document, and calls createStore create a framework-independent data stream. The create defined in react.ts is because the returned useStore is a Hooks, so it has React environment characteristics, hence the name.

The first line of this function calls createStore create the basic store, because it is an internal API for the framework, so the name is also called 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
) => /** ... */

Next, all the codes are creating useStore . Let's take a look at its internal implementation:

Simply put, it is to use subscribe monitor changes, and to force refresh the current component when needed, and pass in the latest state to useStore . So the first step is of course to create the forceUpdate function:

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

state by calling the API and pass it to the selector, and call equalityFn (this function can be customized) to determine whether the state has changed:

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

If the status changes, update currentSliceRef.current :

useIsomorphicLayoutEffect(() => {
  if (hasNewStateSlice) {
    currentSliceRef.current = newStateSlice as StateSlice
  }
  stateRef.current = state
  selectorRef.current = selector
  equalityFnRef.current = equalityFn
  erroredRef.current = false
})

useIsomorphicLayoutEffect isomorphic API routines common frame, at the front end environment useLayoutEffect , the node environment useEffect :

Explain the functions of currentSliceRef and newStateSlice Let's look at the final return value of useStore

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

The discovery logic is as follows: if the state changes, return the new state, otherwise return the old state, which ensures that when the compare function judges equal, the references to the returned objects are exactly the same. This is the core implementation of immutable data. In addition, we can also learn the skills of reading source code, that is, to skip reading frequently.

So how to update the store when the selector changes? There is also a piece of core code in the middle, which calls subscribe , I believe you have guessed it, the following is the core code snippet:

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
}, [])

This code first from api.subscribe(listener) look, which makes any setState will trigger listener executed, and listener use api.getState() get the latest state , and get the last compare function equalityFnRef execution to see if the value has changed before and after the judgment, if it is changed Update currentSliceRef and do a forced refresh (call forceUpdate ).

Implementation of context

Noting the context syntax, it is possible to create multiple store instances that do not interfere with each other:

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

const { Provider, useStore } = createContext()

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

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

First of all, we know create does not interfere with each other. The problem is that there is create returned by useStore , and there is no <Provider> , so how to construct the above API?

First, Provider stores the create returned by useStore :

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

Then useStore itself does not actually implement the data stream function, but takes and returns the <Provider> provided by 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
  )
}

So the core logic is still in the current create function, context.ts just uses ReactContext to useStore into the component, and using the ReactContext feature, this injection can have multiple instances without affecting each other.

middleware

Middleware doesn't really need to be implemented. For example, look at this example of redux middleware:

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

You can change the usage of zustand to reducer, which actually uses the functional concept. The redux function itself can get set, get, api . If you want to keep the API unchanged, you can just return the callback as it is. If you want to change the usage, return a specific structure. It's that simple.

In order to deepen our understanding, let's take a look at the source code of redux middleware:

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 as redux API: dispatch is essentially calling set .

Summarize

zustand is a sophisticated React data flow management tool. The layering independent of its own framework is reasonable, and the middleware is cleverly implemented, which is worth learning.

The discussion address is: Intensive Reading "zustand Source Code" Issue #392 dt-fe/weekly

If you would like to join the discussion, please click here , there are new topics every week, published on weekends or Mondays. Front-end intensive reading - help you filter reliable content.

attention to 161ee0129a8d52 front-end intensive reading WeChat public

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

Copyright notice: Free to reprint - non-commercial - non-derivative - keep attribution ( Creative Commons 3.0 License )