Infer Immer middleware type for setState from setState instead of getState

[chrisvander]

The Immer middleware infers the type nextStateOrUpdater for setState from getState:

type StoreImmer<S> = S extends {
    getState: () => infer T;
    setState: infer SetState;
} ? SetState extends (...a: infer A) => infer Sr ? {
    setState(nextStateOrUpdater: T | Partial<T> | ((state: Draft<T>) => void), shouldReplace?: boolean | undefined, ...a: SkipTwo<A>): Sr;
} : never : never;

This is typically not a problem unless you're using a middleware like chrisvander/zustand-computed, where the state returned by getState differs from what is allowed to be set via setState. In the case of the computed middleware, getState returns additional keys beyond what the user is allowed to modify. Can we change the Immer state to infer off of the SetState function? E.g. use something like Parameters<SetState>[0].

[dai-shi]

Parameters<SetState>[0] sounds fairly reasonable, but then we can't support devtools action name with immer middleware, can we?

[chrisvander]

The middleware declaration uses a second generic parameter, A. Does that mess anything up?

type StoreCompute<S, A> = S extends {
  getState: () => infer T
  setState: infer SetState
}
  ? Write<StoreApi<T & A>, { setState: SetState; getState: () => T & A }>
  : never

type WithCompute<S, A> = Write<S, StoreCompute<S, A>>

declare module "zustand" {
  interface StoreMutators<S, A> {
    "chrisvander/zustand-computed": WithCompute<S, A>
  }
}

[dai-shi]

https://tsplay.dev/mL3OaN
It doesn't seem only an issue with immer middleware. All other middleware don't seem to work.
Can you try something simple and see if it works as expected?

[chrisvander]

So the setState being inferred from Immer is still a problem, but you're right that the generics aren't being passed down appropriately.

Try this playground.

On line 116 you'll notice that set contains the getState type with the computed store, which I'm trying to avoid. Regardless, I manually typed the Immer middleware for useStoreImmer and useStoreImmer2 and that mitigated the problem you highlighted above. It seems that the Mos type is not being passed down.

[dai-shi]

https://tsplay.dev/WGgq9W
Now, I know why. ComputedStore has to be inferred before passing to the next middleware.

[chrisvander]

Beautiful. Thanks. I suppose I could use currying, either way I see now that the computed state needs to be inferred first.