Zustand 핵심로직

MountionRiver·2025년 1월 25일
zustand

zustand의 핵심 로직에 대해서 알아보자

핵심로직 및 타입설정(vanilla)

Zustand의 핵심 로직 파일은 약 40줄의 핵심로직과 50 줄정도의 타입 설정으로 이루어져있다.

핵심로직

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) => {
  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)     // 리스너 등록
  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

위 로직은 최종적으로 createStore를 동작시키는 함수이다. createStore는 사용자가 인자를 제공할 경우 createStoreImpl(createState)를 즉시 실행, 없으면 createStoreImpl 반환하는 함수 이다.

그럼이제 createStoreImpl이 뭔지 알아보자

createStoreImpl

createStoreImpl은 setState, getState, getInitialState, subscribe 네가지 함수를 가지고 있다.
각각의 함수는

  • setState : 상태를 변경시키는 함수
  • getState : 현재 상태를 반환하는 함수
  • getInitialState : 초기 상태를 반환하는 함수
  • subscribe : 리스너를 추가하거나 제거하는 함수

setState

  const setState: StoreApi<TState>['setState'] = (partial, replace) => {
    // 새로운 상태 계산 - partial이 함수면 실행, 아니면 그대로 사용
    const nextState =
      typeof partial === 'function'
        ? (partial as (state: TState) => TState)(state)
        : partial
    
    // 상태가 변경되었을 때만 처리
    if (!Object.is(nextState, state)) {
      const previousState = state
      // replace가 true이거나 nextState가 객체가 아닐 경우 대체, 
      // 그렇지 않으면 기존 상태와 병합
      state =
        (replace ?? (typeof nextState !== 'object' || nextState === null))
          ? (nextState as TState)
          : Object.assign({}, state, nextState)
      // 모든 리스너에게 상태 변경 알림
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

이 함수는 새로운 상태를 인자로 줄 경우 함수를 실행하여 상태를 변경한 후, 모든 리스너에게 상태 변경을 알린다. 만약에 상태가 변경되지 않았을 경우 그대로 사용한다.

getState, getInitialState, subscribe

 // 현재 상태 반환
  const getState: StoreApi<TState>['getState'] = () => state
 
  // 초기 상태 반환
  const getInitialState: StoreApi<TState>['getInitialState'] = () => initialState
 
  // 리스너 구독 함수
  const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener)
  // 구독 해제 함수 반환
    return () => listeners.delete(listener)
  }

이는 각각 현재 상태, 초기상태, 리스너 구독 함수를 추가 반환 하는 함수이다.

initialState의 동작방식은 

const initialState = (state = createState(setState, getState, api))
  1. createState 함수에 setState, getState, api를 인자로 전달
  2. createState 함수가 실행되어 초기 상태 객체 반환
  3. 반환된 객체를 state 변수에 할당
  4. 이 state가 initialState가 된다.

이 createState는 상태를 변경하는 것이 아닌 생성하는 것이고, 생성한 이후 접근하는 방법은 존재하지 않기 때문에 초기값이 보존된다.

타입설정

이제 타입 설정에 대해 알아보자.

SetStateInternal

SetStateInternal는 상태 변경 업데이트를 위한 오버로드다.

type SetStateInternal<T> = {
  // 부분 업데이트를 위한 오버로드
  _(partial: T | Partial<T> | { _(state: T): T | Partial<T> }['_'], replace?: false): void
  // 전체 업데이트를 위한 오버로드
  _(state: T | { _(state: T): T }['_'], replace: true): void
}['_']

StoreApi

이제 상태 api를 정의하는 인터페이스이다. 위에서 본것처럼 네가지의 함수가 있다.

export interface StoreApi<T> {
  setState: SetStateInternal<T>
  getState: () => T
  getInitialState: () => T
  subscribe: (listener: (state: T, prevState: T) => void) => () => void
}

Mutate

이 타입은 스토어에 미들웨어를 체이닝하는 방식을 구현한다.

// 타입 유틸리티: T에서 K를 찾고 없으면 F 반환
type Get<T, K, F> = K extends keyof T ? T[K] : F

// 뮤테이터 체인 적용 타입
export type Mutate<S, Ms> = number extends Ms['length' & keyof Ms]
  ? S
  : Ms extends []
    ? S
    : Ms extends [[infer Mi, infer Ma], ...infer Mrs]
      ? Mutate<StoreMutators<S, Ma>[Mi & StoreMutatorIdentifier], Mrs>
      : never

사실 이렇게만 보면 이게 어떻게 쓰이는지 잘 구분이 안된다. 예시 코드를 하나 보자

const store = create(
  persist(
    (set) => ({
      count: 0,
      increment: () => set((state) => ({ count: state.count + 1 }))
    }),
    { name: 'count-storage' }
  )
)

이 코드는 Zustand 에서 사용하는 코드이다. 이때 위 Mutate 코드가 뮤테이터 체인 적용을 시켜,

type Result = Mutate
  StoreApi<{count: number}>, 
  [['persist', { name: string }]]
>

로 처리되어 persist 기능이 포함된 스토어 타입이 된다.

StateCreator

위 타입정의는 상태 저장소의 타입을 정의하는 함수이다.

  • T, Mis, Mos, U 의 타입을 받아서 setState, getState, store의 타입을 정의한다.
  • setState, getState, store은 타입이 사전에 정의 되면 정의된 타입과 미들웨어가 적용된 스토어의 타입을 지정하고, State가 없으면 never 타입 반환 함으로써 타입안정성을 강화한다.
export type StateCreator<
  T,    // 기본 상태 타입 (예: { count: number })
  Mis extends [StoreMutatorIdentifier, unknown][] = [],  // 입력 미들웨어 배열
  Mos extends [StoreMutatorIdentifier, unknown][] = [],  // 출력 미들웨어 배열
  U = T  // 반환 타입 (기본값은 T)
> = ((
  setState: Get<Mutate<StoreApi<T>, Mis>, 'setState', never>,
  getState: Get<Mutate<StoreApi<T>, Mis>, 'getState', never>,
  store: Mutate<StoreApi<T>, Mis>,
) => U) & { $$storeMutators?: Mos }

export interface StoreMutators<S, A> {}
export type StoreMutatorIdentifier = keyof StoreMutators<unknown, unknown>

CreateStore, CreateStoreImpl

마지막으로 CreateStore, CreateStoreImpl의 타입을 정의하고, Mutate 타입을 통해 미들웨어의 체이닝 방식을 구현한다.

type CreateStore = {
  <T, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
    initializer: StateCreator<T, [], Mos>,
  ): Mutate<StoreApi<T>, Mos>

  <T>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
    initializer: StateCreator<T, [], Mos>,
  ) => Mutate<StoreApi<T>, Mos>
}

type CreateStoreImpl = <
  T,
  Mos extends [StoreMutatorIdentifier, unknown][] = [],
>(
  initializer: StateCreator<T, [], Mos>,
) => Mutate<StoreApi<T>, Mos>

타입 코드

type SetStateInternal<T> = {
  _(
    partial: T | Partial<T> | { _(state: T): T | Partial<T> }['_'],
    replace?: false,
  ): void
  _(state: T | { _(state: T): T }['_'], replace: true): void
}['_']

export interface StoreApi<T> {
  setState: SetStateInternal<T>
  getState: () => T
  getInitialState: () => T
  subscribe: (listener: (state: T, prevState: T) => void) => () => void
}

type Get<T, K, F> = K extends keyof T ? T[K] : F

export type Mutate<S, Ms> = number extends Ms['length' & keyof Ms]
  ? S
  : Ms extends []
    ? S
    : Ms extends [[infer Mi, infer Ma], ...infer Mrs]
      ? Mutate<StoreMutators<S, Ma>[Mi & StoreMutatorIdentifier], Mrs>
      : never

export type StateCreator<
  T,
  Mis extends [StoreMutatorIdentifier, unknown][] = [],
  Mos extends [StoreMutatorIdentifier, unknown][] = [],
  U = T,
> = ((
  setState: Get<Mutate<StoreApi<T>, Mis>, 'setState', never>,
  getState: Get<Mutate<StoreApi<T>, Mis>, 'getState', never>,
  store: Mutate<StoreApi<T>, Mis>,
) => U) & { $$storeMutators?: Mos }

// eslint-disable-next-line @typescript-eslint/no-unused-vars, @typescript-eslint/no-empty-object-type
export interface StoreMutators<S, A> {}
export type StoreMutatorIdentifier = keyof StoreMutators<unknown, unknown>

type CreateStore = {
  <T, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
    initializer: StateCreator<T, [], Mos>,
  ): Mutate<StoreApi<T>, Mos>

  <T>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
    initializer: StateCreator<T, [], Mos>,
  ) => Mutate<StoreApi<T>, Mos>
}

type CreateStoreImpl = <
  T,
  Mos extends [StoreMutatorIdentifier, unknown][] = [],
>(
  initializer: StateCreator<T, [], Mos>,
) => Mutate<StoreApi<T>, Mos>
profile
MountionRiver
이전 포스트

Zustand presist를 뜯어보자.

다음 포스트

Zustand의 리액트와의 통합

0개의 댓글