diff --git a/docs/typescript.md b/docs/typescript.md
new file mode 100644
index 0000000000..af599af9ff
--- /dev/null
+++ 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>`.
diff --git a/docs/v4-migration.md b/docs/v4-migration.md
new file mode 100644
index 0000000000..be8c03aac9
--- /dev/null
+++ 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) => ..., ...)`.
diff --git a/readme.md b/readme.md
index cfd7a22d0a..7c8cf4cefa 100644
--- 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.
-
-```tsx
-import { combine } from 'zustand/middleware'
-
-const useStore = create(
-  combine(
-    { bears: 0 },
-    (set) => ({ increase: (by: number) => set((state) => ({ bears: state.bears + by })) })
-  ),
-)
-```
+A more complete TypeScript guide is [here](https://github.com/pmndrs/zustand/blob/main/docs/typescript.md).
 
-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).
diff --git a/src/context.ts b/src/context.ts
index f5266a4e57..44d0645c27 100644
--- 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]
     )
   }
diff --git a/src/middleware.ts b/src/middleware.ts
index 90642820a8..c1787bac3d 100644
--- 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'
diff --git a/src/middleware/combine.ts b/src/middleware/combine.ts
index 42899b5baf..8ca4ecbc0c 100644
--- 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))
diff --git a/src/middleware/devtools.ts b/src/middleware/devtools.ts
index 72caf95957..f19f93f34d 100644
--- 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>
->(
-  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>
+type Devtools = <
+  T extends State,
+  Mps extends [StoreMutatorIdentifier, unknown][] = [],
+  Mcs extends [StoreMutatorIdentifier, unknown][] = []
 >(
-  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, [], []>>
+
+type PopArgument<T extends (...a: never[]) => unknown> = T extends (
+  ...a: [...infer A, infer _]
+) => infer R
+  ? (...a: A) => R
+  : never
+
+export type NamedSet<T extends State> = WithDevtools<StoreApi<T>>['setState']
+
+const devtoolsImpl: DevtoolsImpl = (fn, options) => (set, get, api) => {
+  type S = ReturnType<typeof fn>
 
-    const extension = extensionConnector.connect(devtoolsOptions)
+  const devtoolsOptions =
+    options === undefined
+      ? {}
+      : typeof options === 'string'
+      ? { name: options }
+      : options
 
-    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()
+  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 initialState = fn(api.setState, get, api)
-    extension.init(initialState)
+  const extension = extensionConnector.connect(devtoolsOptions)
 
-    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)
-      }
-    }
+  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)
 
-    ;(
-      extension as unknown as {
-        // FIXME https://github.com/reduxjs/redux-devtools/issues/1097
-        subscribe: (
-          listener: (message: Message) => void
-        ) => (() => void) | undefined
+  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
       }
-    ).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
-              }
+      ;(originalDispatch as any)(...a)
+    }
+  }
 
-              if (!(api as any).dispatchFromDevtools) return
-              if (typeof (api as any).dispatch !== 'function') return
-              ;(api as any).dispatch(action)
-            }
+  ;(
+    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'
           )
+          return
+        }
+        return parseJsonThen<{ type: unknown; state?: PartialState<S> }>(
+          message.payload,
+          (action) => {
+            if (action.type === '__setState') {
+              setStateFromDevtools(action.state as PartialState<S>)
+              return
+            }
 
-        case 'DISPATCH':
-          switch (message.payload.type) {
-            case 'RESET':
-              setStateFromDevtools(initialState)
-              return extension.init(api.getState())
+            if (!(api as any).dispatchFromDevtools) return
+            if (typeof (api as any).dispatch !== 'function') return
+            ;(api as any).dispatch(action)
+          }
+        )
 
-            case 'COMMIT':
-              return extension.init(api.getState())
+      case 'DISPATCH':
+        switch (message.payload.type) {
+          case 'RESET':
+            setStateFromDevtools(initialState)
+            return extension.init(api.getState())
 
-            case 'ROLLBACK':
-              return parseJsonThen<S>(message.state, (state) => {
-                setStateFromDevtools(state)
-                extension.init(api.getState())
-              })
+          case 'COMMIT':
+            return extension.init(api.getState())
 
-            case 'JUMP_TO_STATE':
-            case 'JUMP_TO_ACTION':
-              return parseJsonThen<S>(message.state, (state) => {
-                setStateFromDevtools(state)
-              })
+          case 'ROLLBACK':
+            return parseJsonThen<S>(message.state, (state) => {
+              setStateFromDevtools(state)
+              extension.init(api.getState())
+            })
 
-            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 'JUMP_TO_STATE':
+          case 'JUMP_TO_ACTION':
+            return parseJsonThen<S>(message.state, (state) => {
+              setStateFromDevtools(state)
+            })
 
-            case 'PAUSE_RECORDING':
-              return (isRecording = !isRecording)
+          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
           }
-          return
-      }
-    })
 
-    return initialState
-  }
+          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
diff --git a/src/middleware/immer.ts b/src/middleware/immer.ts
new file mode 100644
index 0000000000..7c302caf75
--- /dev/null
+++ 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
diff --git a/src/middleware/persist.ts b/src/middleware/persist.ts
index 49f1436efa..650da0a2ee 100644
--- 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,
-    }
-
-    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)
-    }
-
-    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 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,
+  }
 
-    const thenableSerialize = toThenable(options.serialize)
+  let hasHydrated = false
+  const hydrationListeners = new Set<PersistListener<S>>()
+  const finishHydrationListeners = new Set<PersistListener<S>>()
+  let storage: StateStorage | undefined
 
-    const setItem = (): Thenable<void> => {
-      const state = options.partialize({ ...get() })
+  try {
+    storage = options.getStorage()
+  } catch (e) {
+    // prevent error if the storage is not defined (e.g. when server side rendering a page)
+  }
 
-      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
+  const thenableSerialize = toThenable(options.serialize)
 
-    // rehydrate initial state with existing stored state
-    const hydrate = () => {
-      if (!storage) return
+  const setItem = (): Thenable<void> => {
+    const state = options.partialize({ ...get() })
 
-      hasHydrated = false
-      hydrationListeners.forEach((cb) => cb(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 postRehydrationCallback =
-        options.onRehydrateStorage?.(get()) || undefined
+  const savedSetState = api.setState
 
-      // 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`
+  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
               )
-            } else {
-              return deserializedStorageValue.state
             }
+            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,
         }
+      })
+      .then((migratedState) => {
+        stateFromStorage = options.merge(
+          migratedState as S,
+          get() ?? configResult
+        )
 
-        if (newOptions.getStorage) {
-          storage = newOptions.getStorage()
-        }
-      },
-      clearStorage: () => {
-        storage?.removeItem(options.name)
-      },
-      rehydrate: () => hydrate() as Promise<void>,
-      hasHydrated: () => hasHydrated,
-      onHydrate: (cb) => {
-        hydrationListeners.add(cb)
+        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)
+      })
+  }
 
-        return () => {
-          hydrationListeners.delete(cb)
-        }
-      },
-      onFinishHydration: (cb) => {
-        finishHydrationListeners.add(cb)
+  ;(api as StoreApi<S> & StorePersist<S, S>).persist = {
+    setOptions: (newOptions) => {
+      options = {
+        ...options,
+        ...newOptions,
+      }
 
-        return () => {
-          finishHydrationListeners.delete(cb)
-        }
-      },
-    }
+      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
+}
 
-    hydrate()
+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]>
 
-    return stateFromStorage || configResult
+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
diff --git a/src/middleware/redux.ts b/src/middleware/redux.ts
index 8148b75c49..8bffda989e 100644
--- 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
 
-/**
- * @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
+interface Action {
+  type: unknown
 }
 
-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 WithRedux<S, A> = Write<Cast<S, object>, StoreRedux<Cast<A, Action>>>
 
-type Write<T extends object, U extends object> = Omit<T, keyof U> & U
-type Cast<T, U> = T extends U ? T : U
+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]]>
 
-type WithRedux<S, A> = Write<Cast<S, object>, StoreRedux<Cast<A, Action>>>
+declare module '../vanilla' {
+  interface StoreMutators<S, A> {
+    'zustand/redux': WithRedux<S, A>
+  }
+}
 
-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
-
-    return { dispatch: (...a) => api.dispatch(...a), ...initial }
+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
diff --git a/src/middleware/subscribeWithSelector.ts b/src/middleware/subscribeWithSelector.ts
index 171acafd33..70ba70522e 100644
--- 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
diff --git a/src/react.ts b/src/react.ts
index 909b2fd079..e5af6c7700 100644
--- 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
diff --git a/src/vanilla.ts b/src/vanilla.ts
index 3cdd9695c3..4c7996fc67 100644
--- 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
-
-function createStore<
-  TState extends State,
-  CustomSetState,
-  CustomGetState,
-  CustomStoreApi extends StoreApi<TState>
->(
-  createState: StateCreator<
-    TState,
-    CustomSetState,
-    CustomGetState,
-    CustomStoreApi
-  >
-): CustomStoreApi
-
-function createStore<TState extends State>(
-  createState: StateCreator<TState, SetState<TState>, GetState<TState>, any>
-): StoreApi<TState>
-
-function createStore<
-  TState extends State,
-  CustomSetState,
-  CustomGetState,
-  CustomStoreApi extends StoreApi<TState>
+  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 }
+
+// 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>
+
+type PopArgument<T extends (...a: never[]) => unknown> = T extends (
+  ...a: [...infer A, infer _]
+) => infer R
+  ? (...a: A) => R
+  : never
+
+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
 }
 
-export default createStore
-
-// eslint-disable-next-line @typescript-eslint/no-unused-vars
-export interface StoreMutators<S, A> {}
-export type StoreMutatorIdentifier = keyof StoreMutators<unknown, unknown>
+const createStore = ((createState) =>
+  createState ? createStoreImpl(createState) : createStoreImpl) as CreateStore
 
-export type Mutate<S, Ms> = Ms extends []
-  ? S
-  : Ms extends [[infer Mi, infer Ma], ...infer Mrs]
-  ? Mutate<StoreMutators<S, Ma>[Mi & StoreMutatorIdentifier], Mrs>
-  : never
+export default createStore
diff --git a/tests/context.test.tsx b/tests/context.test.tsx
index 7427cf5631..ff2f0afbaa 100644
--- 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)
diff --git a/tests/devtools.test.tsx b/tests/devtools.test.tsx
index 6f9dadfd7e..880555333e 100644
--- 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(
diff --git a/tests/middlewareTypes.test.tsx b/tests/middlewareTypes.test.tsx
index afe02f6849..6e29d1cdb6 100644
--- 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
+  })
+})
diff --git a/tests/persistAsync.test.tsx b/tests/persistAsync.test.tsx
index 3b888f1ce5..30b2d7c4dd 100644
--- 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 {
diff --git a/tests/persistSync.test.tsx b/tests/persistSync.test.tsx
index 6ef1bb42fc..798c98745e 100644
--- 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 {
diff --git a/tests/types.test.tsx b/tests/types.test.tsx
index d088af213e..69aaf230d3 100644
--- 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
+})