Zustand 분석 1 - 상태 관리 코드를 살펴보자

Zustand는 크게 3 가지로 분석해볼 수 있다.

 

1. state 관리를 위한 vaniila.ts와 react.ts

2. mutate를 위한 middleware

3. 최적화를 위한 shallow

 

이번 섹션에서는 1번인 상태 관리 코드를 우선적으로 살펴보고자 한다.

 

상태 관리 코드는 2가지 파일에 나눠서 관리 중이다. 

첫 번째는 vanilla.ts 파일이고, 두 번째는 react.ts 파일이다.

react.ts 코드는 JS로 관리하는 상태관리 코드를 `useSyncExternalStore`라는 훅을 통해 리액트의 고유 상태로 넣어주는 코드이다. 

때문에, vanilla.ts 파일을 자세히 살펴보면 react.ts 파일을 해석하는 데 오래 걸리지 않는다.

 

`vanilla.ts`

이 파일을 살펴보면 크게 2가지 함수가 존재함을 확인할 수 있다.

 

먼저 `createStoreImpl` 함수이며, 이 함수는 store 생성 함수의 구현부이다.

// 타입 정의
type CreateStoreImpl = <
  T,
  Mos extends [StoreMutatorIdentifier, unknown][] = [],
>(
  initializer: StateCreator<T, [], Mos>, // 상태 생성 함수
) => Mutate<StoreApi<T>, Mos>;

// 함수
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
}

이 함수는 총 4개의 속성을 가진 객체를 반환한다.

setState는 상태를 변경해주는 함수, getState는 상태를 가져오는 함수, getInitialState는 초깃값을 가져오는 함수, subscribe는 상태를 구독할 수 있도록 해주는 함수다.

이 때, getInitialState는 SSR을 위해서 사용되는 함수이다. (`useSyncExternalStore`의 세번째 인자값에 사용된다.)

 

그 다음으로 vanilla.ts는 createStore라는 함수를 정의했다.

// 타입 정의 (함수 시그니처)
// CreateStore 타입은 Store를 생성하는 함수의 형태를 정의합니다.
// 상태 생성자를 인자로 받아 Mutate된 StoreApi를 반환합니다. 
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>;
};

// 함수 정의
export const createStore = ((createState) =>
  createState ? createStoreImpl(createState) : createStoreImpl) as CreateStore

이 함수는 매개변수가 없으면 createStoreImpl 함수를 그대로 반환하고, 매개변수로 createState 값이 있으면 createStoreImpl를 호출한 객체값을 반환받는다.

 

그렇다면 createStore에서 매개변수로 받게 되는 StateCreator 타입은 무엇일까? 그리고, Mutate 타입은 무엇일까?

 

먼저 Mutate 타입이다.

// Mutate 타입은 StoreMutators를 기반으로 상태를 변형하는 역할을 합니다.
// 다양한 StoreMutators의 길이에 따라 상태 변형 로직을 적용합니다.
export type Mutate<S, Ms> =
  // Ms가 동적 배열이면 상태 S 그대로 반환
  number extends Ms['length' & keyof Ms]
    ? S
    : // Ms가 빈 배열이면 상태 S 그대로 반환
    Ms extends []
    ? S
    : // Ms가 [[Mi, Ma], ...Mrs] 형태의 배열이면 변형을 적용 (재귀적으로 Mutate 호출)
    Ms extends [[infer Mi, infer Ma], ...infer Mrs]
    ? Mutate<StoreMutators<S, Ma>[Mi & StoreMutatorIdentifier], Mrs>
    : // 위의 모든 조건이 해당되지 않으면 never 반환
      never;
      
      
// StoreMutators는 상태와 Arguments을 인자로 받아 상태를 변형하는 로직을 정의하는 타입입니다.
// 기본적으로 빈 타입이지만, 커스텀하게 확장될 수 있습니다.
// 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>; // StoreMutators의 키를 사용하여 식별자 정의

`Mutate` 타입은 `Mutator`(미들웨어)를 적용할 때의 상태 타입을 유추할 수 있도록 해준다.

동적 배열인 경우 `S` 그대로 반환하고, 빈 배열(미들웨어 적용 안한 경우)도 `S`, 정적 배열인 경우 재귀적으로 `Mutate`를 적용하여 타입을 유추하고, 그외의 경우 never로 타입을 결정한다.

Mi: 각 mutator의 식별자 (Mutator identifier)
Ma: mutator에 전달될 인자 (Mutator arguments)
Mrs: 나머지 mutator들의 배열 (Mutators rest)

 

이번엔 `StateCreator` 타입이다.

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

Mis와 Mos는 각각 Mutators in과 상태 생성자 이후에 적용되는 mutators와 상태 생성자의 입력에 영향을 주는 mutators와 상태 생성자 이후에 적용되는 mutators를 나타낸다.

 

즉, 간단하게 설명하면 다음과 같다.

 

• Mis: 상태 생성자가 사용할 때 `setState`, `getState` 등을 변형함. (상태를 생성하거나 업데이트 할 때)
• Mos: 상태 생성자가 반환한 상태에 추가적으로 변형을 적용함. (상태 생성자가 반환한 후 추가적인 기능을 적용하거나 상태를 결합)

 

 

Zustand에서 mutators는 상태 업데이트나 조회 동작을 커스터마이징하거나, 상태에 특정 기능을 추가하는 데 사용된다.
Mis와 Mos와 연관된 대표적인 mutator들을 각각 몇 가지 예시로 설명해보자면 다음과 같다.

Mis에 적용할 수 있는 Mutator들
devtools
- 역할: 상태 관리에 DevTools를 연결하여 개발 중 상태 변화를 시각적으로 추적할 수 있게 해준다.
- 동작: setState, getState를 감싸서 상태 변화가 일어날 때마다 DevTools에 그 상태를 기록한다.

immer
- 역할: setState에서 Immer 라이브러리를 사용하여 불변성을 신경 쓰지 않고 상태를 직접 수정할 수 있게 해준다.
- 동작: setState 함수를 감싸서 내부적으로 Immer를 사용하여 상태의 불변성을 유지하면서도 직관적인 상태 수정이 가능하게 한다.

persist
- 역할: 상태를 localStorage나 sessionStorage에 자동으로 저장하고, 다시 로드할 수 있게 해준다.
- 동작: setState와 getState를 감싸서 상태가 변경될 때마다 이를 지정된 저장소에 저장하고, 애플리케이션이 다시 로드될 때 해당 저장소에서 상태를 가져온다.

Mos에 적용할 수 있는 Mutator들
combine
- 역할: 여러 상태를 하나로 결합하여 관리할 수 있게 해준다.
- 동작: 상태 생성자가 반환하는 객체를 변형하여 여러 상태를 하나의 스토어로 결합할 수 있다.

subscribeWithSelector
- 역할: 상태의 특정 부분에 대해서만 선택적으로 구독할 수 있는 기능을 제공한다.
- 동작: 반환된 상태에서 특정 부분을 구독하고, 해당 부분이 변경될 때만 리렌더링 등의 동작을 발생시킨다.

이러한 mutator들을 사용하면 상태의 관리와 업데이트를 더욱 유연하게 제어할 수 있다.
Mis는 상태 생성 시점에서 API 함수들을 변형하는 데 쓰이고, Mos는 상태가 반환된 후 추가적으로 변형하는 데 사용된다.

 

 

`react.ts`

이 함수에는 3가지 함수를 정의해 두었는데 가장 중요한 부분은 `useStore` 훅이다.

// 타입 정의
export function useStore<S extends ReadonlyStoreApi<unknown>>(
  api: S,
): ExtractState<S>;

export function useStore<S extends ReadonlyStoreApi<unknown>, U>(
  api: S,
  selector: (state: ExtractState<S>) => U,
): U;

// 함수
const identity = <T>(arg: T): T => arg;
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;
}

이 훅은 useSyncExternalStore를 활용하여 외부의 상태를 리액트의 상태로 연동시키도록 도와준다.

 

매개변수는 총 2가지이다.

1. `api`는 상태 관리 API이다. 이 API는 `getState`, `subscribe`, 그리고 `getInitialState` 메서드를 속성으로 제공할 수 있어야 한다.

 

2. `selector`는 `TState`에서 필요한 부분만 추출하는 함수다. 기본적으로 `identity`로 지정되어 있는데, 이는 상태의 전체를 반환하는 기본 선택자다. 즉, 사용자가 선택자를 제공하지 않으면 기본적으로 상태 전체를 구독하게 된다.

 

최종적으로 선택된 상태 조각을 반환합니다. 컴포넌트는 이 상태를 사용하여 필요한 부분만 구독하고 렌더링할 수 있습니다.

 

그 다음에 확인할 함수는 `createImpl` 함수이다.

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

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

  Object.assign(useBoundStore, api);

  return useBoundStore;
};

 

1. createState

• createImpl는 상태를 만드는 함수를 받는다.
  이 함수는 StateCreator 타입으로, `setState`, `getState`, `store`를 매개변수로 받는 함수이다.

2. createStore(createState)

• `createStore`는 `createState`를 사용하여 상태를 초기화하고, 상태 관리 API를 생성하는 함수이다.
  이 API는 `getState`, `subscribe`, `setState` 등의 메서드를 포함한다.

3. useBoundStore

• 이 변수는 `useStore` 훅을 래핑한 함수며, `selector`를 인자로 받아 `useStore`에서 상태를 구독하고, 선택된 상태 조각을 반환한다.
• `useBoundStore`는 상태 API(api)와 `useStore` 훅을 결합한 형태로, 선택자 함수가 없을 경우 기본적으로 전체 상태를 구독하게 됩니다.

4. Object.assign(useBoundStore, api)

• `useBoundStore` 함수에 `api` 객체의 모든 메서드를 복사하여 추가한다. 이렇게 하면 `useBoundStore` 자체가 상태 관리 API의 메서드를 모두 사용할 수 있게 된다.
  예를 들어, `useBoundStore.getState()`나 `useBoundStore.setState()` 같은 API 호출이 가능해진다.

5. return useBoundStore

• 최종적으로 useBoundStore를 반환한다. 이 함수는 상태를 구독하고 사용할 수 있을 뿐만 아니라 상태 관리 API의 모든 메서드를 함께 제공하는 상태 관리 훅이 된다.

마지막으로 create 함수는 그냥 vanilla.ts의 createStore와 유사하다.

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

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

 

이렇게 zustand가 `useState`와 같은 리액트 훅을 사용하지 않고 어떻게 리액트 내부에서 전역 상태관리를 할 수 있는지 대략적으로 확인해보았다. 아직 미들웨어를 확인해보지 못해서 Mutator 관련된 타입들을 정확히 이해하지 못했는데, 다음엔 미들웨어를 확인해보면서 좀 더 이해를 해봐야겠다.