breaking(types): Add higher kinded mutator types (#725)

* imaginary code that uses uSES * revert backward compatibility code as this is not going to be v4 * use use-sync-external-store * revert to react 17 * handle error by our own * v4.0.0-alpha.2 * fix&refactor a bit * update uSES experimental package * remove error propagation hack * update size snapshot * update uSES and add dts * split react.ts and no export wild * split useStore impl * context to follow the new api, export wild again * v4.0.0-alpha.3 * add missing await * update uSES * update uSES * uses uSES extra! * v4.0.0-alpha.3 * fix(types): Rename from UseStore to UseBoundStore * breaking(types): drop deprecated UseStore type * breaking(core): drop v2 hook compatibility * breaking(middleware): drop deprecated persist options * breaking(core): drop deprecated store.subscribe with selector * update uSES * fix update uSES * v4.0.0-alpha.5 * combine subscribe type * intentional undefined type * add useDebugValue * update uSES * update uSES types * breaking(middleware): make persist options.removeItem required * update uSES * v4.0.0-alpha.6 * fix(readme): remove memoization section which is no longer valid with uSES * feat(readme): add new createStore/useStore usage * update useSES * update uSES and deps * v4.0.0-alpha.7 * update uSES * update uSES * shave bytes * vanilla: add higher kinded mutator types * persist: add higher kinded mutator types * persist: try to minimize diff * use `PopArgument` in vanilla too * update uSES * use overloads instead of `createWithState` * avoid symbols * add new types to middlewares * add new types react * add new types to context * fix persist types * add immer * migrate middleware type tests * fix react type, export `UseBoundStore` * migrate vanilla type tests * rename `_createStore` to `createStoreImpl` * Default to no mutations in `StateCreator` * migrate context.test.tsx * fix devtools.test.tsx type erros * context: remove callsignature in useStoreApi * context: remove `UseContextStore` type * context: fix useBoundStore type * context: keep `UseContextStore` for tooltip just don't export it * react: remove duplicate overload in create * export `WithPersist` * devtools: preserve try/catch * devtools: preserve window check * add a test case for v3 style create * devtools: preverse test fix from base branch * remove StoreApiWithFoo types, don't export WithFoo types * style * devtools: preverse `originalIsRecording` change * fix bug in devtools * empty commit * 4.0.0-beta.1 * fix lint * style * export immer fix tests * style, minor fixes * devtools: fix test * Update tests/devtools.test.tsx * breaking(middleware/devtools): use official devtools extension types * type object.create * avoid emitting @redux-devtools/extension * fix type with any * refactor * fix yarn lock * temporary fix #829 * v4.0.0-beta.2 * fix lint * lock date-fns version * test middleware subtyping * fix errors in conflict resolution * lock testing-library/react alpha version * more correct (and strict) persist types * migrate tests * wip release notes * fix devtools merge with base Co-authored-by: Daishi Kato <dai-shi@users.noreply.github.com> * add a test case for persist with partialize option * update readme * fix lint * immer: mutate `store.setState` * fix devtools merge with base * immer: fix mutations order * changes in readme * move and rename v4 migration md * add `combine` usage in readme * typos * create separate md for typescript, add common recipes * minor fixes * devtools: minor type fix (probably I copy pasted from persist and forgot to remove `U`) * add more migrations * context: fix import Co-authored-by: Daishi Kato <dai-shi@users.noreply.github.com> * devtools: minor refactor rearrange code for better diff * fix lint: run prettier * wip * uSES rc.1 * getServerState for #886, no types yet * uSES v1 * devtools: remove deprecations and warnings * fix tests * v4.0.0-beta.2 * wip * migrate tests * persist: keep diff minimal * fix merge in package.json and yarn.lock * fix merge for persist * don't use `import type` * docs(typescript): add slices pattern * fix selector & equals types for inference with useCallback, see issue #812 * add test for setState with replace * remove undefined selector overload * make immer more generic * make devtools replace-friendly and more correctly typed * migrate tests * make setState bivariant to make the state covariant * devtools: return the result of `setState` * devtools: make the fallback branch in `StoreSetStateWithAction` bivariant too (I forgot to make this bivaraint) * remove strict replace * fix lint Co-authored-by: daishi <daishi@axlight.com> Co-authored-by: Daishi Kato <dai-shi@users.noreply.github.com>
2025-12-08 19:45:52 +00:00 · 2022-04-18 05:44:00 +05:30 · 2022-04-18 05:44:00 +05:30 · 99cddcf76f
commit 99cddcf76f
parent 76eeacb1c4
19 changed files with 1449 additions and 941 deletions
--- a/docs/typescript.md
+++ b/docs/typescript.md
@ -0,0 +1,382 @@
+# TypeScript Guide
+
+## Basic usage
+
+When using TypeScript you just have to make a tiny change that instead of writing `create(...)` you'll have to write `create<T>()(...)` where `T` would be type of the state so as to annotate it. Example...
+
+```ts
+import create from "zustand"
+
+interface BearState {
+  bears: number
+  increase: (by: number) => void
+}
+
+const useStore = create<BearState>()((set) => ({
+  bears: 0,
+  increase: (by) => set((state) => ({ bears: state.bears + by })),
+}))
+```
+
+<details>
+  <summary>Why can't we just simply infer the type from initial state?</summary>
+
+  **TLDR**: Because state generic `T` is invariant.
+  
+  Consider this minimal version `create`...
+
+  ```ts
+  declare const create: <T>(f: (get: () => T) => T) => T
+
+  const x = create((get) => ({
+    foo: 0,
+    bar: () => get()
+  }))
+  // `x` is inferred as `unknown` instead of
+  // interface X {
+  //   foo: number,
+  //   bar: () => X
+  // }
+  ```
+
+  Here if you look at the type of `f` in `create` ie `(get: () => T) => T` it "gives" `T` as it returns `T` but then it also "takes" `T` via `get` so where does `T` come from TypeScript thinks... It's a like that chicken or egg problem. At the end TypeScript gives up and infers `T` as `unknown`.
+
+  So as long as the generic to be inferred is invariant TypeScript won't be able to infer it. Another simple example would be this...
+
+  ```ts
+  declare const createFoo: <T>(f: (t: T) => T) => T
+  const x = createFoo(_ => "hello")
+  ```
+
+  Here again `x` is `unknown` instead of `string`.
+  
+  Now one can argue it's impossible to write an implementation for `createFoo`, and that's true. But then it's also impossible to write Zustand's `create`... Wait but Zustand exists? So what do I mean by that?
+
+  The thing is Zustand is lying in it's type, the simplest way to prove it by showing unsoundness. Consider this example...
+
+  ```ts
+  import create from "zustand/vanilla"
+
+  const useStore = create<{ foo: number }>()((_, get) => ({
+    foo: get().foo,
+  }))
+  ```
+
+  This code compiles, but guess what happens when you run it? You'll get an exception "Uncaught TypeError: Cannot read properties of undefined (reading 'foo') because after all `get` would return `undefined` before the initial state is created (hence kids don't call `get` when creating the initial state). But the types tell that get is `() => { foo: number }` which is exactly the lie I was taking about, `get` is that eventually but first it's `() => undefined`.
+
+  Okay we're quite deep in the rabbit hole haha, long story short zustand has a bit crazy runtime behavior that can't be typed in a sound way and inferrable way. We could make it inferrable with the right TypeScript features that don't exist today. And hey that tiny bit of unsoundness is not a problem.
+</details>
+
+<details>
+  <summary>Why that currying `()(...)`?</summary>
+  
+  **TLDR**: It's a workaround for [microsoft/TypeScript#10571](https://github.com/microsoft/TypeScript/issues/10571).
+
+  Imagine you have a scenario like this...
+
+  ```ts
+  declare const withError: <T, E>(p: Promise<T>) =>
+    Promise<[error: undefined, value: T] | [error: E, value: undefined]>
+  declare const doSomething: () => Promise<string>
+
+  const main = async () => {
+    let [error, value] = await withError(doSomething())
+  }
+  ```
+
+  Here `T` is inferred as `string` and `E` is inferred as `unknown`. Now for some reason you want to annotate `E` as `Foo` because you're certain what shape of error `doSomething()` would throw. But too bad you can't do that, you can either pass all generics or none. So now along with annotating `E` as `Foo` you'll also have to annotate `T` as `string` which gets inferred anyway. So what to do? What you do is make a curried version of `withError` that does nothing in runtime, it's purpose is to just allow you annotate `E`...
+
+  ```ts
+  declare const withError: {
+    <E>(): <T>(p: Promise<T>) =>
+      Promise<[error: undefined, value: T] | [error: E, value: undefined]>
+    <T, E>(p: Promise<T>):
+      Promise<[error: undefined, value: T] | [error: E, value: undefined]>
+  }
+  declare const doSomething: () => Promise<string>
+  interface Foo { bar: string }
+
+  const main = async () => {
+    let [error, value] = await withError<Foo>()(doSomething())
+  }
+  ```
+
+  And now `T` gets inferred and you get to annotate `E` too. Zustand has the same use case we want to annotate the state (the first type parameter) but allow the rest type parameters to get inferred.
+</details>
+
+Alternatively you can also use `combine` which infers the state instead of you having to type it...
+
+```ts
+import create from "zustand"
+import { combine } from "zustand/middleware"
+
+const useStore = create(combine({ bears: 0 }, (set) => ({
+  increase: (by: number) => set((state) => ({ bears: state.bears + by })),
+}))
+```
+
+<details>
+  <summary>But be a little careful...</summary>
+
+  We achieve the inference by lying a little in the types of `set`, `get` and `store` that you receive as parameters. The lie is that they're typed in a way as if the state is the first parameter only when in fact the state is the shallow-merge (`{ ...a, ...b }`) of both first parameter and the second parameter's return. So for example `get` from the second parameter has type `() => { bears: number }` and that's a lie as it should be `() => { bears: number, increase: (by: number) => void }`. And `useStore` still has the correct type, ie for example `useStore.getState` is typed as `() => { bears: number, increase: (by: number) => void }`.
+
+  It's not a lie lie because `{ bears: number }` is still a subtype `{ bears: number, increase: (by: number) => void }`, so in most cases there won't be a problem. Just you have to be careful while using replace. For eg `set({ bears: 0 }, true)` would compile but will be unsound as it'll delete the `increase` function. (If you set from "outside" ie `useStore.setState({ bears: 0 }, true)` then it won't compile because the "outside" store knows that `increase` is missing.) Another instance where you should be careful you're doing `Object.keys`, `Object.keys(get())` will return `["bears", "increase"]` and not `["bears"]` (the return type of `get` can make you fall for this).
+
+  So `combine` trades-off a little type-safety for the convience of not having to write a type for state. Hence you should use `combine` accordingly, usually it's not a big deal and it's okay to use it.
+</details>
+
+## Using middlewares
+
+You don't have to do anything special to use middlewares in TypeScript.
+
+```ts
+import create from "zustand"
+import { devtools, persist } from "zustand/middleware"
+
+interface BearState {
+  bears: number
+  increase: (by: number) => void
+}
+
+const useStore = create<BearState>()(devtools(persist((set) => ({
+  bears: 0,
+  increase: (by) => set((state) => ({ bears: state.bears + by })),
+}))))
+```
+
+Just make sure you're using them immediately inside `create` so as to make the contextual inference work. Doing something even remotely fancy like the following `myMiddlewares` would require more advanced types.
+
+```ts
+import create from "zustand"
+import { devtools, persist } from "zustand/middleware"
+
+const myMiddlewares = f => devtools(persist(f))
+
+interface BearState {
+  bears: number
+  increase: (by: number) => void
+}
+
+const useStore = create<BearState>()(myMiddlewares((set) => ({
+  bears: 0,
+  increase: (by) => set((state) => ({ bears: state.bears + by })),
+})))
+```
+
+## Authoring middlewares and advanced usage
+
+Imagine you had to write this hypothetical middleware...
+
+```js
+import create from "zustand"
+
+const foo = (f, bar) => (set, get, store) => {
+  store.foo = bar
+  return f(set, get, store);
+}
+
+const useStore = create(foo(() => ({ bears: 0 }), "hello"))
+console.log(store.foo.toUpperCase())
+```
+
+Yes, if you didn't know Zustand middlewares do and are allowed to mutate the store. But how could we possibly encode the mutation on the type-level? That is to say how could do we type `foo` so that this code compiles?
+
+For an usual statically typed language this is impossible, but thanks to TypeScript, Zustand has something called an "higher kinded mutator" that makes this possible. If you're dealing with complex type problems like typing a middleware or using the `StateCreator` type, then you'll have to understand this implementation detail, for that check out [#710](https://github.com/pmndrs/zustand/issues/710).
+
+If you're eager to know what the answer is to this particular problem then it's [here](#middleware-that-changes-the-store-type).
+
+## Common recipes
+
+### Middleware that does not change the store type
+
+```ts
+import create, { State, StateCreator, StoreMutatorIdentifier, Mutate, StoreApi } from "zustand"
+
+type Logger =
+  < T extends State
+  , Mps extends [StoreMutatorIdentifier, unknown][] = []
+  , Mcs extends [StoreMutatorIdentifier, unknown][] = []
+  >
+    ( f: StateCreator<T, Mps, Mcs>
+    , name?: string
+    ) =>
+      StateCreator<T, Mps, Mcs>
+
+type LoggerImpl =
+  <T extends State>
+    ( f: PopArgument<StateCreator<T, [], []>>
+    , name?: string
+    ) =>
+      PopArgument<StateCreator<T, [], []>>
+
+const loggerImpl: LoggerImpl = (f, name) => (set, get, store) => {
+  type T = ReturnType<typeof f>
+  const loggedSet: typeof set = (...a) => {
+    set(...a)
+    console.log(...(name ? [`${name}:`] : []), get())
+  }
+  store.setState = loggedState
+
+  return f(loggedSet, get, store)
+}
+
+export const logger = loggerImpl as unknown as Foo
+
+type PopArgument<T extends (...a: never[]) => unknown> =
+  T extends (...a: [...infer A, infer _]) => infer R
+    ? (...a: A) => R
+    : never
+
+// ---
+
+const useStore = create<BearState>()(logger((set) => ({
+  bears: 0,
+  increase: (by) => set((state) => ({ bears: state.bears + by })),
+}), "bear-store"))
+```
+
+### Middleware that changes the store type
+
+```js
+import create, { State, StateCreator, StoreMutatorIdentifier, Mutate, StoreApi } from "zustand"
+
+type Foo =
+  < T extends State
+  , A
+  , Mps extends [StoreMutatorIdentifier, unknown][] = []
+  , Mcs extends [StoreMutatorIdentifier, unknown][] = []
+  >
+    ( f: StateCreator<T, [...Mps, ['foo', A]], Mcs>
+    , bar: A
+    ) =>
+      StateCreator<T, Mps, [['foo', A], ...Mcs]>
+
+declare module 'zustand' {
+  interface StoreMutators<S, A> {
+    foo: Write<Cast<S, object> { foo: A }>
+  }
+}
+
+type FooImpl =
+  <T extends State, A>
+    ( f: PopArgument<StateCreator<T, [], []>>
+    , bar: A
+    ) => PopArgument<StateCreator<T, [], []>>
+
+const fooImpl: FooImpl = (f, bar) => (set, get, _store) => {
+  type T = ReturnType<typeof f>
+  type A = typeof bar
+
+  const store = _store as Mutate<StoreApi<T>, [['foo', A]]>
+  store.foo = bar
+  return f(set, get, _store)
+}
+
+export const foo = fooImpl as unknown as Foo
+
+type PopArgument<T extends (...a: never[]) => unknown> =
+  T extends (...a: [...infer A, infer _]) => infer R
+    ? (...a: A) => R
+    : never
+
+type Write<T extends object, U extends object> =
+  Omit<T, keyof U> & U
+
+type Cast<T, U> =
+  T extends U ? T : U;
+
+// ---
+
+const useStore = create(foo(() => ({ bears: 0 }), "hello"))
+console.log(store.foo.toUpperCase())
+```
+
+### `create` without curried workaround
+
+The recommended way to use `create` is using the curried workaround ie `create<T>()(...)` because this enabled you to infer the store type. But for some reason if you don't want to use the workaround then you can pass the type parameters like the following. Note that in some cases this acts as an assertion instead of annotation, so it's not recommended.
+
+```ts
+import create, { State, StateCreator, StoreMutatorIdentifier, Mutate, StoreApi } from "zustand"
+
+interface BearState {
+  bears: number
+  increase: (by: number) => void
+}
+
+const useStore = create<
+  BearState,
+  [
+    ['zustand/persist', BearState],
+    ['zustand/devtools', never]
+  ]
+>(devtools(persist((set) => ({
+  bears: 0,
+  increase: (by) => set((state) => ({ bears: state.bears + by })),
+})))
+```
+
+### Independent slices pattern
+
+```ts
+import create, { State, StateCreator, StoreMutatorIdentifier, Mutate, StoreApi } from "zustand"
+
+interface BearSlice {
+  bears: number
+  addBear: () => void
+}
+const createBearSlice: StateCreator<BearSlice, [], []> = (set) => ({
+  bears: 0,
+  addBear: () => set((state) => ({ bears: state.bears + 1 }))
+})
+
+interface FishSlice {
+  fishes: number
+  addFish: () => void
+}
+const createFishSlice: StateCreator<FishSlice, [], []> = (set) => ({
+  fishes: 0,
+  addFish: () => set((state) => ({ fishes: state.fishes + 1 }))
+})
+
+const useStore = create<BearSlice & FishSlice>()((...a) => ({
+  ...createBearSlice(...a),
+  ...createFishSlice(...a)
+}))
+```
+
+If you have some middlewares then replace `StateCreator<MySlice, [], []>` with `StateCreator<MySlice, Mutators, []>`. Eg if you're using `devtools` then it'll be `StateCreator<MySlice, [["zustand/devtools", never]], []>`.
+
+Also you can even write `StateCreator<MySlice>` instead of `StateCreator<MySlice, [], []>` as the second and third parameter have `[]` as their default value.
+
+### Interdependent slices pattern
+
+```ts
+import create, { State, StateCreator, StoreMutatorIdentifier, Mutate, StoreApi } from "zustand"
+
+interface BearSlice {
+  bears: number
+  addBear: () => void
+  eatFish: () => void
+}
+const createBearSlice: StateCreator<BearSlice & FishSlice, [], [], BearSlice> = (set) => ({
+  bears: 0,
+  addBear: () => set((state) => ({ bears: state.bears + 1 })),
+  eatFish: () => set((state) => ({ fishes: state.fishes - 1 }))
+})
+
+interface FishSlice {
+  fishes: number
+  addFish: () => void
+}
+const createFishSlice: StateCreator<BearSlice & FishSlice, [], [], FishSlice> = (set) => ({
+  fishes: 0,
+  addFish: () => set((state) => ({ fishes: state.fishes + 1 }))
+})
+
+const useStore = create<BearSlice & FishSlice>()((...a) => ({
+  ...createBearSlice(...a),
+  ...createFishSlice(...a)
+}))
+```
+
+If you have some middlewares then replace `StateCreator<MyState, [], [], MySlice>` with `StateCreator<MyState, Mutators, [], MySlice>`. Eg if you're using `devtools` then it'll be `StateCreator<MyState, [["zustand/devtools", never]], [], MySlice>`.
--- a/docs/v4-migration.md
+++ b/docs/v4-migration.md
@ -0,0 +1,205 @@
+# v4 Migrations
+
+If you're not using the typed version (either via TypeScript or via JSDoc) then there are no breaking changes for you and hence no migration is needed either.
+
+Also it's recommended to first read the new [TypeScript Guide](https://github.com/pmndrs/zustand/blob/main/docs/typescript.md), it'll be easier to understand the migration.
+
+In addition to this migration guide you can also check the diff of the test files in the repo from v3 to v4.
+
+## `create` (from `zustand` and `zustand/vanilla`)
+
+### Change
+
+```diff
+- create:
+-   < State
+-   , StoreSetState = StoreApi<State>["set"]
+-   , StoreGetState = StoreApi<State>["get"]
+-   , Store = StoreApi<State>
+-   >
+-     (f: ...) => ...
+ create:
+   { <State>(): (f: ...) => ...
+   , <State, Mutators>(f: ...) => ...
+   }
+```
+
+### Migration
+
+If you're not passing any type parameters to `create` then there is no migration needed. If you're using a "leaf" middleware like `combine` or `redux` then remove all type parameters from `create`. Else replace `create<T, ...>(...)` with `create<T>()(...)`.
+
+## `StateCreator` (from `zustand` and `zustand/vanilla`)
+
+### Change
+
+```diff
+- type StateCreator
+-   < State
+-   , StoreSetState = StoreApi<State>["set"]
+-   , StoreGetState = StoreApi<State>["get"]
+-   , Store = StoreApi<State>
+-   > =
+-     ...
+ type StateCreator
+   < State
+   , InMutators extends [StoreMutatorIdentifier, unknown][] = []
+   , OutMutators extends [StoreMutatorIdentifier, unknown][] = []
+   , Return = State
+   > =
+     ...
+```
+
+### Migration
+
+If you're using `StateCreator` you're likely authoring a middleware or using the "slices" pattern, for that check the TypeScript Guide's ["Authoring middlewares and advanced usage"](https://github.com/pmndrs/zustand/blob/main/docs/typescript.md#authoring-middlewares-and-advanced-usage) and ["Common recipes"](https://github.com/pmndrs/zustand/blob/main/docs/typescript.md#authoring-middlewares-and-advanced-usage) sections.
+
+## `PartialState` (from `zustand` and `zustand/vanilla`)
+
+### Change
+
+```diff
+- type PartialState
+-   < T extends State
+-   , K1 extends keyof T = keyof T
+-   , K2 extends keyof T = K1
+-   , K3 extends keyof T = K2
+-   , K4 extends keyof T = K3
+-   > =
+-   | (Pick<T, K1> | Pick<T, K2> | Pick<T, K3> | Pick<T, K4> | T)
+-   | ((state: T) => Pick<T, K1> | Pick<T, K2> | Pick<T, K3> | Pick<T, K4> | T)
+ type PartialState<T> =
+   | Partial<T>
+   | ((state: T) => Partial<T>)
+```
+
+### Migration
+
+Replace `PartialState<T, ...>` with `PartialState<T>` and preferably turn on [`--exactOptionalPropertyTypes`](https://www.typescriptlang.org/tsconfig#exactOptionalPropertyTypes)/
+
+We're no longer using the trick to disallow `{ foo: undefined }` to be assigned to `Partial<{ foo: string }>` instead now we're relying on the users to turn on `--exactOptionalPropertyTypes`.
+
+## `useStore` (from `zustand` and `zustand/react`)
+
+### Change
+
+```diff
+- useStore:
+-   { <State>(store: StoreApi<State>): State
+-   , <State, StateSlice>
+-       ( store: StoreApi<State>
+-       , selector: StateSelector<State, StateSlice>,
+-       , equals?: EqualityChecker<StateSlice>
+-       ): StateSlice
+-   }
+ useStore:
+   <Store, StateSlice = ExtractState<Store>>
+     ( store: Store
+     , selector?: StateSelector<State, StateSlice>,
+     , equals?: EqualityChecker<StateSlice>
+     )
+       => StateSlice
+```
+
+### Migration
+
+If you're not passing any type parameters to `useStore` then there is no migration needed. If you are then it's recommended to remove them, or pass the store type instead of the state type as the first parameter.
+
+## `UseBoundStore` (from `zustand` and `zustand/react`)
+
+### Change 
+
+```diff
+- type UseBoundStore<
+-   State,
+-   Store = StoreApi<State>
+- > =
+-   & { (): T
+-     , <StateSlice>
+-         ( selector: StateSelector<State, StateSlice>
+-         , equals?: EqualityChecker<StateSlice>
+-         ): U
+-     }
+-   & Store
+ type UseBoundStore<Store> =
+   & (<StateSlice = ExtractState<S>>
+       ( selector?: (state: ExtractState<S>) => StateSlice
+       , equals?: (a: StateSlice, b: StateSlice) => boolean
+       ) => StateSlice
+     )
+   & S
+```
+
+### Migration
+
+Replace `UseBoundStore<T>` with `UseBoundStore<StoreApi<T>>` and `UseBoundStore<T, S>` with `UseBoundStore<S>`
+
+## `UseContextStore` (from `zustand/context`)
+
+### Change
+
+```diff
+- type UseContextStore
+```
+
+### Migration
+
+Use `typeof MyContext.useStore` instead
+
+## `createContext` (from `zustand/context`)
+
+### Change
+
+```diff
+  createContext:
+-   <State, Store = StoreApi<State>>() => ...
+   <Store>() => ...
+```
+
+### Migration
+
+Replace `createContext<T>()` with `createContext<StoreApi<T>>()` and `createContext<T, S>()` with `createContext<S>()`.
+
+## `combine`, `devtools`, `persist`, `subscribeWithSelector` (from `zustand/middleware`)
+
+### Change
+
+```diff
+- combine:
+-   <T, U>(...) => ...
+ combine:
+   <T, U, Mps, Mcs>(...) => ...
+
+- devtools:
+-   <T>(...) => ...
+ devtools:
+   <T, Mps, Mcs>(...) => ...
+
+- persist:
+-   <T, U>(...) => ...
+ persist:
+   <T, U, Mps, Mcs>(...) => ...
+
+- subscribeWithSelector:
+-   <T>(...) => ...
+ subscribeWithSelector:
+   <T, Mps, Mcs>(...) => ...
+```
+
+### Migration
+
+If you're not passing any type parameters then there is no migration needed. If you're passing any type parameters, remove them as are inferred.
+
+## `redux` (from `zustand/middleware`)
+
+### Change
+
+```diff
+- redux:
+-   <T, A>(...) => ...
+ redux:
+   <T, A, Mps, Mcs>(...) => ...
+```
+
+### Migration
+
+If you're not passing any type parameters then there is no migration needed. If you're passing type parameters them remove them and annotate the second (action) parameter. That is replace `redux<T, A>((state, action) => ..., ...)` with `redux((state, action: A) => ..., ...)`.
--- a/readme.md
+++ b/readme.md
@ -196,30 +196,6 @@ const unsub4 = useStore.subscribe(state => [state.paw, state.fur], console.log,
 const unsub5 = useStore.subscribe(state => state.paw, console.log, { fireImmediately: true })
 ```

-<details>
-<summary>How to type store with `subscribeWithSelector` in TypeScript</summary>
-
-```ts
-import create, { Mutate, GetState, SetState, StoreApi } from 'zustand'
-import { subscribeWithSelector } from 'zustand/middleware'
-
-type BearState = {
-  paw: boolean
-  snout: boolean
-  fur: boolean
-}
-const useStore = create<
-  BearState,
-  SetState<BearState>,
-  GetState<BearState>,
-  Mutate<StoreApi<BearState>, [["zustand/subscribeWithSelector", never]]>
->(subscribeWithSelector(() => ({ paw: true, snout: true, fur: true })))
-```
-
-For more complex typing with multiple middlewares,
-Please refer [middlewareTypes.test.tsx](./tests/middlewareTypes.test.tsx).
-</details>
-
 ## Using zustand without React

 Zustands core can be imported and used without the React dependency. The only difference is that the create function does not return a hook, but the api utilities.
@ -306,36 +282,6 @@ const useStore = create(
 )
 ```

-<details>
-<summary>How to pipe middlewares</summary>
-
-```js
-import create from "zustand"
-import produce from "immer"
-import pipe from "ramda/es/pipe"
-
-/* log and immer functions from previous example */
-/* you can pipe as many middlewares as you want */
-const createStore = pipe(log, immer, create)
-
-const useStore = createStore(set => ({
-  bears: 1,
-  increasePopulation: () => set(state => ({ bears: state.bears + 1 }))
-}))
-
-export default useStore
-```
-
-For a TS example see the following [discussion](https://github.com/pmndrs/zustand/discussions/224#discussioncomment-118208)
-</details>
-
-<details>
-<summary>How to type immer middleware in TypeScript</summary>
-
-There is a reference implementation in [middlewareTypes.test.tsx](./tests/middlewareTypes.test.tsx) with some use cases.
-You can use any simplified variant based on your requirement.
-</details>
-
 ## Persist middleware

 You can persist your store's data using any kind of storage.
@ -581,48 +527,32 @@ const Component = () => {
      >
        <Button />
      </Provider>
-  )
-}
+    )
+  }
  ```
 </details>

-## Typing your store and `combine` middleware
+## TypeScript Usage

-```tsx
-// You can use `type`
-type BearState = {
-  bears: number
-  increase: (by: number) => void
-}
+Basic typescript usage doesn't require anything special except for writing `create<State>()(...)` instead of `create(...)`...
+
+```ts
+import create from "zustand"
+import { devtools, persist } from "zustand/middleware"

-// Or `interface`
 interface BearState {
  bears: number
  increase: (by: number) => void
 }

-// And it is going to work for both
-const useStore = create<BearState>(set => ({
+const useStore = create<BearState>()(devtools(persist((set) => ({
  bears: 0,
-  increase: (by) => set(state => ({ bears: state.bears + by })),
-}))
+  increase: (by) => set((state) => ({ bears: state.bears + by })),
+}))))
 ```

-Or, use `combine` and let tsc infer types. This merges two states shallowly.
+A more complete TypeScript guide is [here](https://github.com/pmndrs/zustand/blob/main/docs/typescript.md).

-```tsx
-import { combine } from 'zustand/middleware'
-
-const useStore = create(
-  combine(
-    { bears: 0 },
-    (set) => ({ increase: (by: number) => set((state) => ({ bears: state.bears + by })) })
-  ),
-)
-```
-
-Typing with multiple middleware might require some TypeScript knowledge. Refer some working examples in [middlewareTypes.test.tsx](./tests/middlewareTypes.test.tsx).
-  
 ## Best practices
  
 * You may wonder how to organize your code for better maintenance: [Splitting the store into seperate slices](https://github.com/pmndrs/zustand/wiki/Splitting-the-store-into-separate-slices).
--- a/src/context.ts
+++ b/src/context.ts
@ -14,30 +14,29 @@ import {
  useStore,
 } from 'zustand'

-/**
- * @deprecated Use `typeof MyContext.useStore` instead.
- */
-export type UseContextStore<T extends State> = {
-  (): T
-  <U>(selector: StateSelector<T, U>, equalityFn?: EqualityChecker<U>): U
+type UseContextStore<S extends StoreApi<State>> = {
+  (): ExtractState<S>
+  <U>(
+    selector: StateSelector<ExtractState<S>, U>,
+    equalityFn?: EqualityChecker<U>
+  ): U
 }

-function createContext<
-  TState extends State,
-  CustomStoreApi extends StoreApi<TState> = StoreApi<TState>
->() {
-  const ZustandContext = reactCreateContext<CustomStoreApi | undefined>(
-    undefined
-  )
+type ExtractState<S> = S extends { getState: () => infer T } ? T : never
+
+type WithoutCallSignature<T> = { [K in keyof T]: T[K] }
+
+function createContext<S extends StoreApi<State>>() {
+  const ZustandContext = reactCreateContext<S | undefined>(undefined)

  const Provider = ({
    createStore,
    children,
  }: {
-    createStore: () => CustomStoreApi
+    createStore: () => S
    children: ReactNode
  }) => {
-    const storeRef = useRef<CustomStoreApi>()
+    const storeRef = useRef<S>()

    if (!storeRef.current) {
      storeRef.current = createStore()
@ -50,8 +49,8 @@ function createContext<
    )
  }

-  const useBoundStore: UseContextStore<TState> = <StateSlice>(
-    selector?: StateSelector<TState, StateSlice>,
+  const useBoundStore = (<StateSlice = ExtractState<S>>(
+    selector?: StateSelector<ExtractState<S>, StateSlice>,
    equalityFn?: EqualityChecker<StateSlice>
  ) => {
    const store = useContext(ZustandContext)
@ -62,17 +61,12 @@ function createContext<
    }
    return useStore(
      store,
-      selector as StateSelector<TState, StateSlice>,
+      selector as StateSelector<ExtractState<S>, StateSlice>,
      equalityFn
    )
-  }
+  }) as UseContextStore<S>

-  const useStoreApi = (): {
-    getState: CustomStoreApi['getState']
-    setState: CustomStoreApi['setState']
-    subscribe: CustomStoreApi['subscribe']
-    destroy: CustomStoreApi['destroy']
-  } => {
+  const useStoreApi = () => {
    const store = useContext(ZustandContext)
    if (!store) {
      throw new Error(
@ -80,12 +74,13 @@ function createContext<
      )
    }
    return useMemo(
-      () => ({
-        getState: store.getState,
-        setState: store.setState,
-        subscribe: store.subscribe,
-        destroy: store.destroy,
-      }),
+      () =>
+        ({
+          getState: store.getState,
+          setState: store.setState,
+          subscribe: store.subscribe,
+          destroy: store.destroy,
+        } as WithoutCallSignature<S>),
      [store]
    )
  }
--- a/src/middleware.ts
+++ b/src/middleware.ts
@ -3,3 +3,4 @@ export * from './middleware/devtools'
 export * from './middleware/subscribeWithSelector'
 export * from './middleware/combine'
 export * from './middleware/persist'
+export * from './middleware/immer'
--- a/src/middleware/combine.ts
+++ b/src/middleware/combine.ts
@ -1,25 +1,18 @@
-import { GetState, SetState, State, StoreApi } from '../vanilla'
-import { NamedSet } from './devtools'
+import { State, StateCreator, StoreMutatorIdentifier } from '../vanilla'

-type Combine<T, U> = Omit<T, keyof U> & U
+type Write<T, U> = Omit<T, keyof U> & U

-export const combine =
-  <PrimaryState extends State, SecondaryState extends State>(
-    initialState: PrimaryState,
-    create: (
-      // Note: NamedSet added for convenience
-      set: SetState<PrimaryState> & NamedSet<PrimaryState>,
-      get: GetState<PrimaryState>,
-      api: StoreApi<PrimaryState>
-    ) => SecondaryState
-  ) =>
-  (
-    set: SetState<Combine<PrimaryState, SecondaryState>>,
-    get: GetState<Combine<PrimaryState, SecondaryState>>,
-    api: StoreApi<Combine<PrimaryState, SecondaryState>>
-  ) =>
-    Object.assign(
-      {},
-      initialState,
-      create(set as any, get as any, api as any)
-    ) as Combine<PrimaryState, SecondaryState>
+type Combine = <
+  T extends State,
+  U extends State,
+  Mps extends [StoreMutatorIdentifier, unknown][] = [],
+  Mcs extends [StoreMutatorIdentifier, unknown][] = []
+>(
+  initialState: T,
+  additionalStateCreator: StateCreator<T, Mps, Mcs, U>
+) => StateCreator<Write<T, U>, Mps, Mcs>
+
+export const combine: Combine =
+  (initialState, create) =>
+  (...a) =>
+    Object.assign({}, initialState, (create as any)(...a))
--- a/src/middleware/devtools.ts
+++ b/src/middleware/devtools.ts
@ -1,5 +1,12 @@
 import type {} from '@redux-devtools/extension'
-import { GetState, PartialState, SetState, State, StoreApi } from '../vanilla'
+import {
+  PartialState,
+  SetState,
+  State,
+  StateCreator,
+  StoreApi,
+  StoreMutatorIdentifier,
+} from '../vanilla'

 declare module '../vanilla' {
  // eslint-disable-next-line @typescript-eslint/no-unused-vars
@ -17,11 +24,36 @@ type Message = {

 type Write<T extends object, U extends object> = Omit<T, keyof U> & U
 type Cast<T, U> = T extends U ? T : U
+type TakeTwo<T> = T extends []
+  ? [undefined, undefined]
+  : T extends [unknown]
+  ? [...a0: T, a1: undefined]
+  : T extends [unknown?]
+  ? [...a0: T, a1: undefined]
+  : T extends [unknown, unknown]
+  ? T
+  : T extends [unknown, unknown?]
+  ? T
+  : T extends [unknown?, unknown?]
+  ? T
+  : T extends [infer A0, infer A1, ...unknown[]]
+  ? [A0, A1]
+  : T extends [infer A0, (infer A1)?, ...unknown[]]
+  ? [A0, A1?]
+  : T extends [(infer A0)?, (infer A1)?, ...unknown[]]
+  ? [A0?, A1?]
+  : never

 type WithDevtools<S> = Write<Cast<S, object>, StoreSetStateWithAction<S>>

-type StoreSetStateWithAction<S> = S extends { getState: () => infer T }
-  ? S & { setState: NamedSet<Cast<T, object>> }
+type StoreSetStateWithAction<S> = S extends {
+  setState: (...a: infer A) => infer Sr
+}
+  ? {
+      setState(
+        ...a: [...a: TakeTwo<A>, actionType?: string | { type: unknown }]
+      ): Sr
+    }
  : never

 interface DevtoolsOptions {
@ -42,226 +74,183 @@ interface DevtoolsOptions {
      }
 }

-export type NamedSet<T extends State> = {
-  <
-    K1 extends keyof T,
-    K2 extends keyof T = K1,
-    K3 extends keyof T = K2,
-    K4 extends keyof T = K3
-  >(
-    partial: PartialState<T, K1, K2, K3, K4>,
-    replace?: boolean,
-    name?: string | { type: unknown }
-  ): void
-}
-/**
- * @deprecated Use `Mutate<StoreApi<T>, [["zustand/devtools", never]]>`.
- * See tests/middlewaresTypes.test.tsx for usage with multiple middlewares.
- */
-export type StoreApiWithDevtools<T extends State> = StoreApi<T> & {
-  setState: NamedSet<T>
-}
-
-export function devtools<
-  S extends State,
-  CustomSetState extends SetState<S>,
-  CustomGetState extends GetState<S>,
-  CustomStoreApi extends StoreApi<S>
+type Devtools = <
+  T extends State,
+  Mps extends [StoreMutatorIdentifier, unknown][] = [],
+  Mcs extends [StoreMutatorIdentifier, unknown][] = []
 >(
-  fn: (set: NamedSet<S>, get: CustomGetState, api: CustomStoreApi) => S
-): (
-  set: CustomSetState,
-  get: CustomGetState,
-  api: CustomStoreApi & StoreApiWithDevtools<S>
-) => S
-/**
- * @deprecated Passing `name` as directly will be not allowed in next major.
- * Pass the `name` in an object `{ name: ... }` instead
- */
-export function devtools<
-  S extends State,
-  CustomSetState extends SetState<S> = SetState<S>,
-  CustomGetState extends GetState<S> = GetState<S>,
-  CustomStoreApi extends StoreApi<S> = StoreApi<S>
->(
-  fn: (set: NamedSet<S>, get: CustomGetState, api: CustomStoreApi) => S,
-  options?: string
-): (
-  set: CustomSetState,
-  get: CustomGetState,
-  api: CustomStoreApi & StoreApiWithDevtools<S>
-) => S
-export function devtools<
-  S extends State,
-  CustomSetState extends SetState<S>,
-  CustomGetState extends GetState<S>,
-  CustomStoreApi extends StoreApi<S>
->(
-  fn: (set: NamedSet<S>, get: CustomGetState, api: CustomStoreApi) => S,
+  initializer: StateCreator<T, [...Mps, ['zustand/devtools', never]], Mcs>,
  options?: DevtoolsOptions
-): (
-  set: CustomSetState,
-  get: CustomGetState,
-  api: CustomStoreApi & StoreApiWithDevtools<S>
-) => S
-export function devtools<
-  S extends State,
-  CustomSetState extends SetState<S>,
-  CustomGetState extends GetState<S>,
-  CustomStoreApi extends StoreApi<S>
->(
-  fn: (set: NamedSet<S>, get: CustomGetState, api: CustomStoreApi) => S,
-  options?: string | DevtoolsOptions
-) {
-  return (
-    set: CustomSetState,
-    get: CustomGetState,
-    api: CustomStoreApi & StoreApiWithDevtools<S>
-  ): S => {
-    const devtoolsOptions =
-      options === undefined
-        ? {}
-        : typeof options === 'string'
-        ? { name: options }
-        : options
+) => StateCreator<T, Mps, [['zustand/devtools', never], ...Mcs]>

-    let extensionConnector: typeof window['__REDUX_DEVTOOLS_EXTENSION__']
-    try {
-      extensionConnector = window.__REDUX_DEVTOOLS_EXTENSION__
-    } catch {
-      // ignored
-    }
+declare module '../vanilla' {
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  interface StoreMutators<S, A> {
+    'zustand/devtools': WithDevtools<S>
+  }
+}

-    if (!extensionConnector) {
-      if (__DEV__ && typeof window !== 'undefined') {
-        console.warn(
-          '[zustand devtools middleware] Please install/enable Redux devtools extension'
-        )
-      }
-      return fn(set, get, api)
-    }
+type DevtoolsImpl = <T extends State>(
+  storeInitializer: PopArgument<StateCreator<T, [], []>>,
+  options: DevtoolsOptions
+) => PopArgument<StateCreator<T, [], []>>

-    const extension = extensionConnector.connect(devtoolsOptions)
+type PopArgument<T extends (...a: never[]) => unknown> = T extends (
+  ...a: [...infer A, infer _]
+) => infer R
+  ? (...a: A) => R
+  : never

-    let isRecording = true
-    ;(api.setState as NamedSet<S>) = (state, replace, nameOrAction) => {
-      set(state, replace)
-      if (!isRecording) return
-      extension.send(
-        nameOrAction === undefined
-          ? { type: devtoolsOptions.anonymousActionType || 'anonymous' }
-          : typeof nameOrAction === 'string'
-          ? { type: nameOrAction }
-          : nameOrAction,
-        get()
+export type NamedSet<T extends State> = WithDevtools<StoreApi<T>>['setState']
+
+const devtoolsImpl: DevtoolsImpl = (fn, options) => (set, get, api) => {
+  type S = ReturnType<typeof fn>
+
+  const devtoolsOptions =
+    options === undefined
+      ? {}
+      : typeof options === 'string'
+      ? { name: options }
+      : options
+
+  let extensionConnector
+  try {
+    extensionConnector = window.__REDUX_DEVTOOLS_EXTENSION__
+  } catch {
+    // ignored
+  }
+
+  if (!extensionConnector) {
+    if (__DEV__ && typeof window !== 'undefined') {
+      console.warn(
+        '[zustand devtools middleware] Please install/enable Redux devtools extension'
      )
    }
-    const setStateFromDevtools: SetState<S> = (...a) => {
-      const originalIsRecording = isRecording
-      isRecording = false
-      set(...a)
-      isRecording = originalIsRecording
+    return fn(set, get, api)
+  }
+
+  const extension = extensionConnector.connect(devtoolsOptions)
+
+  let isRecording = true
+  ;(api.setState as NamedSet<S>) = (state, replace, nameOrAction) => {
+    const r = set(state, replace)
+    if (!isRecording) return r
+    extension.send(
+      nameOrAction === undefined
+        ? { type: devtoolsOptions.anonymousActionType || 'anonymous' }
+        : typeof nameOrAction === 'string'
+        ? { type: nameOrAction }
+        : nameOrAction,
+      get()
+    )
+    return r
+  }
+  const setStateFromDevtools: SetState<S> = (...a) => {
+    const originalIsRecording = isRecording
+    isRecording = false
+    set(...a)
+    isRecording = originalIsRecording
+  }
+
+  const initialState = fn(api.setState, get, api)
+  extension.init(initialState)
+
+  if (
+    (api as any).dispatchFromDevtools &&
+    typeof (api as any).dispatch === 'function'
+  ) {
+    let didWarnAboutReservedActionType = false
+    const originalDispatch = (api as any).dispatch
+    ;(api as any).dispatch = (...a: any[]) => {
+      if (
+        __DEV__ &&
+        a[0].type === '__setState' &&
+        !didWarnAboutReservedActionType
+      ) {
+        console.warn(
+          '[zustand devtools middleware] "__setState" action type is reserved ' +
+            'to set state from the devtools. Avoid using it.'
+        )
+        didWarnAboutReservedActionType = true
+      }
+      ;(originalDispatch as any)(...a)
    }
+  }

-    const initialState = fn(api.setState, get, api)
-    extension.init(initialState)
-
-    if (
-      (api as any).dispatchFromDevtools &&
-      typeof (api as any).dispatch === 'function'
-    ) {
-      let didWarnAboutReservedActionType = false
-      const originalDispatch = (api as any).dispatch
-      ;(api as any).dispatch = (...a: any[]) => {
-        if (
-          __DEV__ &&
-          a[0].type === '__setState' &&
-          !didWarnAboutReservedActionType
-        ) {
-          console.warn(
-            '[zustand devtools middleware] "__setState" action type is reserved ' +
-              'to set state from the devtools. Avoid using it.'
+  ;(
+    extension as unknown as {
+      // FIXME https://github.com/reduxjs/redux-devtools/issues/1097
+      subscribe: (
+        listener: (message: Message) => void
+      ) => (() => void) | undefined
+    }
+  ).subscribe((message: any) => {
+    switch (message.type) {
+      case 'ACTION':
+        if (typeof message.payload !== 'string') {
+          console.error(
+            '[zustand devtools middleware] Unsupported action format'
          )
-          didWarnAboutReservedActionType = true
+          return
        }
-        ;(originalDispatch as any)(...a)
-      }
-    }
-
-    ;(
-      extension as unknown as {
-        // FIXME https://github.com/reduxjs/redux-devtools/issues/1097
-        subscribe: (
-          listener: (message: Message) => void
-        ) => (() => void) | undefined
-      }
-    ).subscribe((message) => {
-      switch (message.type) {
-        case 'ACTION':
-          if (typeof message.payload !== 'string') {
-            console.error(
-              '[zustand devtools middleware] Unsupported action format'
-            )
-            return
-          }
-          return parseJsonThen<{ type: unknown; state?: PartialState<S> }>(
-            message.payload,
-            (action) => {
-              if (action.type === '__setState') {
-                setStateFromDevtools(action.state as PartialState<S>)
-                return
-              }
-
-              if (!(api as any).dispatchFromDevtools) return
-              if (typeof (api as any).dispatch !== 'function') return
-              ;(api as any).dispatch(action)
-            }
-          )
-
-        case 'DISPATCH':
-          switch (message.payload.type) {
-            case 'RESET':
-              setStateFromDevtools(initialState)
-              return extension.init(api.getState())
-
-            case 'COMMIT':
-              return extension.init(api.getState())
-
-            case 'ROLLBACK':
-              return parseJsonThen<S>(message.state, (state) => {
-                setStateFromDevtools(state)
-                extension.init(api.getState())
-              })
-
-            case 'JUMP_TO_STATE':
-            case 'JUMP_TO_ACTION':
-              return parseJsonThen<S>(message.state, (state) => {
-                setStateFromDevtools(state)
-              })
-
-            case 'IMPORT_STATE': {
-              const { nextLiftedState } = message.payload
-              const lastComputedState =
-                nextLiftedState.computedStates.slice(-1)[0]?.state
-              if (!lastComputedState) return
-              setStateFromDevtools(lastComputedState)
-              extension.send(
-                null as any, // FIXME no-any
-                nextLiftedState
-              )
+        return parseJsonThen<{ type: unknown; state?: PartialState<S> }>(
+          message.payload,
+          (action) => {
+            if (action.type === '__setState') {
+              setStateFromDevtools(action.state as PartialState<S>)
              return
            }

-            case 'PAUSE_RECORDING':
-              return (isRecording = !isRecording)
+            if (!(api as any).dispatchFromDevtools) return
+            if (typeof (api as any).dispatch !== 'function') return
+            ;(api as any).dispatch(action)
          }
-          return
-      }
-    })
+        )

-    return initialState
-  }
+      case 'DISPATCH':
+        switch (message.payload.type) {
+          case 'RESET':
+            setStateFromDevtools(initialState)
+            return extension.init(api.getState())
+
+          case 'COMMIT':
+            return extension.init(api.getState())
+
+          case 'ROLLBACK':
+            return parseJsonThen<S>(message.state, (state) => {
+              setStateFromDevtools(state)
+              extension.init(api.getState())
+            })
+
+          case 'JUMP_TO_STATE':
+          case 'JUMP_TO_ACTION':
+            return parseJsonThen<S>(message.state, (state) => {
+              setStateFromDevtools(state)
+            })
+
+          case 'IMPORT_STATE': {
+            const { nextLiftedState } = message.payload
+            const lastComputedState =
+              nextLiftedState.computedStates.slice(-1)[0]?.state
+            if (!lastComputedState) return
+            setStateFromDevtools(lastComputedState)
+            extension.send(
+              null as any, // FIXME no-any
+              nextLiftedState
+            )
+            return
+          }
+
+          case 'PAUSE_RECORDING':
+            return (isRecording = !isRecording)
+        }
+        return
+    }
+  })
+
+  return initialState
 }
+export const devtools = devtoolsImpl as unknown as Devtools

 const parseJsonThen = <T>(stringified: string, f: (parsed: T) => void) => {
  let parsed: T | undefined
--- a/src/middleware/immer.ts
+++ b/src/middleware/immer.ts
@ -0,0 +1,77 @@
+// eslint-disable-next-line import/named
+import { Draft, produce } from 'immer'
+import { State, StateCreator, StoreMutatorIdentifier } from '../vanilla'
+
+type Immer = <
+  T extends State,
+  Mps extends [StoreMutatorIdentifier, unknown][] = [],
+  Mcs extends [StoreMutatorIdentifier, unknown][] = []
+>(
+  initializer: StateCreator<T, [...Mps, ['zustand/immer', never]], Mcs>
+) => StateCreator<T, Mps, [['zustand/immer', never], ...Mcs]>
+
+declare module 'zustand' {
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  interface StoreMutators<S, A> {
+    ['zustand/immer']: WithImmer<S>
+  }
+}
+
+type Write<T extends object, U extends object> = Omit<T, keyof U> & U
+type SkipTwo<T> = T extends []
+  ? []
+  : T extends [unknown]
+  ? []
+  : T extends [unknown?]
+  ? []
+  : T extends [unknown, unknown, ...infer A]
+  ? A
+  : T extends [unknown, unknown?, ...infer A]
+  ? A
+  : T extends [unknown?, unknown?, ...infer A]
+  ? A
+  : never
+
+type WithImmer<S> = S extends {
+  getState: () => infer T
+  setState: infer SetState
+}
+  ? Write<
+      S,
+      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
+
+type PopArgument<T extends (...a: never[]) => unknown> = T extends (
+  ...a: [...infer A, infer _]
+) => infer R
+  ? (...a: A) => R
+  : never
+
+type ImmerImpl = <T extends State>(
+  storeInitializer: PopArgument<StateCreator<T, [], []>>
+) => PopArgument<StateCreator<T, [], []>>
+
+const immerImpl: ImmerImpl = (initializer) => (set, get, store) => {
+  type T = ReturnType<typeof initializer>
+
+  store.setState = (updater, replace, ...a) => {
+    const nextState = (
+      typeof updater === 'function' ? produce(updater as any) : updater
+    ) as ((s: T) => T) | T | Partial<T>
+
+    return set(nextState as any, replace, ...a)
+  }
+
+  return initializer(store.setState, get, store)
+}
+
+export const immer = immerImpl as unknown as Immer
--- a/src/middleware/persist.ts
+++ b/src/middleware/persist.ts
@ -1,8 +1,9 @@
-import { GetState, SetState, State, StoreApi } from '../vanilla'
-
-type DeepPartial<T> = {
-  [P in keyof T]?: DeepPartial<T[P]>
-}
+import {
+  State,
+  StateCreator,
+  StoreApi,
+  StoreMutatorIdentifier,
+} from '../vanilla'

 export type StateStorage = {
  getItem: (name: string) => string | null | Promise<string | null>
@ -10,12 +11,9 @@ export type StateStorage = {
  removeItem: (name: string) => void | Promise<void>
 }

-type StorageValue<S> = { state: DeepPartial<S>; version?: number }
+type StorageValue<S> = { state: S; version?: number }

-export type PersistOptions<
-  S,
-  PersistedState extends Partial<S> = Partial<S>
-> = {
+export type PersistOptions<S, PersistedState = S> = {
  /** Name of the storage (must be unique) */
  name: string
  /**
@ -48,13 +46,15 @@ export type PersistOptions<
   *
   * @params state The state's value
   */
-  partialize?: (state: S) => DeepPartial<S>
+  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?: Error) => void) | void
+  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.
@ -64,37 +64,17 @@ export type PersistOptions<
   * 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>
+  migrate?: (persistedState: unknown, version: number) => S | Promise<S>
  /**
   * 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: any, currentState: S) => S
+  merge?: (persistedState: unknown, currentState: S) => S
 }

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

-/**
- * @deprecated Use `Mutate<StoreApi<T>, [["zustand/persist", Partial<T>]]>`.
- * See tests/middlewaresTypes.test.tsx for usage with multiple middlewares.
- */
-export type StoreApiWithPersist<S extends State> = StoreApi<S> &
-  StorePersist<S, Partial<S>>
-
-declare module '../vanilla' {
-  interface StoreMutators<S, A> {
-    'zustand/persist': WithPersist<S, A>
-  }
-}
-
-type Write<T extends object, U extends object> = Omit<T, keyof U> & U
-type Cast<T, U> = T extends U ? T : U
-
-type WithPersist<S, A> = S extends { getState: () => infer T }
-  ? Write<S, StorePersist<Cast<T, State>, A>>
-  : never
-
-interface StorePersist<S extends State, Ps> {
+type StorePersist<S extends State, Ps> = {
  persist: {
    setOptions: (options: Partial<PersistOptions<S, Ps>>) => void
    clearStorage: () => void
@ -144,191 +124,210 @@ const toThenable =
    }
  }

-export const persist =
-  <
-    S extends State,
-    CustomSetState extends SetState<S> = SetState<S>,
-    CustomGetState extends GetState<S> = GetState<S>,
-    CustomStoreApi extends StoreApi<S> = StoreApi<S>
-  >(
-    config: (
-      set: CustomSetState,
-      get: CustomGetState,
-      api: CustomStoreApi
-    ) => S,
-    baseOptions: PersistOptions<S>
-  ) =>
-  (
-    set: CustomSetState,
-    get: CustomGetState,
-    api: CustomStoreApi & StoreApiWithPersist<S>
-  ): S => {
-    let options = {
-      getStorage: () => localStorage,
-      serialize: JSON.stringify as (state: StorageValue<S>) => string,
-      deserialize: JSON.parse as (str: string) => StorageValue<Partial<S>>,
-      partialize: (state: S) => state,
-      version: 0,
-      merge: (persistedState: any, currentState: S) => ({
-        ...currentState,
-        ...persistedState,
-      }),
-      ...baseOptions,
-    }
+const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
+  type S = ReturnType<typeof config>
+  let options = {
+    getStorage: () => localStorage,
+    serialize: JSON.stringify as (state: StorageValue<S>) => string,
+    deserialize: JSON.parse as (str: string) => StorageValue<S>,
+    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: StateStorage | undefined
+  let hasHydrated = false
+  const hydrationListeners = new Set<PersistListener<S>>()
+  const finishHydrationListeners = new Set<PersistListener<S>>()
+  let storage: StateStorage | undefined

-    try {
-      storage = options.getStorage()
-    } catch (e) {
-      // prevent error if the storage is not defined (e.g. when server side rendering a page)
-    }
+  try {
+    storage = options.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(
-            `[zustand persist middleware] Unable to update item '${options.name}', the given storage is currently unavailable.`
-          )
-          set(...args)
-        }) as CustomSetState,
-        get,
-        api
-      )
-    }
-
-    const thenableSerialize = toThenable(options.serialize)
-
-    const setItem = (): Thenable<void> => {
-      const state = options.partialize({ ...get() })
-
-      let errorInSync: Error | undefined
-      const thenable = thenableSerialize({ state, version: options.version })
-        .then((serializedValue) =>
-          (storage as StateStorage).setItem(options.name, serializedValue)
+  if (!storage) {
+    return config(
+      (...args) => {
+        console.warn(
+          `[zustand persist middleware] Unable to update item '${options.name}', the given storage is currently unavailable.`
        )
-        .catch((e) => {
-          errorInSync = e
-        })
-      if (errorInSync) {
-        throw errorInSync
-      }
-      return thenable
-    }
-
-    const savedSetState = api.setState
-
-    api.setState = (state, replace) => {
-      savedSetState(state, replace)
-      void setItem()
-    }
-
-    const configResult = config(
-      ((...args) => {
        set(...args)
-        void setItem()
-      }) as CustomSetState,
+      },
      get,
      api
    )
-
-    // 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
-
-      hasHydrated = false
-      hydrationListeners.forEach((cb) => cb(get()))
-
-      const postRehydrationCallback =
-        options.onRehydrateStorage?.(get()) || undefined
-
-      // bind is used to avoid `TypeError: Illegal invocation` error
-      return toThenable(storage.getItem.bind(storage))(options.name)
-        .then((storageValue) => {
-          if (storageValue) {
-            return options.deserialize(storageValue)
-          }
-        })
-        .then((deserializedStorageValue) => {
-          if (deserializedStorageValue) {
-            if (
-              typeof deserializedStorageValue.version === 'number' &&
-              deserializedStorageValue.version !== options.version
-            ) {
-              if (options.migrate) {
-                return options.migrate(
-                  deserializedStorageValue.state,
-                  deserializedStorageValue.version
-                )
-              }
-              console.error(
-                `State loaded from storage couldn't be migrated since no migrate function was provided`
-              )
-            } else {
-              return deserializedStorageValue.state
-            }
-          }
-        })
-        .then((migratedState) => {
-          stateFromStorage = options.merge(
-            migratedState as S,
-            get() ?? configResult
-          )
-
-          set(stateFromStorage as S, true)
-          return setItem()
-        })
-        .then(() => {
-          postRehydrationCallback?.(stateFromStorage, undefined)
-          hasHydrated = true
-          finishHydrationListeners.forEach((cb) => cb(stateFromStorage as S))
-        })
-        .catch((e: Error) => {
-          postRehydrationCallback?.(undefined, e)
-        })
-    }
-
-    api.persist = {
-      setOptions: (newOptions) => {
-        options = {
-          ...options,
-          ...newOptions,
-        }
-
-        if (newOptions.getStorage) {
-          storage = newOptions.getStorage()
-        }
-      },
-      clearStorage: () => {
-        storage?.removeItem(options.name)
-      },
-      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)
-        }
-      },
-    }
-
-    hydrate()
-
-    return stateFromStorage || configResult
  }
+
+  const thenableSerialize = toThenable(options.serialize)
+
+  const setItem = (): Thenable<void> => {
+    const state = options.partialize({ ...get() })
+
+    let errorInSync: Error | undefined
+    const thenable = thenableSerialize({ state, version: options.version })
+      .then((serializedValue) =>
+        (storage as StateStorage).setItem(options.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()
+  }
+
+  const configResult = config(
+    (...args) => {
+      set(...args)
+      void setItem()
+    },
+    get,
+    api
+  )
+
+  // 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
+
+    hasHydrated = false
+    hydrationListeners.forEach((cb) => cb(get()))
+
+    const postRehydrationCallback =
+      options.onRehydrateStorage?.(get()) || undefined
+
+    // bind is used to avoid `TypeError: Illegal invocation` error
+    return toThenable(storage.getItem.bind(storage))(options.name)
+      .then((storageValue) => {
+        if (storageValue) {
+          return options.deserialize(storageValue)
+        }
+      })
+      .then((deserializedStorageValue) => {
+        if (deserializedStorageValue) {
+          if (
+            typeof deserializedStorageValue.version === 'number' &&
+            deserializedStorageValue.version !== options.version
+          ) {
+            if (options.migrate) {
+              return options.migrate(
+                deserializedStorageValue.state,
+                deserializedStorageValue.version
+              )
+            }
+            console.error(
+              `State loaded from storage couldn't be migrated since no migrate function was provided`
+            )
+          } else {
+            return deserializedStorageValue.state
+          }
+        }
+      })
+      .then((migratedState) => {
+        stateFromStorage = options.merge(
+          migratedState as S,
+          get() ?? configResult
+        )
+
+        set(stateFromStorage as S, true)
+        return setItem()
+      })
+      .then(() => {
+        postRehydrationCallback?.(stateFromStorage, undefined)
+        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.getStorage) {
+        storage = newOptions.getStorage()
+      }
+    },
+    clearStorage: () => {
+      storage?.removeItem(options.name)
+    },
+    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)
+      }
+    },
+  }
+
+  hydrate()
+
+  return stateFromStorage || configResult
+}
+
+type Persist = <
+  T extends State,
+  Mps extends [StoreMutatorIdentifier, unknown][] = [],
+  Mcs extends [StoreMutatorIdentifier, unknown][] = [],
+  U = Partial<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 extends object, U extends object> = Omit<T, keyof U> & U
+type Cast<T, U> = T extends U ? T : U
+
+type WithPersist<S, A> = S extends { getState: () => infer T }
+  ? Write<S, StorePersist<Cast<T, State>, A>>
+  : never
+
+type PersistImpl = <T extends State>(
+  storeInitializer: PopArgument<StateCreator<T, [], []>>,
+  options: PersistOptions<T, T>
+) => PopArgument<StateCreator<T, [], []>>
+
+type PopArgument<T extends (...a: never[]) => unknown> = T extends (
+  ...a: [...infer A, infer _]
+) => infer R
+  ? (...a: A) => R
+  : never
+
+export const persist = persistImpl as unknown as Persist
--- a/src/middleware/redux.ts
+++ b/src/middleware/redux.ts
@ -1,31 +1,15 @@
-import { GetState, SetState, State, StoreApi } from '../vanilla'
+import { State, StateCreator, StoreMutatorIdentifier } from '../vanilla'
 import { NamedSet } from './devtools'

-type DevtoolsType = {
-  prefix: string
-  subscribe: (dispatch: any) => () => void
-  unsubscribe: () => void
-  send: (action: string, state: any) => void
-  init: (state: any) => void
-  error: (payload: any) => void
+type Write<T extends object, U extends object> = Omit<T, keyof U> & U
+type Cast<T, U> = T extends U ? T : U
+
+interface Action {
+  type: unknown
 }

-/**
- * @deprecated Use `Mutate<StoreApi<T & { dispatch: (a: A) => A }>, [["zustand/redux", A]]>`.
- * See tests/middlewaresTypes.test.tsx for usage with multiple middlewares.
- */
-export type StoreApiWithRedux<
-  T extends State,
-  A extends { type: unknown }
-> = StoreApi<T & { dispatch: (a: A) => A }> & {
-  dispatch: (a: A) => A
-  dispatchFromDevtools: boolean
-}
-
-declare module '../vanilla' {
-  interface StoreMutators<S, A> {
-    'zustand/redux': WithRedux<S, A>
-  }
+interface ReduxState<A extends Action> {
+  dispatch: StoreRedux<A>['dispatch']
 }

 interface StoreRedux<A extends Action> {
@ -33,30 +17,43 @@ interface StoreRedux<A extends Action> {
  dispatchFromDevtools: true
 }

-interface Action {
-  type: unknown
-}
-
-type Write<T extends object, U extends object> = Omit<T, keyof U> & U
-type Cast<T, U> = T extends U ? T : U
-
 type WithRedux<S, A> = Write<Cast<S, object>, StoreRedux<Cast<A, Action>>>

-export const redux =
-  <S extends State, A extends { type: unknown }>(
-    reducer: (state: S, action: A) => S,
-    initial: S
-  ) =>
-  (
-    set: SetState<S & { dispatch: (a: A) => A }>,
-    get: GetState<S & { dispatch: (a: A) => A }>,
-    api: StoreApiWithRedux<S, A> & { devtools?: DevtoolsType }
-  ): S & { dispatch: (a: A) => A } => {
-    api.dispatch = (action: A) => {
-      ;(set as NamedSet<S>)((state: S) => reducer(state, action), false, action)
-      return action
-    }
-    api.dispatchFromDevtools = true
+type Redux = <
+  T extends State,
+  A extends Action,
+  Cms extends [StoreMutatorIdentifier, unknown][] = []
+>(
+  reducer: (state: T, action: A) => T,
+  initialState: T
+) => StateCreator<Write<T, ReduxState<A>>, Cms, [['zustand/redux', A]]>

-    return { dispatch: (...a) => api.dispatch(...a), ...initial }
+declare module '../vanilla' {
+  interface StoreMutators<S, A> {
+    'zustand/redux': WithRedux<S, A>
  }
+}
+
+type PopArgument<T extends (...a: never[]) => unknown> = T extends (
+  ...a: [...infer A, infer _]
+) => infer R
+  ? (...a: A) => R
+  : never
+
+type ReduxImpl = <T extends State, A extends Action>(
+  reducer: (state: T, action: A) => T,
+  initialState: T
+) => PopArgument<StateCreator<T & ReduxState<A>, [], []>>
+
+const reduxImpl: ReduxImpl = (reducer, initial) => (set, get, api) => {
+  type S = typeof initial
+  type A = Parameters<typeof reducer>[1]
+  ;(api as any).dispatch = (action: A) => {
+    ;(set as NamedSet<S>)((state: S) => reducer(state, action), false, action)
+    return action
+  }
+  ;(api as any).dispatchFromDevtools = true
+
+  return { dispatch: (...a) => (api as any).dispatch(...a), ...initial }
+}
+export const redux = reduxImpl as Redux
--- a/src/middleware/subscribeWithSelector.ts
+++ b/src/middleware/subscribeWithSelector.ts
@ -1,15 +1,30 @@
 import {
-  EqualityChecker,
-  GetState,
-  SetState,
  State,
+  StateCreator,
  StateListener,
-  StateSelector,
-  StateSliceListener,
-  StoreApi,
+  StoreMutatorIdentifier,
  Subscribe,
 } from '../vanilla'

+type SubscribeWithSelector = <
+  T extends State,
+  Mps extends [StoreMutatorIdentifier, unknown][] = [],
+  Mcs extends [StoreMutatorIdentifier, unknown][] = []
+>(
+  initializer: StateCreator<
+    T,
+    [...Mps, ['zustand/subscribeWithSelector', never]],
+    Mcs
+  >
+) => StateCreator<T, Mps, [['zustand/subscribeWithSelector', never], ...Mcs]>
+
+type Write<T extends object, U extends object> = Omit<T, keyof U> & U
+type Cast<T, U> = T extends U ? T : U
+
+type WithSelectorSubscribe<S> = S extends { getState: () => infer T }
+  ? Write<S, StoreSubscribeWithSelector<Cast<T, State>>>
+  : never
+
 declare module '../vanilla' {
  // eslint-disable-next-line @typescript-eslint/no-unused-vars
  interface StoreMutators<S, A> {
@ -17,10 +32,6 @@ declare module '../vanilla' {
  }
 }

-type WithSelectorSubscribe<S> = S extends { getState: () => infer T }
-  ? Omit<S, 'subscribe'> & StoreSubscribeWithSelector<Extract<T, State>>
-  : never
-
 interface StoreSubscribeWithSelector<T extends State> {
  subscribe: {
    (listener: (selectedState: T, previousSelectedState: T) => void): () => void
@ -35,38 +46,19 @@ interface StoreSubscribeWithSelector<T extends State> {
  }
 }

-/**
- * @deprecated Use `Mutate<StoreApi<T>, [["zustand/subscribeWithSelector", never]]>`.
- * See tests/middlewaresTypes.test.tsx for usage with multiple middlewares.
- */
-export type StoreApiWithSubscribeWithSelector<T extends State> = StoreApi<T> & {
-  subscribe: {
-    (listener: StateListener<T>): () => void
-    <StateSlice>(
-      selector: StateSelector<T, StateSlice>,
-      listener: StateSliceListener<StateSlice>,
-      options?: {
-        equalityFn?: EqualityChecker<StateSlice>
-        fireImmediately?: boolean
-      }
-    ): () => void
-  }
-}
+type SubscribeWithSelectorImpl = <T extends State>(
+  storeInitializer: PopArgument<StateCreator<T, [], []>>
+) => PopArgument<StateCreator<T, [], []>>

-export const subscribeWithSelector =
-  <
-    S extends State,
-    CustomSetState extends SetState<S> = SetState<S>,
-    CustomGetState extends GetState<S> = GetState<S>,
-    CustomStoreApi extends StoreApi<S> = StoreApi<S>
-  >(
-    fn: (set: CustomSetState, get: CustomGetState, api: CustomStoreApi) => S
-  ) =>
-  (
-    set: CustomSetState,
-    get: CustomGetState,
-    api: CustomStoreApi & StoreApiWithSubscribeWithSelector<S>
-  ): S => {
+type PopArgument<T extends (...a: never[]) => unknown> = T extends (
+  ...a: [...infer A, infer _]
+) => infer R
+  ? (...a: A) => R
+  : never
+
+const subscribeWithSelectorImpl: SubscribeWithSelectorImpl =
+  (fn) => (set, get, api) => {
+    type S = ReturnType<typeof fn>
    const origSubscribe = api.subscribe as Subscribe<S>
    api.subscribe = ((selector: any, optListener: any, options: any) => {
      let listener: StateListener<S> = selector // if no selector
@ -89,3 +81,5 @@ export const subscribeWithSelector =
    const initialState = fn(set, get, api)
    return initialState
  }
+export const subscribeWithSelector =
+  subscribeWithSelectorImpl as unknown as SubscribeWithSelector
--- a/src/react.ts
+++ b/src/react.ts
@ -2,18 +2,18 @@ import { useDebugValue } from 'react'
 import { useSyncExternalStoreWithSelector } from 'use-sync-external-store/shim/with-selector'
 import createStore, {
  EqualityChecker,
-  GetState,
-  SetState,
+  Mutate,
  State,
  StateCreator,
  StateSelector,
  StoreApi,
+  StoreMutatorIdentifier,
 } from './vanilla'

-export function useStore<T extends State>(api: StoreApi<T>): T
-export function useStore<T extends State, U>(
-  api: StoreApi<T>,
-  selector: StateSelector<T, U>,
+export function useStore<S extends StoreApi<State>>(api: S): ExtractState<S>
+export function useStore<S extends StoreApi<State>, U>(
+  api: S,
+  selector: StateSelector<ExtractState<S>, U>,
  equalityFn?: EqualityChecker<U>
 ): U
 export function useStore<TState extends State, StateSlice>(
@ -33,42 +33,28 @@ export function useStore<TState extends State, StateSlice>(
  return slice
 }

-export type UseBoundStore<
-  T extends State,
-  CustomStoreApi extends StoreApi<T> = StoreApi<T>
-> = {
-  (): T
-  <U>(selector: StateSelector<T, U>, equalityFn?: EqualityChecker<U>): U
-} & CustomStoreApi
+type ExtractState<S> = S extends { getState: () => infer T } ? T : never

-function create<
-  TState extends State,
-  CustomSetState,
-  CustomGetState,
-  CustomStoreApi extends StoreApi<TState>
->(
-  createState:
-    | StateCreator<TState, CustomSetState, CustomGetState, CustomStoreApi>
-    | CustomStoreApi
-): UseBoundStore<TState, CustomStoreApi>
+export type UseBoundStore<S extends StoreApi<State>> = {
+  (): ExtractState<S>
+  <U>(
+    selector: StateSelector<ExtractState<S>, U>,
+    equals?: EqualityChecker<U>
+  ): U
+} & S

-function create<TState extends State>(
-  createState:
-    | StateCreator<TState, SetState<TState>, GetState<TState>, any>
-    | StoreApi<TState>
-): UseBoundStore<TState, StoreApi<TState>>
+type Create = {
+  <T extends State, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
+    initializer: StateCreator<T, [], Mos>
+  ): UseBoundStore<Mutate<StoreApi<T>, Mos>>
+  <T extends State>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
+    initializer: StateCreator<T, [], Mos>
+  ) => UseBoundStore<Mutate<StoreApi<T>, Mos>>
+  <S extends StoreApi<State>>(store: S): UseBoundStore<S>
+}

-function create<
-  TState extends State,
-  CustomSetState,
-  CustomGetState,
-  CustomStoreApi extends StoreApi<TState>
->(
-  createState:
-    | StateCreator<TState, CustomSetState, CustomGetState, CustomStoreApi>
-    | CustomStoreApi
-): UseBoundStore<TState, CustomStoreApi> {
-  const api: CustomStoreApi =
+const createImpl = <T extends State>(createState: StateCreator<T, [], []>) => {
+  const api =
    typeof createState === 'function' ? createStore(createState) : createState

  const useBoundStore: any = (selector?: any, equalityFn?: any) =>
@ -79,4 +65,8 @@ function create<
  return useBoundStore
 }

+const create = (<T extends State>(
+  createState: StateCreator<T, [], []> | undefined
+) => (createState ? createImpl(createState) : createImpl)) as Create
+
 export default create
--- a/src/vanilla.ts
+++ b/src/vanilla.ts
@ -1,21 +1,7 @@
 export type State = object
-// types inspired by setState from React, see:
-// https://github.com/DefinitelyTyped/DefinitelyTyped/blob/6c49e45842358ba59a508e13130791989911430d/types/react/v16/index.d.ts#L489-L495
-
-/**
- * @deprecated Use the builtin `Partial<T>` instead of `PartialState<T>`.
- * Additionally turn on `--exactOptionalPropertyTypes` tsc flag.
- * `PartialState` will be removed in next major
- */
-export type PartialState<
-  T extends State,
-  K1 extends keyof T = keyof T,
-  K2 extends keyof T = K1,
-  K3 extends keyof T = K2,
-  K4 extends keyof T = K3
-> =
-  | (Pick<T, K1> | Pick<T, K2> | Pick<T, K3> | Pick<T, K4> | T)
-  | ((state: T) => Pick<T, K1> | Pick<T, K2> | Pick<T, K3> | Pick<T, K4> | T)
+export type PartialState<T extends State> =
+  | Partial<T>
+  | ((state: T) => Partial<T>)
 export type StateSelector<T extends State, U> = (state: T) => U
 export type EqualityChecker<T> = (state: T, newState: T) => boolean
 export type StateListener<T> = (state: T, previousState: T) => void
@ -28,16 +14,11 @@ export type Subscribe<T extends State> = {
 }

 export type SetState<T extends State> = {
-  <
-    K1 extends keyof T,
-    K2 extends keyof T = K1,
-    K3 extends keyof T = K2,
-    K4 extends keyof T = K3
-  >(
-    partial: PartialState<T, K1, K2, K3, K4>,
-    replace?: boolean
+  _(
+    partial: T | Partial<T> | ((state: T) => T | Partial<T>),
+    replace?: boolean | undefined
  ): void
-}
+}['_']
 export type GetState<T extends State> = () => T
 export type Destroy = () => void
 export type StoreApi<T extends State> = {
@ -46,44 +27,56 @@ export type StoreApi<T extends State> = {
  subscribe: Subscribe<T>
  destroy: Destroy
 }
+
 export type StateCreator<
  T extends State,
-  CustomSetState = SetState<T>,
-  CustomGetState = GetState<T>,
-  CustomStoreApi extends StoreApi<T> = StoreApi<T>
-> = (set: CustomSetState, get: CustomGetState, api: CustomStoreApi) => T
+  Mis extends [StoreMutatorIdentifier, unknown][] = [],
+  Mos extends [StoreMutatorIdentifier, unknown][] = [],
+  U = T
+> = ((
+  setState: Get<Mutate<StoreApi<T>, Mis>, 'setState', undefined>,
+  getState: Get<Mutate<StoreApi<T>, Mis>, 'getState', undefined>,
+  store: Mutate<StoreApi<T>, Mis>,
+  $$storeMutations: Mis
+) => U) & { $$storeMutators?: Mos }

-function createStore<
-  TState extends State,
-  CustomSetState,
-  CustomGetState,
-  CustomStoreApi extends StoreApi<TState>
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+export interface StoreMutators<S, A> {}
+export type StoreMutatorIdentifier = keyof StoreMutators<unknown, unknown>
+
+export type Mutate<S, Ms> = Ms extends []
+  ? S
+  : Ms extends [[infer Mi, infer Ma], ...infer Mrs]
+  ? Mutate<StoreMutators<S, Ma>[Mi & StoreMutatorIdentifier], Mrs>
+  : never
+
+type Get<T, K, F = never> = K extends keyof T ? T[K] : F
+
+type CreateStore = {
+  <T extends State, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
+    initializer: StateCreator<T, [], Mos>
+  ): Mutate<StoreApi<T>, Mos>
+
+  <T extends State>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
+    initializer: StateCreator<T, [], Mos>
+  ) => Mutate<StoreApi<T>, Mos>
+}
+
+type CreateStoreImpl = <
+  T extends State,
+  Mos extends [StoreMutatorIdentifier, unknown][] = []
 >(
-  createState: StateCreator<
-    TState,
-    CustomSetState,
-    CustomGetState,
-    CustomStoreApi
-  >
-): CustomStoreApi
+  initializer: StateCreator<T, [], Mos>
+) => Mutate<StoreApi<T>, Mos>

-function createStore<TState extends State>(
-  createState: StateCreator<TState, SetState<TState>, GetState<TState>, any>
-): StoreApi<TState>
+type PopArgument<T extends (...a: never[]) => unknown> = T extends (
+  ...a: [...infer A, infer _]
+) => infer R
+  ? (...a: A) => R
+  : never

-function createStore<
-  TState extends State,
-  CustomSetState,
-  CustomGetState,
-  CustomStoreApi extends StoreApi<TState>
->(
-  createState: StateCreator<
-    TState,
-    CustomSetState,
-    CustomGetState,
-    CustomStoreApi
-  >
-): CustomStoreApi {
+const createStoreImpl: CreateStoreImpl = (createState) => {
+  type TState = ReturnType<typeof createState>
  let state: TState
  const listeners: Set<StateListener<TState>> = new Set()

@ -113,22 +106,15 @@ function createStore<

  const destroy: Destroy = () => listeners.clear()
  const api = { setState, getState, subscribe, destroy }
-  state = createState(
-    setState as unknown as CustomSetState,
-    getState as unknown as CustomGetState,
-    api as unknown as CustomStoreApi
+  state = (createState as PopArgument<typeof createState>)(
+    setState,
+    getState,
+    api
  )
-  return api as unknown as CustomStoreApi
+  return api as any
 }

+const createStore = ((createState) =>
+  createState ? createStoreImpl(createState) : createStoreImpl) as CreateStore
+
 export default createStore
-
-// eslint-disable-next-line @typescript-eslint/no-unused-vars
-export interface StoreMutators<S, A> {}
-export type StoreMutatorIdentifier = keyof StoreMutators<unknown, unknown>
-
-export type Mutate<S, Ms> = Ms extends []
-  ? S
-  : Ms extends [[infer Mi, infer Ma], ...infer Mrs]
-  ? Mutate<StoreMutators<S, Ma>[Mi & StoreMutatorIdentifier], Mrs>
-  : never
--- a/tests/context.test.tsx
+++ b/tests/context.test.tsx
@ -6,12 +6,9 @@ import {
  useState,
 } from 'react'
 import { render } from '@testing-library/react'
-import create, { GetState, SetState } from 'zustand'
+import create, { StoreApi } from 'zustand'
 import createContext from 'zustand/context'
-import {
-  StoreApiWithSubscribeWithSelector,
-  subscribeWithSelector,
-} from 'zustand/middleware'
+import { subscribeWithSelector } from 'zustand/middleware'

 const consoleError = console.error
 afterEach(() => {
@ -24,7 +21,7 @@ type CounterState = {
 }

 it('creates and uses context store', async () => {
-  const { Provider, useStore } = createContext<CounterState>()
+  const { Provider, useStore } = createContext<StoreApi<CounterState>>()

  const createStore = () =>
    create<CounterState>((set) => ({
@ -48,7 +45,7 @@ it('creates and uses context store', async () => {
 })

 it('uses context store with selectors', async () => {
-  const { Provider, useStore } = createContext<CounterState>()
+  const { Provider, useStore } = createContext<StoreApi<CounterState>>()

  const createStore = () =>
    create<CounterState>((set) => ({
@ -74,12 +71,7 @@ it('uses context store with selectors', async () => {

 it('uses context store api', async () => {
  const createStore = () =>
-    create<
-      CounterState,
-      SetState<CounterState>,
-      GetState<CounterState>,
-      StoreApiWithSubscribeWithSelector<CounterState>
-    >(
+    create<CounterState>()(
      subscribeWithSelector((set) => ({
        count: 0,
        inc: () => set((state) => ({ count: state.count + 1 })),
@ -87,7 +79,7 @@ it('uses context store api', async () => {
    )

  type CustomStore = ReturnType<typeof createStore>
-  const { Provider, useStoreApi } = createContext<CounterState, CustomStore>()
+  const { Provider, useStoreApi } = createContext<CustomStore>()

  function Counter() {
    const storeApi = useStoreApi()
@ -140,7 +132,7 @@ it('throws error when not using provider', async () => {
    }
  }

-  const { useStore } = createContext<CounterState>()
+  const { useStore } = createContext<StoreApi<CounterState>>()
  function Component() {
    useStore()
    return <div>no error</div>
@ -155,7 +147,7 @@ it('throws error when not using provider', async () => {
 })

 it('useCallback with useStore infers types correctly', async () => {
-  const { useStore } = createContext<CounterState>()
+  const { useStore } = createContext<StoreApi<CounterState>>()
  function _Counter() {
    const _x = useStore(useCallback((state) => state.count, []))
    expectAreTypesEqual<typeof _x, number>().toBe(true)
--- a/tests/devtools.test.tsx
+++ b/tests/devtools.test.tsx
@ -86,10 +86,10 @@ describe('When state changes...', () => {
      { type: 'testSetStateName' },
      { count: 10, foo: 'bar' }
    )
-    api.setState({ count: 5 }, true)
+    api.setState({ count: 5, foo: 'baz' }, true)
    expect(extension.send).toHaveBeenLastCalledWith(
      { type: 'anonymous' },
-      { count: 5 }
+      { count: 5, foo: 'baz' }
    )
  })
 })
@ -439,7 +439,12 @@ describe('when it receives an message of type...', () => {
 })

 describe('with redux middleware', () => {
-  let api: StoreApi<{ count: number }>
+  let api: StoreApi<{
+    count: number
+    dispatch: (
+      action: { type: 'INCREMENT' } | { type: 'DECREMENT' }
+    ) => { type: 'INCREMENT' } | { type: 'DECREMENT' }
+  }>

  it('works as expected', () => {
    api = create(
--- a/tests/middlewareTypes.test.tsx
+++ b/tests/middlewareTypes.test.tsx
@ -1,48 +1,13 @@
-import { produce } from 'immer'
-import type { Draft } from 'immer'
-import create, {
-  GetState,
-  Mutate,
-  SetState,
-  State,
-  StateCreator,
-  StoreApi,
-} from 'zustand'
+import create, { State, StoreApi } from 'zustand'
 import {
-  PersistOptions,
  combine,
  devtools,
+  immer,
  persist,
  redux,
  subscribeWithSelector,
 } from 'zustand/middleware'
-
-const immer =
-  <
-    T extends State,
-    CustomSetState extends SetState<T>,
-    CustomGetState extends GetState<T>,
-    CustomStoreApi extends StoreApi<T>
-  >(
-    config: StateCreator<
-      T,
-      (partial: ((draft: Draft<T>) => void) | T, replace?: boolean) => void,
-      CustomGetState,
-      CustomStoreApi
-    >
-  ): StateCreator<T, CustomSetState, CustomGetState, CustomStoreApi> =>
-  (set, get, api) =>
-    config(
-      (partial, replace) => {
-        const nextState =
-          typeof partial === 'function'
-            ? produce(partial as (state: Draft<T>) => T)
-            : (partial as T)
-        return set(nextState, replace)
-      },
-      get,
-      api
-    )
+import createVanilla from 'zustand/vanilla'

 type CounterState = {
  count: number
@ -78,7 +43,7 @@ describe('counter state spec (single middleware)', () => {
  })

  it('immer', () => {
-    const useStore = create<CounterState>(
+    const useStore = create<CounterState>()(
      immer((set, get) => ({
        count: 0,
        inc: () =>
@ -97,6 +62,10 @@ describe('counter state spec (single middleware)', () => {
      return <></>
    }
    TestComponent
+
+    const _testSubtyping: StoreApi<State> = createVanilla(
+      immer(() => ({ count: 0 }))
+    )
  })

  it('redux', () => {
@ -121,16 +90,15 @@ describe('counter state spec (single middleware)', () => {
      return <></>
    }
    TestComponent
+
+    const _testSubtyping: StoreApi<State> = createVanilla(
+      redux((x) => x, { count: 0 })
+    )
  })

  it('devtools', () => {
    __DEV__ = false
-    const useStore = create<
-      CounterState,
-      SetState<CounterState>,
-      GetState<CounterState>,
-      Mutate<StoreApi<CounterState>, [['zustand/devtools', never]]>
-    >(
+    const useStore = create<CounterState>()(
      devtools(
        (set, get) => ({
          count: 0,
@ -150,15 +118,14 @@ describe('counter state spec (single middleware)', () => {
      return <></>
    }
    TestComponent
+
+    const _testSubtyping: StoreApi<State> = createVanilla(
+      devtools(() => ({ count: 0 }))
+    )
  })

  it('subscribeWithSelector', () => {
-    const useStore = create<
-      CounterState,
-      SetState<CounterState>,
-      GetState<CounterState>,
-      Mutate<StoreApi<CounterState>, [['zustand/subscribeWithSelector', never]]>
-    >(
+    const useStore = create<CounterState>()(
      subscribeWithSelector((set, get) => ({
        count: 1,
        inc: () => set({ count: get().count + 1 }, false),
@ -178,6 +145,10 @@ describe('counter state spec (single middleware)', () => {
      return <></>
    }
    TestComponent
+
+    const _testSubtyping: StoreApi<State> = createVanilla(
+      subscribeWithSelector(() => ({ count: 0 }))
+    )
  })

  it('combine', () => {
@ -196,18 +167,14 @@ describe('counter state spec (single middleware)', () => {
      return <></>
    }
    TestComponent
+
+    const _testSubtyping: StoreApi<State> = createVanilla(
+      combine({ count: 0 }, () => ({}))
+    )
  })

  it('persist', () => {
-    const useStore = create<
-      CounterState,
-      SetState<CounterState>,
-      GetState<CounterState>,
-      Mutate<
-        StoreApi<CounterState>,
-        [['zustand/persist', Partial<CounterState>]]
-      >
-    >(
+    const useStore = create<CounterState>()(
      persist(
        (set, get) => ({
          count: 1,
@ -227,10 +194,41 @@ describe('counter state spec (single middleware)', () => {
      return <></>
    }
    TestComponent
+
+    const _testSubtyping: StoreApi<State> = createVanilla(
+      persist(() => ({ count: 0 }))
+    )
+  })
+
+  it('persist with partialize', () => {
+    const useStore = create<CounterState>()(
+      persist(
+        (set, get) => ({
+          count: 1,
+          inc: () => set({ count: get().count + 1 }, false),
+        }),
+        { name: 'prefix', partialize: (s) => s.count }
+      )
+    )
+    const TestComponent = () => {
+      useStore((s) => s.count) * 2
+      useStore((s) => s.inc)()
+      useStore().count * 2
+      useStore().inc()
+      useStore.getState().count * 2
+      useStore.getState().inc()
+      useStore.persist.hasHydrated()
+      useStore.persist.setOptions({
+        // @ts-expect-error to test if the partialized state is inferred as number
+        partialize: () => 'not-a-number',
+      })
+      return <></>
+    }
+    TestComponent
  })

  it('persist without custom api (#638)', () => {
-    const useStore = create<CounterState>(
+    const useStore = create<CounterState>()(
      persist(
        (set, get) => ({
          count: 1,
@ -263,12 +261,7 @@ describe('counter state spec (double middleware)', () => {

  it('devtools & immer', () => {
    __DEV__ = false
-    const useStore = create<
-      CounterState,
-      SetState<CounterState>,
-      GetState<CounterState>,
-      Mutate<StoreApi<CounterState>, [['zustand/devtools', never]]>
-    >(
+    const useStore = create<CounterState>()(
      devtools(
        immer((set, get) => ({
          count: 0,
@ -297,8 +290,8 @@ describe('counter state spec (double middleware)', () => {
    __DEV__ = false
    const useStore = create(
      devtools(
-        redux<{ count: number }, { type: 'INC' }>(
-          (state, action) => {
+        redux(
+          (state, action: { type: 'INC' }) => {
            switch (action.type) {
              case 'INC':
                return { ...state, count: state.count + 1 }
@ -371,15 +364,7 @@ describe('counter state spec (double middleware)', () => {

  it('devtools & subscribeWithSelector', () => {
    __DEV__ = false
-    const useStore = create<
-      CounterState,
-      SetState<CounterState>,
-      GetState<CounterState>,
-      Mutate<
-        StoreApi<CounterState>,
-        [['zustand/subscribeWithSelector', never], ['zustand/devtools', never]]
-      >
-    >(
+    const useStore = create<CounterState>()(
      devtools(
        subscribeWithSelector((set, get) => ({
          count: 1,
@ -407,18 +392,7 @@ describe('counter state spec (double middleware)', () => {

  it('devtools & persist', () => {
    __DEV__ = false
-    const useStore = create<
-      CounterState,
-      SetState<CounterState>,
-      GetState<CounterState>,
-      Mutate<
-        StoreApi<CounterState>,
-        [
-          ['zustand/persist', Partial<CounterState>],
-          ['zustand/devtools', never]
-        ]
-      >
-    >(
+    const useStore = create<CounterState>()(
      devtools(
        persist(
          (set, get) => ({
@ -456,18 +430,7 @@ describe('counter state spec (triple middleware)', () => {

  it('devtools & persist & immer', () => {
    __DEV__ = false
-    const useStore = create<
-      CounterState,
-      SetState<CounterState>,
-      GetState<CounterState>,
-      Mutate<
-        StoreApi<CounterState>,
-        [
-          ['zustand/persist', Partial<CounterState>],
-          ['zustand/devtools', never]
-        ]
-      >
-    >(
+    const useStore = create<CounterState>()(
      devtools(
        persist(
          immer((set, get) => ({
@ -527,19 +490,7 @@ describe('counter state spec (triple middleware)', () => {

  it('devtools & subscribeWithSelector & persist', () => {
    __DEV__ = false
-    const useStore = create<
-      CounterState,
-      SetState<CounterState>,
-      GetState<CounterState>,
-      Mutate<
-        StoreApi<CounterState>,
-        [
-          ['zustand/subscribeWithSelector', never],
-          ['zustand/persist', Partial<CounterState>],
-          ['zustand/devtools', never]
-        ]
-      >
-    >(
+    const useStore = create<CounterState>()(
      devtools(
        subscribeWithSelector(
          persist(
@ -583,19 +534,7 @@ describe('counter state spec (quadruple middleware)', () => {

  it('devtools & subscribeWithSelector & persist & immer (#616)', () => {
    __DEV__ = false
-    const useStore = create<
-      CounterState,
-      SetState<CounterState>,
-      GetState<CounterState>,
-      Mutate<
-        StoreApi<CounterState>,
-        [
-          ['zustand/subscribeWithSelector', never],
-          ['zustand/persist', Partial<CounterState>],
-          ['zustand/devtools', never]
-        ]
-      >
-    >(
+    const useStore = create<CounterState>()(
      devtools(
        subscribeWithSelector(
          persist(
@ -633,24 +572,9 @@ describe('counter state spec (quadruple middleware)', () => {

 describe('more complex state spec with subscribeWithSelector', () => {
  it('#619, #632', () => {
-    type MyState = {
-      foo: boolean
-    }
    const useStore = create(
      subscribeWithSelector(
-        // NOTE: Adding type annotation to inner middleware works.
-        persist<
-          MyState,
-          SetState<MyState>,
-          GetState<MyState>,
-          Mutate<
-            StoreApi<MyState>,
-            [
-              ['zustand/subscribeWithSelector', never],
-              ['zustand/persist', Partial<MyState>]
-            ]
-          >
-        >(
+        persist(
          () => ({
            foo: true,
          }),
@ -676,12 +600,7 @@ describe('more complex state spec with subscribeWithSelector', () => {
    type MyState = {
      foo: number | null
    }
-    const useStore = create<
-      MyState,
-      SetState<MyState>,
-      GetState<MyState>,
-      Mutate<StoreApi<MyState>, [['zustand/subscribeWithSelector', never]]>
-    >(
+    const useStore = create<MyState>()(
      subscribeWithSelector(
        () =>
          ({
@ -708,13 +627,8 @@ describe('more complex state spec with subscribeWithSelector', () => {
      authenticated: boolean
      authenticate: (username: string, password: string) => Promise<void>
    }
-    // NOTE: This is a simplified middleware type without persist api
-    type MyPersist = (
-      config: StateCreator<MyState>,
-      options: PersistOptions<MyState>
-    ) => StateCreator<MyState>
-    const useStore = create<MyState>(
-      (persist as MyPersist)(
+    const useStore = create<MyState>()(
+      persist(
        (set) => ({
          token: undefined,
          authenticated: false,
@ -737,3 +651,41 @@ describe('more complex state spec with subscribeWithSelector', () => {
    TestComponent
  })
 })
+
+describe('create with explicitly annotated mutators', () => {
+  it('subscribeWithSelector & persist', () => {
+    const useStore = create<
+      CounterState,
+      [
+        ['zustand/subscribeWithSelector', never],
+        ['zustand/persist', CounterState]
+      ]
+    >(
+      subscribeWithSelector(
+        persist(
+          (set, get) => ({
+            count: 0,
+            inc: () => set({ count: get().count + 1 }, false),
+          }),
+          { name: 'count' }
+        )
+      )
+    )
+    const TestComponent = () => {
+      useStore((s) => s.count) * 2
+      useStore((s) => s.inc)()
+      useStore().count * 2
+      useStore().inc()
+      useStore.getState().count * 2
+      useStore.getState().inc()
+      useStore.subscribe(
+        (state) => state.count,
+        (count) => console.log(count * 2)
+      )
+      useStore.setState({ count: 0 }, false)
+      useStore.persist.hasHydrated()
+      return <></>
+    }
+    TestComponent
+  })
+})
--- a/tests/persistAsync.test.tsx
+++ b/tests/persistAsync.test.tsx
@ -228,7 +228,11 @@ describe('persist middleware with async configuration', () => {
      removeItem: () => {},
    }

-    const useStore = create(
+    const useStore = create<{
+      count: number
+      name: string
+      setName: (name: string) => void
+    }>()(
      persist(
        (set) => ({
          count: 0,
@ -407,7 +411,8 @@ describe('persist middleware with async configuration', () => {
      persist(() => ({ count: 0, actions: { unstorableMethod } }), {
        name: 'test-storage',
        getStorage: () => storage,
-        merge: (persistedState, currentState) => {
+        merge: (_persistedState, currentState) => {
+          const persistedState = _persistedState as any
          delete persistedState.actions

          return {
--- a/tests/persistSync.test.tsx
+++ b/tests/persistSync.test.tsx
@ -271,7 +271,8 @@ describe('persist middleware with sync configuration', () => {
      persist(() => ({ count: 0, actions: { unstorableMethod } }), {
        name: 'test-storage',
        getStorage: () => storage,
-        merge: (persistedState, currentState) => {
+        merge: (_persistedState, currentState) => {
+          const persistedState = _persistedState as any
          delete persistedState.actions

          return {
--- a/tests/types.test.tsx
+++ b/tests/types.test.tsx
@ -30,11 +30,11 @@ it('can use exposed types', () => {
    }
  }
  const selector: StateSelector<ExampleState, number> = (state) => state.num
-  const partial: PartialState<ExampleState, 'num' | 'numGet'> = {
+  const partial: PartialState<ExampleState> = {
    num: 2,
    numGet: () => 2,
  }
-  const partialFn: PartialState<ExampleState, 'num' | 'numGet'> = (state) => ({
+  const partialFn: PartialState<ExampleState> = (state) => ({
    ...state,
    num: 2,
  })
@ -73,7 +73,7 @@ it('can use exposed types', () => {

  function checkAllTypes(
    _getState: GetState<ExampleState>,
-    _partialState: PartialState<ExampleState, 'num' | 'numGet'>,
+    _partialState: PartialState<ExampleState>,
    _setState: SetState<ExampleState>,
    _state: State,
    _stateListener: StateListener<ExampleState>,
@ -83,7 +83,7 @@ it('can use exposed types', () => {
    _destroy: Destroy,
    _equalityFn: EqualityChecker<ExampleState>,
    _stateCreator: StateCreator<ExampleState>,
-    _useStore: UseBoundStore<ExampleState>
+    _useStore: UseBoundStore<StoreApi<ExampleState>>
  ) {
    expect(true).toBeTruthy()
  }
@ -128,7 +128,6 @@ it('should have correct (partial) types for setState', () => {
  // ok, should not error
  store.setState({ count: 1 })
  store.setState({})
-  store.setState(() => undefined)
  store.setState((previous) => previous)

  // @ts-expect-error type undefined is not assignable to type number
@ -155,7 +154,7 @@ it('should allow for different partial keys to be returnable from setState', ()
    }
    return { count: 0 }
  })
-  store.setState<'count', 'something'>((previous) => {
+  store.setState((previous) => {
    if (previous.count === 0) {
      return { count: 1 }
    }
@ -166,10 +165,26 @@ it('should allow for different partial keys to be returnable from setState', ()
  })

  // @ts-expect-error Type '{ something: boolean; count?: undefined; }' is not assignable to type 'State'.
-  store.setState<'count', 'something'>((previous) => {
+  store.setState((previous) => {
    if (previous.count === 0) {
      return { count: 1 }
    }
    return { something: true }
  })
 })
+
+it('state is covariant', () => {
+  const store = create<{ count: number; foo: string }>()(() => ({
+    count: 0,
+    foo: '',
+  }))
+
+  const _testIsCovariant: StoreApi<{ count: number }> = store
+
+  // @ts-expect-error should not compile
+  const _testIsNotContravariant: StoreApi<{
+    count: number
+    foo: string
+    baz: string
+  }> = store
+})