Repository: pmndrs/zustand
File: src/middleware.ts (TypeScript)

Code before:
        // this is a fake function just to avoid catch is not a function error
        catch: () => undefined,
      }
    }
  }
  // rehydrate initial state with existing stored state
  ;(() => {
    const postRehydrationCallback = onRehydrateStorage?.(get()) || undefined

    let storageValuePending = makeThenable(storage.getItem(name))
    storageValuePending
      .then((storageValue) => {
        if (storageValue) {
          return makeThenable(deserialize(storageValue))
        } else {
          return makeThenable(null)
        }
      })
      .then((deserializedStorageValue) => {
        if (deserializedStorageValue) {
          if (deserializedStorageValue.version !== version) {
            return makeThenable(
              migrate?.(
                deserializedStorageValue.state,
                deserializedStorageValue.version
              )
            )
          } else {
            set(deserializedStorageValue.state)
            return makeThenable(null)
          }
        } else {
          return makeThenable(null)
        }
      })
      .then((migratedState) => {
        if (migratedState) {
          set(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)
      })

Review comment (suggestion):
I think this needs to be
```suggestion
            
            return makeThenable(
              migrate?.(
                deserializedStorageValue.state,
                deserializedStorageValue.version
              )
              .then((migratedState) => makeThenable(set(migratedState)))
              .then(() => makeThenable(setItem()))
```
-- Edit: Added `makeThenable` calls.

Code after:
          `Persist middleware: unable to update ${name}, the given storage is currently unavailable.`
        )
        set(...args)
      },
      get,
      api
    )
  }

  const thenableSerialize = toThenable(serialize)

  const setItem = (): Thenable<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])
    }

    let errorInSync: Error | undefined
    const thenable = thenableSerialize({ state, version })
      .then((serializedValue) =>
        (storage as StateStorage).setItem(name, serializedValue)
      )
      .catch((e) => {
        errorInSync = e
      })
    if (errorInSync) {
      throw errorInSync
    }
    return thenable
  }

  const savedSetState = api.setState

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

  // rehydrate initial state with existing stored state

  // 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 stateFromStorageInSync: S | undefined
  const postRehydrationCallback = onRehydrateStorage?.(get()) || undefined