Fix for sync rehydrating zustand store

middleware.ts

import {
  GetState,
  PartialState,
  SetState,
  State,
  StateCreator,
  StoreApi,
} from './vanilla'

export const redux = <S extends State, A extends { type: unknown }>(
  reducer: (state: S, action: A) => S,
  initial: S
) => (
  set: SetState<S>,
  get: GetState<S>,
  api: StoreApi<S> & {
    dispatch?: (a: A) => A
    devtools?: any
  }
): S & { dispatch: (a: A) => A } => {
  api.dispatch = (action: A) => {
    set((state: S) => reducer(state, action))
    if (api.devtools) {
      api.devtools.send(api.devtools.prefix + action.type, get())
    }
    return action
  }
  return { dispatch: api.dispatch, ...initial }
}

type NamedSet<T extends State> = {
  <K extends keyof T>(
    partial: PartialState<T, K>,
    replace?: boolean,
    name?: string
  ): void
}

export const devtools = <S extends State>(
  fn: (set: NamedSet<S>, get: GetState<S>, api: StoreApi<S>) => S,
  prefix?: string
) => (
  set: SetState<S>,
  get: GetState<S>,
  api: StoreApi<S> & { dispatch?: unknown; devtools?: any }
): S => {
  let extension
  try {
    extension =
      (window as any).__REDUX_DEVTOOLS_EXTENSION__ ||
      (window as any).top.__REDUX_DEVTOOLS_EXTENSION__
  } catch {}

  if (!extension) {
    if (
      process.env.NODE_ENV === 'development' &&
      typeof window !== 'undefined'
    ) {
      console.warn('Please install/enable Redux devtools extension')
    }
    api.devtools = null
    return fn(set, get, api)
  }
  const namedSet: NamedSet<S> = (state, replace, name) => {
    set(state, replace)
    if (!api.dispatch) {
      api.devtools.send(api.devtools.prefix + (name || 'action'), get())
    }
  }
  const initialState = fn(namedSet, get, api)
  if (!api.devtools) {
    const savedSetState = api.setState
    api.setState = <K extends keyof S>(
      state: PartialState<S, K>,
      replace?: boolean
    ) => {
      savedSetState(state, replace)
      api.devtools.send(api.devtools.prefix + 'setState', api.getState())
    }
    api.devtools = extension.connect({ name: prefix })
    api.devtools.prefix = prefix ? `${prefix} > ` : ''
    api.devtools.subscribe((message: any) => {
      if (message.type === 'DISPATCH' && message.state) {
        const ignoreState =
          message.payload.type === 'JUMP_TO_ACTION' ||
          message.payload.type === 'JUMP_TO_STATE'
        if (!api.dispatch && !ignoreState) {
          api.setState(JSON.parse(message.state))
        } else {
          savedSetState(JSON.parse(message.state))
        }
      } else if (
        message.type === 'DISPATCH' &&
        message.payload?.type === 'COMMIT'
      ) {
        api.devtools.init(api.getState())
      } else if (
        message.type === 'DISPATCH' &&
        message.payload?.type === 'IMPORT_STATE'
      ) {
        const actions = message.payload.nextLiftedState?.actionsById
        const computedStates =
          message.payload.nextLiftedState?.computedStates || []

        computedStates.forEach(
          ({ state }: { state: PartialState<S> }, index: number) => {
            const action = actions[index] || api.devtools.prefix + 'setState'

            if (index === 0) {
              api.devtools.init(state)
            } else {
              savedSetState(state)
              api.devtools.send(action, api.getState())
            }
          }
        )
      }
    })
    api.devtools.init(initialState)
  }
  return initialState
}

type Combine<T, U> = Omit<T, keyof U> & U
export const combine = <
  PrimaryState extends State,
  SecondaryState extends State
>(
  initialState: PrimaryState,
  create: (
    set: SetState<PrimaryState>,
    get: GetState<PrimaryState>,
    api: StoreApi<PrimaryState>
  ) => SecondaryState
): StateCreator<Combine<PrimaryState, SecondaryState>> => (set, get, api) =>
  Object.assign(
    {},
    initialState,
    create(
      (set as unknown) as SetState<PrimaryState>,
      (get as unknown) as GetState<PrimaryState>,
      (api as unknown) as StoreApi<PrimaryState>
    )
  )

type StateStorage = {
  getItem: (name: string) => string | null | Promise<string | null>
  setItem: (name: string, value: string) => void | Promise<void>
}
type StorageValue<S> = { state: S; version: number }
type PersistOptions<S> = {
  /** Name of the storage (must be unique) */
  name: string
  /**
   * A function returning a storage.
   * The storage must fit `window.localStorage`'s api (or an async version of it).
   * For example the storage could be `AsyncStorage` from React Native.
   *
   * @default () => localStorage
   */
  getStorage?: () => StateStorage
  /**
   * Use a custom serializer.
   * The returned string will be stored in the storage.
   *
   * @default JSON.stringify
   */
  serialize?: (state: StorageValue<S>) => string | Promise<string>
  /**
   * Use a custom deserializer.
   *
   * @param str The storage's current value.
   * @default JSON.parse
   */
  deserialize?: (str: string) => StorageValue<S> | Promise<StorageValue<S>>
  /**
   * Prevent some items from being stored.
   */
  blacklist?: (keyof S)[]
  /**
   * Only store the listed properties.
   */
  whitelist?: (keyof S)[]
  /**
   * 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?: Error) => 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: any, version: number) => S | Promise<S>
}

export const persist = <S extends State>(
  config: StateCreator<S>,
  options: PersistOptions<S>
) => (set: SetState<S>, get: GetState<S>, api: StoreApi<S>): S => {
  console.log('middleware.ts: start persist()')
  const {
    name,
    getStorage = () => localStorage,
    serialize = JSON.stringify,
    deserialize = JSON.parse,
    blacklist,
    whitelist,
    onRehydrateStorage,
    version = 0,
    migrate,
  } = options || {}

  let storage: StateStorage | undefined

  try {
    storage = getStorage()
  } catch (e) {
    // prevent error if the storage is not defined (e.g. when server side rendering a page)
  }

  if (!storage) {
    return config(
      (...args) => {
        console.warn(
          `Persist middleware: unable to update ${name}, the given storage is currently unavailable.`
        )
        set(...args)
      },
      get,
      api
    )
  }

  const setItem = (): Promise<void> | void => {
    const state = { ...get() }

    if (whitelist) {
      ;(Object.keys(state) as (keyof S)[]).forEach((key) => {
        !whitelist.includes(key) && delete state[key]
      })
    }
    if (blacklist) {
      blacklist.forEach((key) => delete state[key])
    }

    if (storage) {
      const serializeResult = serialize({ state, version })
      return serializeResult instanceof Promise
        ? serializeResult.then((value) => storage.setItem(name, value))
        : storage.setItem(name, serializeResult)
    }
  }

  const savedSetState = api.setState

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

  const makeThenable = (value: any) => {
    if (value instanceof Promise) {
      return value
    } else {
      return {
        then: (callback) => callback(value),
        // this is a fake function just to avoid catch is not a function error
        catch: () => undefined,
      }
    }
  }

  let stateFromStorage
    // rehydrate initial state with existing stored state
  ;(() => {
    const postRehydrationCallback = onRehydrateStorage?.(get()) || undefined

    try {
      makeThenable(storage.getItem(name))
        .then((storageValue) => {
          if (storageValue) {
            return makeThenable(deserialize(storageValue))
          } else {
            return makeThenable(null)
          }
        })
        .then((deserializedStorageValue) => {
          if (deserializedStorageValue) {
            if (deserializedStorageValue.version !== version) {
              stateFromStorage = migrate?.(
                deserializedStorageValue.state,
                deserializedStorageValue.version
              )

              set(stateFromStorage)
              return makeThenable(setItem())
            } else {
              stateFromStorage = deserializedStorageValue.state
              set(stateFromStorage)
              return makeThenable(null)
            }
          } else {
            return makeThenable(null)
          }
        })
        .then((migratedState) => {
          if (migratedState) {
            return makeThenable(setItem())
          } else {
            return makeThenable(null)
          }
        })
        .then(() => {
          postRehydrationCallback?.(get(), undefined)
          return makeThenable(null)
        })
        // this handler will only catch promise related errors
        .catch((e) => {
          postRehydrationCallback?.(undefined, e)
        })
    } catch (e) {
      postRehydrationCallback?.(undefined, e)
    }
  })()

  const configRes = config(
    (...args) => {
      set(...args)
      void setItem()
    },
    get,
    api
  )
  const mergedRes = { ...configRes, ...(stateFromStorage || {}) }
  return mergedRes
}

The relevant changes are in line 311 and 346-348.

Side note: I would have thought that the set() call in line 312 is no longer necessary, but without it it doesn't work at all. No idea why.

the then chains could be async. so setting stateFromStorage value in line 311 may occur after you returned mergedRes.

I guess that wouldn't be an issue. The return in line 348 would have been run by then, and since the set in line 312 would then run after the state was initialised it would work correctly (as it does now for the async) case.

But yeah, I'm not too sure it is the right approach. Maybe @dai-shi knows more?

Update:

• Add error handling for sync mode
• Fix migrations
• Remove console.logs

Hm, this might work. I'm not so sure. We'd need a PR to add more tests in advance, especially around error handling.

My idea was to make Thenable handle errors for both sync and async: pmndrs/zustand#403 (comment)
This would eliminate the duplication of L325, L328, which might be trivial though.

@dai-shi https://gist.github.com/barbogast/b6bf86f4bacb1da3e42e1f428c805eae contains tests for using the persist middleware in sync mode. I basically copied the existing tests and adapted them. Probably makes sense to add the file to the PR, to make sure it works before merging.

Are there any test cases missing in general?

Would you open a new PR for just adding the tests?

Sure, will do. The tests will fail, though, because #403 is not merged. Gotta go now, will do it later today.

oh, sorry, my misunderstanding. we are adding a sync support newly. what do we do? create a dedicated branch to work collaboratively?

Would work for me. @TheEhsanSarshar What do you think?

Hi everyone.
@barbogast yes please add tests and make a PR on the forked repo.