Zustand presist를 뜯어보자.

MountionRiver·2024년 12월 23일

zustand의 presist가 어떤 작업을 하는지 이제 한줄한줄 알아보자

분류

presist는 크게 4개로 나뉘어져 있다
1. 타입정의
2. 스토리지 생성
3. 핵심 로직
4. 유틸리티 헬퍼 함수

타입 정의

export interface StateStorage {
  getItem: (name: string) => string | null | Promise<string | null>
  setItem: (name: string, value: string) => unknown | Promise<unknown>
  removeItem: (name: string) => unknown | Promise<unknown>
}

Copyexport type StorageValue<S> = {
  state: S
  version?: number
}

export interface PersistStorage<S> {
  getItem: (name: string) => StorageValue<S> | null | Promise<StorageValue<S> | null>
  setItem: (name: string, value: StorageValue<S>) => unknown | Promise<unknown>
  removeItem: (name: string) => unknown | Promise<unknown>
}

type JsonStorageOptions = {
  reviver?: (key: string, value: unknown) => unknown
  replacer?: (key: string, value: unknown) => unknown
}
  
export interface PersistOptions<S, PersistedState = S> {
  name: string
  storage?: PersistStorage<PersistedState> | undefined
  partialize?: (state: S) => PersistedState
  onRehydrateStorage?: (
    state: S,
  ) => ((state?: S, error?: unknown) => void) | void
  version?: number
  migrate?: (
    persistedState: unknown,
    version: number,
  ) => PersistedState | Promise<PersistedState>
  merge?: (persistedState: unknown, currentState: S) => S
  skipHydration?: boolean
}

type PersistListener<S> = (state: S) => void

type StorePersist<S, Ps> = {
  persist: {
    setOptions: (options: Partial<PersistOptions<S, Ps>>) => void
    clearStorage: () => void
    rehydrate: () => Promise<void> | void
    hasHydrated: () => boolean
    onHydrate: (fn: PersistListener<S>) => () => void
    onFinishHydration: (fn: PersistListener<S>) => () => void
    getOptions: () => Partial<PersistOptions<S, Ps>>
  }
}

type Thenable<Value> = {
  then<V>(
    onFulfilled: (value: Value) => V | Promise<V> | Thenable<V>,
  ): Thenable<V>
  catch<V>(
    onRejected: (reason: Error) => V | Promise<V> | Thenable<V>,
  ): Thenable<V>
}

StateStorage

export interface StateStorage {
  getItem: (name: string) => string | null | Promise<string | null>
  setItem: (name: string, value: string) => unknown | Promise<unknown>
  removeItem: (name: string) => unknown | Promise<unknown>
}
  • 상태를 저장하기 위한 기본 스토리지 인터페이스의 정의
  • 세 가지 기능으로 이루어져 있다.
    - getItem: 데이터를 불러오는 함수
    - setItem: 데이터를 저장하는 함수
    - removeItem: 데이터를 삭제하는 함수
  • 그리고 unknown | Promise<unknown> 동기와 비동기를 둘다 받을 수 있게 넣어서 동일한 방식으로 데이터를 사용할수 있다.

StorageValue

Copyexport type StorageValue<S> = {
  state: S
  version?: number
}
  • 저장소에 저장될 데이터의 구조를 정의
  • 실제로 저장될 상태 데이터(state)와 상태의 버전을 나타내는 선택적 숫자 값(version)으로 이루어져 있음
    버전 관리를 통해 데이터 구조가 변경될 때 적절한 마이그레이션을 수행가능

PersistStorage

export interface PersistStorage<S> {
  getItem: (name: string) => StorageValue<S> | null | Promise<StorageValue<S> | null>
  setItem: (name: string, value: StorageValue<S>) => unknown | Promise<unknown>
  removeItem: (name: string) => unknown | Promise<unknown>
}
  • StateStorage의 확장 버전이라고 보면 된다. 기존의 StateStorage가 value 값으로 string 만 받는다면, PersistStorage는 value 값으로 StorageValue<S> 타입을 받는다.
    즉, 단순한 문자열이 아닌, state와 version을 포함한 객체를 저장한다. 객체를 저장함으로서 복잡한 데이터 구조를 저장 할 수 있고, 버전을 관리 할 수 있다.

PersistOptions

export interface PersistOptions<S, PersistedState = S> {
  name: string             // 스토리지에서 사용할 고유한 이름
  storage?: PersistStorage<PersistedState> | undefined  // 커스텀 스토리지
  partialize?: (state: S) => PersistedState  // 저장할 상태를 필터링
  onRehydrateStorage?: (state: S) => ((state?: S, error?: unknown) => void) | void  // 상태 복원 콜백
  version?: number         // 상태의 버전
  migrate?: (persistedState: unknown, version: number) => PersistedState | Promise<PersistedState>  // 버전 마이그레이션
  merge?: (persistedState: unknown, currentState: S) => S  // 상태 병합 방식
  skipHydration?: boolean  // 초기 hydration 스킵 여부
}
  • 상태 지속성(persistence)을 위한 다양한 설정 옵션을 정의하는 인터페이스
  • 여덟 가지 주요 옵션으로 구성되어 있다
    - name: 스토리지 내에서 상태를 식별하기 위한 고유 이름
    - storage: 상태를 저장할 커스텀 스토리지 지정 (기본값은 localStorage)
    - partialize: 저장할 상태를 필터링하는 함수
    - onRehydrateStorage: 상태 복원 전후에 실행될 콜백 함수
    - version: 상태의 버전 관리를 위한 번호
    - migrate: 버전이 변경될 때 상태를 변환하는 마이그레이션 함수
    - merge: 저장된 상태와 현재 상태를 병합하는 방식을 정의하는 함수
    - skipHydration: 초기 상태 복원을 건너뛸지 결정하는 옵션

이를 통해 상태의 저장, 복원, 버전 관리, 마이그레이션 등 다양한 측면을 세밀하게 제어할 수 있다.

JsonStorageOptions

type JsonStorageOptions = {
  reviver?: (key: string, value: unknown) => unknown
  replacer?: (key: string, value: unknown) => unknown
}
  • JSON 직렬화/역직렬화 옵션을 정의
  • 하단의 createJSONStorage 에서 사용
    - reviver: JSON.parse의 두 번째 인자로 사용됨
    - replacer: JSON.stringify의 두 번째 인자로 사용됨

스토리지 생성

export function createJSONStorage<S>(
  getStorage: () => StateStorage,
  options?: JsonStorageOptions,
): PersistStorage<S> | undefined {
  let storage: StateStorage | undefined
  try {
    storage = getStorage()
  } catch {
    // prevent error if the storage is not defined (e.g. when server side rendering a page)
    return
  }
  const persistStorage: PersistStorage<S> = {
    getItem: (name) => {
      const parse = (str: string | null) => {
        if (str === null) {
          return null
        }
        return JSON.parse(str, options?.reviver) as StorageValue<S>
      }
      const str = (storage as StateStorage).getItem(name) ?? null
      if (str instanceof Promise) {
        return str.then(parse)
      }
      return parse(str)
    },
    setItem: (name, newValue) =>
      (storage as StateStorage).setItem(
        name,
        JSON.stringify(newValue, options?.replacer),
      ),
    removeItem: (name) => (storage as StateStorage).removeItem(name),
  }
  return persistStorage
}

이 함수는 local이나 session 같은 기본 웹 스토리지(StateStorage)를 생성하고, 이를 JSON 형식의 객체 형태의 데이터를 저장할 수 있는 저장소(PersistStorage)로 변환해주는 함수이다.

  1. 먼저 인자로 getStorage와 options를 받는다. 이 두개는 각각 아래와 같다.
    • getStorage : 스토리지를 반환하는 함수
    • options를 : JSON 변환 옵션
  2. getStorage를 통해서 스토리지를 가지고온다
  3. PersistStorage 객체 생성후 반환
    • 이 PersistStorage의 키로는 getItem, setItem, removeItem 세가지가 있는데, 각각 아래와 같은 동작을 한다.
      • getItem -> 문자열을 객체로 변환 후 반환
        • 이때 options.reviver를 사용하여 특정 값들을 원하는 형태(예: Date 객체)로 변환
      • setItem -> 객체를 문자열로 변환 후 반환
      • removeItem -> storage의 removeItem을 그대로 사용하여 해당 키의 데이터 삭제

유틸리티 헬퍼 함수

// Promise와 유사한 인터페이스 정의
type Thenable<Value> = {
  then<V>(
    onFulfilled: (value: Value) => V | Promise<V> | Thenable<V>,
  ): Thenable<V>
  catch<V>(
    onRejected: (reason: Error) => V | Promise<V> | Thenable<V>,
  ): Thenable<V>
}

// 일반 값이나 Promise를 Thenable로 변환하는 함수
const toThenable =
  <Result, Input>(
    fn: (input: Input) => Result | Promise<Result> | Thenable<Result>,
  ) =>
  (input: Input): Thenable<Result> => {
    try {
      const result = fn(input)
      // 이미 Promise면 그대로 반환
      if (result instanceof Promise) {
        return result as Thenable<Result>
      }
      // 일반 값이면 Thenable 객체로 변환
      return {
        then(onFulfilled) {
          return toThenable(onFulfilled)(result as Result)
        },
        catch(_onRejected) {
          return this as Thenable<any>
        },
      }
    } catch (e: any) {
      // 에러 발생시 catch로 처리할 수 있는 Thenable 반환
      return {
        then(_onFulfilled) {
          return this as Thenable<any>
        },
        catch(onRejected) {
          return toThenable(onRejected)(e)
        },
      }
    }
  }
  • 일반 값이나 Promise를 Promise와 유사한 인터페이스(Thenable)로 통일하는 유틸리티 함수이다.
  • 동기/비동기 작업을 일관되게 처리할 수 있게 한다.

핵심 로직

const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
  type S = ReturnType<typeof config>
  let options = {
    storage: createJSONStorage<S>(() => localStorage),
    partialize: (state: S) => state,
    version: 0,
    merge: (persistedState: unknown, currentState: S) => ({
      ...currentState,
      ...(persistedState as object),
    }),
    ...baseOptions,
  }

  let hasHydrated = false
  const hydrationListeners = new Set<PersistListener<S>>()
  const finishHydrationListeners = new Set<PersistListener<S>>()
  let storage = options.storage

  if (!storage) {
    return config(
      (...args) => {
        console.warn(
          `[zustand persist middleware] Unable to update item '${options.name}', the given storage is currently unavailable.`,
        )
        set(...(args as Parameters<typeof set>))
      },
      get,
      api,
    )
  }

  const setItem = () => {
    const state = options.partialize({ ...get() })
    return (storage as PersistStorage<S>).setItem(options.name, {
      state,
      version: options.version,
    })
  }

  const savedSetState = api.setState

  api.setState = (state, replace) => {
    savedSetState(state, replace as any)
    void setItem()
  }

  const configResult = config(
    (...args) => {
      set(...(args as Parameters<typeof set>))
      void setItem()
    },
    get,
    api,
  )

  api.getInitialState = () => configResult

  // a workaround to solve the issue of not storing rehydrated state in sync storage
  // the set(state) value would be later overridden with initial state by create()
  // to avoid this, we merge the state from localStorage into the initial state.
  let stateFromStorage: S | undefined

  // rehydrate initial state with existing stored state
  const hydrate = () => {
    if (!storage) return

    // On the first invocation of 'hydrate', state will not yet be defined (this is
    // true for both the 'asynchronous' and 'synchronous' case). Pass 'configResult'
    // as a backup  to 'get()' so listeners and 'onRehydrateStorage' are called with
    // the latest available state.

    hasHydrated = false
    hydrationListeners.forEach((cb) => cb(get() ?? configResult))

    const postRehydrationCallback =
      options.onRehydrateStorage?.(get() ?? configResult) || undefined

    // bind is used to avoid `TypeError: Illegal invocation` error
    return toThenable(storage.getItem.bind(storage))(options.name)
      .then((deserializedStorageValue) => {
        if (deserializedStorageValue) {
          if (
            typeof deserializedStorageValue.version === 'number' &&
            deserializedStorageValue.version !== options.version
          ) {
            if (options.migrate) {
              return [
                true,
                options.migrate(
                  deserializedStorageValue.state,
                  deserializedStorageValue.version,
                ),
              ] as const
            }
            console.error(
              `State loaded from storage couldn't be migrated since no migrate function was provided`,
            )
          } else {
            return [false, deserializedStorageValue.state] as const
          }
        }
        return [false, undefined] as const
      })
      .then((migrationResult) => {
        const [migrated, migratedState] = migrationResult
        stateFromStorage = options.merge(
          migratedState as S,
          get() ?? configResult,
        )

        set(stateFromStorage as S, true)
        if (migrated) {
          return setItem()
        }
      })
      .then(() => {
        // TODO: In the asynchronous case, it's possible that the state has changed
        // since it was set in the prior callback. As such, it would be better to
        // pass 'get()' to the 'postRehydrationCallback' to ensure the most up-to-date
        // state is used. However, this could be a breaking change, so this isn't being
        // done now.
        postRehydrationCallback?.(stateFromStorage, undefined)

        // It's possible that 'postRehydrationCallback' updated the state. To ensure
        // that isn't overwritten when returning 'stateFromStorage' below
        // (synchronous-case only), update 'stateFromStorage' to point to the latest
        // state. In the asynchronous case, 'stateFromStorage' isn't used after this
        // callback, so there's no harm in updating it to match the latest state.
        stateFromStorage = get()
        hasHydrated = true
        finishHydrationListeners.forEach((cb) => cb(stateFromStorage as S))
      })
      .catch((e: Error) => {
        postRehydrationCallback?.(undefined, e)
      })
  }
  
    ;(api as StoreApi<S> & StorePersist<S, S>).persist = {
    setOptions: (newOptions) => {
      options = {
        ...options,
        ...newOptions,
      }

      if (newOptions.storage) {
        storage = newOptions.storage
      }
    },
    clearStorage: () => {
      storage?.removeItem(options.name)
    },
    getOptions: () => options,
    rehydrate: () => hydrate() as Promise<void>,
    hasHydrated: () => hasHydrated,
    onHydrate: (cb) => {
      hydrationListeners.add(cb)

      return () => {
        hydrationListeners.delete(cb)
      }
    },
    onFinishHydration: (cb) => {
      finishHydrationListeners.add(cb)

      return () => {
        finishHydrationListeners.delete(cb)
      }
    },
  }

  if (!options.skipHydration) {
    hydrate()
  }

  return stateFromStorage || configResult
}

이부분이 presist의 핵심 로직이다.
간단하게만 이야기 하면 기본 옵션을 설정하고, 저장소 존재 여부를 확인한 후, 저장된 상태와 현재 상태를 불러와 버전을 비교하여 버전이 다르면 마이그레이션을 수행하고, 저장된 상태와 현재 상태를 병합해서 최종 상태를 만듦으로써 zustand의 상태를 영구 저장소와 동기화하는 함수이다.

이제 천천히 알아보자

초기설정

const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
  type S = ReturnType<typeof config>
  let options = {
    storage: createJSONStorage<S>(() => localStorage), // 기본 저장소
    partialize: (state: S) => state,                  // 상태 필터링
    version: 0,                                       // 버전
    merge: (persistedState: unknown, currentState: S) => ({
      ...currentState,
      ...(persistedState as object),
    }),
    ...baseOptions,  // 사용자 옵션으로 덮어쓰기
  }
  • config와 baseOptions를 받아서 zustand의 기본 API(set, get, api)를 받는 함수 반환
  • options으로 기본값 설정 (저장소, 필터링, 버전, 병합 방식)

상태관리 변수 설정

let hasHydrated = false
const hydrationListeners = new Set<PersistListener<S>>()
const finishHydrationListeners = new Set<PersistListener<S>>()
let storage = options.storage
  • 상태 복원 관련 플래그와 리스너들 초기화

저장소가 존재하지 않을시

if (!storage) {
  return config(
    (...args) => {
      console.warn(...)
      set(...(args as Parameters<typeof set>))
    },
    get,
    api,
  )
}
  • 저장소가 존재하지 않을시 경고 출력후 상태 관리만 수행

상태 저장 함수 정의

const setItem = () => {
  const state = options.partialize({ ...get() })
  return storage.setItem(options.name, {
    state,
    version: options.version,
  })
}
  • 현재 상태를 필터링 해서 저장소에 저장

setState 재정의

const savedSetState = api.setState
api.setState = (state, replace) => {
  savedSetState(state, replace as any)
  void setItem()
}
  • 상태 변경기 자동으로 저장소에도 저장되도록 함

초기 상태 설정:

const configResult = config(
  (...args) => {
    set(...(args as Parameters<typeof set>))
    void setItem()
  },
  get,
  api,
)
api.getInitialState = () => configResult
  • 초기 상태를 설정하고 저장소에 저장

hydrate 함수

const hydrate = () => {
  // 저장소가 없으면 아무것도 반환하지 않음
  if (!storage) return
  
  // 상태 복원 시작을 표시하고 리스너들에게 알림
  hasHydrated = false
  hydrationListeners.forEach((cb) => cb(get() ?? configResult))

  // 상태 복원 전 콜백 설정
  const postRehydrationCallback = 
    options.onRehydrateStorage?.(get() ?? configResult) || undefined
  
  // storage.getItem을 Promise로 변환해서 실행
  return toThenable(storage.getItem.bind(storage))(options.name)
    .then(버전 체크 및 마이그레이션함수)
    .then(상태 병합 및 적용)
    .then(복원 완료 처리)
    .catch(에러 처리)
}
  • 저장된 상태를 불러와 현재 상태와 병합
  • 버전이 다르면 마이그레이션 수행
  • 복원 과정의 시작과 완료를 리스너들에게 알림

persist API 및 설정

(api as StoreApi<S> & StorePersist<S, S>).persist = {
    // 옵션 설정/변경
    setOptions: (newOptions) => {
      options = { ...options, ...newOptions }
      if (newOptions.storage) storage = newOptions.storage
    },
    // 저장소 초기화
    clearStorage: () => storage?.removeItem(options.name),
    // 현재 옵션 가져오기
    getOptions: () => options,
    // 상태 복원 실행
    rehydrate: () => hydrate() as Promise<void>,
    // 복원 완료 여부 확인
    hasHydrated: () => hasHydrated,
    // 복원 시작 리스너 추가/제거
    onHydrate: (cb) => {
      hydrationListeners.add(cb)
      return () => hydrationListeners.delete(cb)
    },
    // 복원 완료 리스너 추가/제거
    onFinishHydration: (cb) => {
      finishHydrationListeners.add(cb)
      return () => finishHydrationListeners.delete(cb)
    },
}

// skipHydration 옵션이 false면 hydrate 실행
if (!options.skipHydration) {
  hydrate()
}

// 저장된 상태나 초기 상태 반환
return stateFromStorage || configResult
  • persist API와 타입 정의를 설정하고 기능 구현
  • skipHydration 옵션에 따라 hydrate 함수 실행 여부 결정
  • 최종적으로 저장된 상태나 초기 상태를 반환

전체 코드

import type {
  StateCreator,
  StoreApi,
  StoreMutatorIdentifier,
} from '../vanilla.ts'

export interface StateStorage {
  getItem: (name: string) => string | null | Promise<string | null>
  setItem: (name: string, value: string) => unknown | Promise<unknown>
  removeItem: (name: string) => unknown | Promise<unknown>
}

export type StorageValue<S> = {
  state: S
  version?: number
}

export interface PersistStorage<S> {
  getItem: (
    name: string,
  ) => StorageValue<S> | null | Promise<StorageValue<S> | null>
  setItem: (name: string, value: StorageValue<S>) => unknown | Promise<unknown>
  removeItem: (name: string) => unknown | Promise<unknown>
}

type JsonStorageOptions = {
  reviver?: (key: string, value: unknown) => unknown
  replacer?: (key: string, value: unknown) => unknown
}

export function createJSONStorage<S>(
  getStorage: () => StateStorage,
  options?: JsonStorageOptions,
): PersistStorage<S> | undefined {
  let storage: StateStorage | undefined
  try {
    storage = getStorage()
  } catch {
    // prevent error if the storage is not defined (e.g. when server side rendering a page)
    return
  }
  const persistStorage: PersistStorage<S> = {
    getItem: (name) => {
      const parse = (str: string | null) => {
        if (str === null) {
          return null
        }
        return JSON.parse(str, options?.reviver) as StorageValue<S>
      }
      const str = (storage as StateStorage).getItem(name) ?? null
      if (str instanceof Promise) {
        return str.then(parse)
      }
      return parse(str)
    },
    setItem: (name, newValue) =>
      (storage as StateStorage).setItem(
        name,
        JSON.stringify(newValue, options?.replacer),
      ),
    removeItem: (name) => (storage as StateStorage).removeItem(name),
  }
  return persistStorage
}

export interface PersistOptions<S, PersistedState = S> {
  /** Name of the storage (must be unique) */
  name: string
  /**
   * Use a custom persist storage.
   *
   * Combining `createJSONStorage` helps creating a persist storage
   * with JSON.parse and JSON.stringify.
   *
   * @default createJSONStorage(() => localStorage)
   */
  storage?: PersistStorage<PersistedState> | undefined
  /**
   * Filter the persisted value.
   *
   * @params state The state's value
   */
  partialize?: (state: S) => PersistedState
  /**
   * A function returning another (optional) function.
   * The main function will be called before the state rehydration.
   * The returned function will be called after the state rehydration or when an error occurred.
   */
  onRehydrateStorage?: (
    state: S,
  ) => ((state?: S, error?: unknown) => void) | void
  /**
   * If the stored state's version mismatch the one specified here, the storage will not be used.
   * This is useful when adding a breaking change to your store.
   */
  version?: number
  /**
   * A function to perform persisted state migration.
   * This function will be called when persisted state versions mismatch with the one specified here.
   */
  migrate?: (
    persistedState: unknown,
    version: number,
  ) => PersistedState | Promise<PersistedState>
  /**
   * A function to perform custom hydration merges when combining the stored state with the current one.
   * By default, this function does a shallow merge.
   */
  merge?: (persistedState: unknown, currentState: S) => S

  /**
   * An optional boolean that will prevent the persist middleware from triggering hydration on initialization,
   * This allows you to call `rehydrate()` at a specific point in your apps rendering life-cycle.
   *
   * This is useful in SSR application.
   *
   * @default false
   */
  skipHydration?: boolean
}

type PersistListener<S> = (state: S) => void

type StorePersist<S, Ps> = {
  persist: {
    setOptions: (options: Partial<PersistOptions<S, Ps>>) => void
    clearStorage: () => void
    rehydrate: () => Promise<void> | void
    hasHydrated: () => boolean
    onHydrate: (fn: PersistListener<S>) => () => void
    onFinishHydration: (fn: PersistListener<S>) => () => void
    getOptions: () => Partial<PersistOptions<S, Ps>>
  }
}

type Thenable<Value> = {
  then<V>(
    onFulfilled: (value: Value) => V | Promise<V> | Thenable<V>,
  ): Thenable<V>
  catch<V>(
    onRejected: (reason: Error) => V | Promise<V> | Thenable<V>,
  ): Thenable<V>
}

const toThenable =
  <Result, Input>(
    fn: (input: Input) => Result | Promise<Result> | Thenable<Result>,
  ) =>
  (input: Input): Thenable<Result> => {
    try {
      const result = fn(input)
      if (result instanceof Promise) {
        return result as Thenable<Result>
      }
      return {
        then(onFulfilled) {
          return toThenable(onFulfilled)(result as Result)
        },
        catch(_onRejected) {
          return this as Thenable<any>
        },
      }
    } catch (e: any) {
      return {
        then(_onFulfilled) {
          return this as Thenable<any>
        },
        catch(onRejected) {
          return toThenable(onRejected)(e)
        },
      }
    }
  }

const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
  type S = ReturnType<typeof config>
  let options = {
    storage: createJSONStorage<S>(() => localStorage),
    partialize: (state: S) => state,
    version: 0,
    merge: (persistedState: unknown, currentState: S) => ({
      ...currentState,
      ...(persistedState as object),
    }),
    ...baseOptions,
  }

  let hasHydrated = false
  const hydrationListeners = new Set<PersistListener<S>>()
  const finishHydrationListeners = new Set<PersistListener<S>>()
  let storage = options.storage

  if (!storage) {
    return config(
      (...args) => {
        console.warn(
          `[zustand persist middleware] Unable to update item '${options.name}', the given storage is currently unavailable.`,
        )
        set(...(args as Parameters<typeof set>))
      },
      get,
      api,
    )
  }

  const setItem = () => {
    const state = options.partialize({ ...get() })
    return (storage as PersistStorage<S>).setItem(options.name, {
      state,
      version: options.version,
    })
  }

  const savedSetState = api.setState

  api.setState = (state, replace) => {
    savedSetState(state, replace as any)
    void setItem()
  }

  const configResult = config(
    (...args) => {
      set(...(args as Parameters<typeof set>))
      void setItem()
    },
    get,
    api,
  )

  api.getInitialState = () => configResult

  // a workaround to solve the issue of not storing rehydrated state in sync storage
  // the set(state) value would be later overridden with initial state by create()
  // to avoid this, we merge the state from localStorage into the initial state.
  let stateFromStorage: S | undefined

  // rehydrate initial state with existing stored state
  const hydrate = () => {
    if (!storage) return

    // On the first invocation of 'hydrate', state will not yet be defined (this is
    // true for both the 'asynchronous' and 'synchronous' case). Pass 'configResult'
    // as a backup  to 'get()' so listeners and 'onRehydrateStorage' are called with
    // the latest available state.

    hasHydrated = false
    hydrationListeners.forEach((cb) => cb(get() ?? configResult))

    const postRehydrationCallback =
      options.onRehydrateStorage?.(get() ?? configResult) || undefined

    // bind is used to avoid `TypeError: Illegal invocation` error
    return toThenable(storage.getItem.bind(storage))(options.name)
      .then((deserializedStorageValue) => {
        if (deserializedStorageValue) {
          if (
            typeof deserializedStorageValue.version === 'number' &&
            deserializedStorageValue.version !== options.version
          ) {
            if (options.migrate) {
              return [
                true,
                options.migrate(
                  deserializedStorageValue.state,
                  deserializedStorageValue.version,
                ),
              ] as const
            }
            console.error(
              `State loaded from storage couldn't be migrated since no migrate function was provided`,
            )
          } else {
            return [false, deserializedStorageValue.state] as const
          }
        }
        return [false, undefined] as const
      })
      .then((migrationResult) => {
        const [migrated, migratedState] = migrationResult
        stateFromStorage = options.merge(
          migratedState as S,
          get() ?? configResult,
        )

        set(stateFromStorage as S, true)
        if (migrated) {
          return setItem()
        }
      })
      .then(() => {
        // TODO: In the asynchronous case, it's possible that the state has changed
        // since it was set in the prior callback. As such, it would be better to
        // pass 'get()' to the 'postRehydrationCallback' to ensure the most up-to-date
        // state is used. However, this could be a breaking change, so this isn't being
        // done now.
        postRehydrationCallback?.(stateFromStorage, undefined)

        // It's possible that 'postRehydrationCallback' updated the state. To ensure
        // that isn't overwritten when returning 'stateFromStorage' below
        // (synchronous-case only), update 'stateFromStorage' to point to the latest
        // state. In the asynchronous case, 'stateFromStorage' isn't used after this
        // callback, so there's no harm in updating it to match the latest state.
        stateFromStorage = get()
        hasHydrated = true
        finishHydrationListeners.forEach((cb) => cb(stateFromStorage as S))
      })
      .catch((e: Error) => {
        postRehydrationCallback?.(undefined, e)
      })
  }

  ;(api as StoreApi<S> & StorePersist<S, S>).persist = {
    setOptions: (newOptions) => {
      options = {
        ...options,
        ...newOptions,
      }

      if (newOptions.storage) {
        storage = newOptions.storage
      }
    },
    clearStorage: () => {
      storage?.removeItem(options.name)
    },
    getOptions: () => options,
    rehydrate: () => hydrate() as Promise<void>,
    hasHydrated: () => hasHydrated,
    onHydrate: (cb) => {
      hydrationListeners.add(cb)

      return () => {
        hydrationListeners.delete(cb)
      }
    },
    onFinishHydration: (cb) => {
      finishHydrationListeners.add(cb)

      return () => {
        finishHydrationListeners.delete(cb)
      }
    },
  }

  if (!options.skipHydration) {
    hydrate()
  }

  return stateFromStorage || configResult
}

type Persist = <
  T,
  Mps extends [StoreMutatorIdentifier, unknown][] = [],
  Mcs extends [StoreMutatorIdentifier, unknown][] = [],
  U = T,
>(
  initializer: StateCreator<T, [...Mps, ['zustand/persist', unknown]], Mcs>,
  options: PersistOptions<T, U>,
) => StateCreator<T, Mps, [['zustand/persist', U], ...Mcs]>

declare module '../vanilla' {
  interface StoreMutators<S, A> {
    'zustand/persist': WithPersist<S, A>
  }
}

type Write<T, U> = Omit<T, keyof U> & U

type WithPersist<S, A> = S extends { getState: () => infer T }
  ? Write<S, StorePersist<T, A>>
  : never

type PersistImpl = <T>(
  storeInitializer: StateCreator<T, [], []>,
  options: PersistOptions<T, T>,
) => StateCreator<T, [], []>

export const persist = persistImpl as unknown as Persist

0개의 댓글