Repository: pmndrs/zustand
Pull request #725 (2021-12-23 21:25:14)

breaking(types): Add higher kinded mutator types
Okay @dai-shi now you have two options:

**A. This PR (not complete yet)**
Makes the change as minimal as possible, just adds the types proposed in #710, zero change in runtime code.

**B. #662**
Not only adds #710 but also rewrites the whole codebase to drastically increase the code quality. Don't mind me but the current code is quite a mess, I like to write code for fun so I thought let's rewrite this whole thing nicely. Now this obviously will produce relatively large and unhelpful diffs, nothing can be done about that. It would be much easier to review if you compare the files locally without the red-green diff.
But now I have a small condition: I'll remove prettier and some lint rules, but don't worry you can add them back but you'll have to do it yourself (you have the edit rights). Consider this a little treat that I'm giving myself :P I hope you won't mind.


Let me know which of the two options you prefer. If you think increasing the code quality is worth taking the pain to review, I'd suggest you go with the second option. I don't mind whichever you pick.


codesandbox[bot]: This pull request is automatically built and testable in [CodeSandbox](https://codesandbox.io).

  To see build info of the built libraries, click [here](https://ci.codesandbox.io/status/pmndrs/zustand/pr/725/builds/203058) or the icon next to each commit SHA.

Latest deployment of this branch, based on commit 51ec8cf583ebccd1911b383e0fbd851bb5dc7832:

|Sandbox| Source |
|--|--|
|[React](https://codesandbox.io/s/react-lmknm)| Configuration |
|[React Typescript](https://codesandbox.io/s/react-typescript-ew9td)| Configuration |
|[React Browserify](https://codesandbox.io/s/react-browserify-tpv96)| Configuration |
|[React Snowpack](https://codesandbox.io/s/react-snowpack-csus9)| Configuration |
|[React Parcel](https://codesandbox.io/s/react-parcel-mtc55)| Configuration |
|[Next.js](https://codesandbox.io/s/next-js-1mx9w)| Configuration |

dai-shi: @devanshj I really appreciate your work on these.
tbh, I like this PR approach. at the same time, I like your rewrite too.

Sorry, I mixed two concerns in my previous comments in other PRs.
The first concern is backward compatibility for users.
The second concern is my ability to review the entire (rewritten) code at once.

So, to summarize my thoughts again:
- We want to keep 100% backward compatible for JS code in v4.0.0
- We want to fix store types with mutator types in v4.0.0 (reverting my work in v3.6.x is rather welcome. type breaking change for this fix is acceptable.)
- Let's keep exported type names backward compatible even if they are little weird. (We don't just rename types for better ones if the types are the same.)
- I actually want (and had even wanted in v3) to drop some exported types in v4.0.0 (deprecate them in v3.7.0) if they don't confuse people too much.
- I want to release v4.0.0 with minimal diff, and rewrite the whole code base step by step without changing exported types after releasing. So, we need to foresee what types to drop in advance for the rewrite.

Again, I prefer this PR approach, but want to take the rewrite idea as much as possible.
Probably, your full rewrite in #662 isn't possible with this approach, but we could try best effort.
devanshj: Got it, so let's go with this PR. Then once this is ready we'll release it in v4.0.0 and we can do a rewrite later if wanted.

Regarding dropping some types, `vanilla.ts` is ready on this PR, you can comment which ones you're thinking of dropping for v4. My suggestions are in #715 for reference.
devanshj: > I wonder if we still have design options, like giving up something.

Even if we give up some things the types still will be complex. And more over the users don't have to face any of the current complexity.

Regarding the `createWithState` thing, the users will be happy because now they don't have to annotated the store type. Take a look at diff of `middlewareTypes.test.tsx` in #662. Here's an [example](https://github.com/pmndrs/zustand/pull/662/files#diff-c65003d390ca71427c9de9367e1d32d897005c245143ad673cad2c73d79005ecL385-L390):

```diff
-     const useStore = create<
-       CounterState,
-       SetState<CounterState>,
-       GetState<CounterState>,
-       StoreApiWithPersist<CounterState> & StoreApiWithDevtools<CounterState>
-     >(
+     const useStore = createWithState<CounterState>()(
```
dai-shi: Yeah, that was ugly. That was not my original expectation when I started.
But, you can't assume all TS users skill level and you alone can't teach them all.
Like, some people prefer annotating types rather than inferring types.
devanshj: > But, you can't assume all TS users skill level and you alone can't teach them all.

I know what you're saying but if they can read the release and migration notes they'll be able to pick it up

> Like, some people prefer annotating types rather than inferring types.

I doubt in this case anyone would want to annotate the store types, but in case they want to, they can `create<T, [[$$persist, U], [$$devtools, never]]>()` a little more unfamiliar but arguably better than before.
devanshj: > oops, did i stop you moving forward?

No no haha, I was taking a little break to work on [`@sthir/predicate`](https://github.com/devanshj/sthir/tree/main/packages/predicate), which is quite complete now.

I'll get back to this soon, hopefully in a day or two!
devanshj: > back to the code. these two lines are hard to understand.
how does it work?
https://www.typescriptlang.org/play?#code/JYOwLgpgTgZghgYwgAgMpgPZQgWQK5hyZQDOAPKgDTICCAfMgN4C+AUGAJ4AOK6WuBIlgCSAEwjhgMYNGQBeZAGsIHDDDTEBhYuTwhFIDAHcQ1PQeMg6QA
I thought I understood whey I first saw, but now I'm not sure.

> why do we need `StoreMutators` and `StoreMutatorIdentifier` for Mutate? Are they only to guarantee the uniqueness?

The idea is we want a way to "pass" types from one module to another module. So to do that what we do is we declare an interface in the module the needs to receive the type and in the modules that needs to send the type we augment the interface, here's an example...

![](https://user-images.githubusercontent.com/30295578/148380764-cc808246-0d15-4789-ac5a-ec295603ae12.png)

So you can think of all export interfaces as "global" any other module can append new properties to them.

Now we need a global "registry" of middleware mutators, hence we define an interface `StoreMutators` which all middleware modules would add append their mutators and then we'd use those mutators in `Mutate`. So when we write...

```
// @file devtool.ts
interface StoreMutators<S, A>
  { [$$devtools]: WithDevtools<S, A>
  }
```

...we're actually passing the `WithDevtools` type to `vanilla.ts` which will use it in `Mutate`

Also we don't really need to worry about uniqueness because if you different modules append two different mutators then typescript would so them error.

![](https://user-images.githubusercontent.com/30295578/148382288-517b5eaa-a9c1-41b3-a839-f97931d0f3b3.png)

But the thing is the user is importing a third party module then if there are name clashes then it becomes a problem because the third party person won't know what names are already registed, so we'll keep a convention to for third party middleware authors to start the mutator name with their npm package name (the exported function name can be anything ofc). So for zustand, it'd become `"zustand/devtools": WithDevtools<S, A>` etc.

My suggestion would be to use symbols for this purpose, but I'm with this convention too.

And lastly we use `StoreMutatorIdentifier` simply has a short hand for writing `keyof StoreMutators<unknown, unknown>` to get a type of all mutators that are registered.






dai-shi: > The idea is we want a way to "pass" types from one module to another module. So to do that what we do is we declare an interface in the module the needs to receive the type and in the modules that needs to send the type we augment the interface, here's an example...

Thanks for the explanation. I think I understand the concept. And, I feel like I should understand it better.
May I ask you a favor? Would you create a small standalone typescript playground with this idea, so that I can play with it?
devanshj: > Would you create a small standalone typescript playground with this idea, so that I can play with it?

Sure, does [this](https://tsplay.dev/mx6A7w) work? Or you need something even more minimal? [Here](https://tsplay.dev/wQe9nw)'s a more minimal version of that only shows how `Mutate` works.
devanshj: > [Here](https://tsplay.dev/N53XoW)'s a more minimal version of that only shows how `Mutate` works.

I think you already understand how it work but here's a JS version of the above type-level code...

```js
const storeMutators = {}
// interface StoreMutators<S, A> {}

storeMutators.withA = (s, a) => Object.assign({}, s, { a: a })
// interface StoreMutators<S, A>
//  { withA: Write<S, { a: A }>
//  }
  
storeMutators.withB = (s, a) => Object.assign({}, s, { b: a })
// interface StoreMutators<S, A>
//  { withB: Write<S, { b: A }>
//  }

const mutate = (s, ms) =>
  ms.length === 0 ? s :
  (([[mi, ma], ...mrs]) => mutate(storeMutators[mi](s, ma), mrs))(ms)
  
// type Mutate<S, Ms> =
//   Ms extends [] ? S :
//   Ms extends [[infer Mi, infer Ma], ...infer Mrs]
//     ? Mutate<
//        StoreMutators<S, Ma>[Mi & StoreMutatorIdentifier],
//        Mrs
//      > :
//  never

const testStore = mutate({ get: () => {} }, [["withA", "hello"], ["withB", "world"]]); 
// type TestStore = Mutate<Store<{ count: number }>, [["withA", string], ["withB", string]]>

console.log(testStore)
```
<details>
<summary>Console Output: </summary>

![console output](https://user-images.githubusercontent.com/30295578/148652058-edfa6d88-f5b3-4e41-8532-3d114fb83920.png)
</details>


devanshj: This is almost ready now
dai-shi: > [Here](https://tsplay.dev/N53XoW)'s a more minimal version of that only shows how `Mutate` works.

This really helps.
I just missed the fact `interface` can be overloaded. 😓 

[This](https://www.typescriptlang.org/play?pretty=false#code/C4TwDgpgBAKhDOwDKwD2AnaBeKBZArsAIbAQA8KG5A3lAMar4B2wAXFE-gLYBGE6UAL4A+ADRQA2hIBEAdwCWwABYBBaeMTp5TAOYBdcTIXKAQuqibt+vcIDcAKAD0jqK6gA9APz3720ugAzIjpoSkwCYjR0eApxFWEoantXY1V2AHUtUljEqCJ2FSE7e0EfP34gkKgwiAiSDBikOISklMUlEwys8ibcngKih1L7UEhqqPIYBJxWi0iIdhghh18WCuDQibqoxubE0tHNqm2MAEkAEwgWeQD5figcAGsIEFQA8ePCeuiyZkemVCyJjiP4AoHFEbgaDbHriXDwaZ4eBQCAAD1ITHOyIkemSUE81Tx7HhKPRVyxkgk2gC91w8nE1NpRAMUAAdOzGQJcNFca4CXTSRiKTUTugLldgDc7ug8Xy8F9siKFQ0crgiMIJHSWdyEbKoOwmBAAG78IkcY38FaHKCZRSTcQAVURAHkuIoyDBxM9Xu8nVAAGRQB0OIA) is what I tried. I understand 100% of this.

Now, it made me think [this](https://www.typescriptlang.org/play?pretty=false#code/C4TwDgpgBAKhDOwDKwD2AnaBeKBZArsAIbAQA8KG5A3lAMar4B2wAXFE-gLYBGE6UAL4A+ADRQA2rSLtE6AJZMA5kPG0e7Tr35CAusIDcAKAD0JqBagA9APxGjoSFEqYyMYVBzUjFxCQjsMMaCxg7g0ATEpBTiuPAeOHFQEAAepEwAJvCSuj5QNs557Emp6VmSigBmOrjiAHQNVTXo8LmW+XiE-mQA6grRSLFieC3CRRwQAG78oY7QffLRMOIAqglQAPJci27iANYQIKiVUGtQAGSnxkA) works. But, it doesn't work in practice in our case???

---
btw, I'm so amazed to see `^?` in the playground. Where can I find such tips in docs?
Also, how do you create the tsplay.dev shorten url?
dai-shi: > Sure, does [this](https://tsplay.dev/mx6A7w) work?

So, [this](https://www.typescriptlang.org/play?#code/DYUwLgBAzmD2BOIDyBXSBeCBjRBDMIAPAN7awoB2YAXBBSgLYBGI8EAvgHwAUAlNwCgIEAO4BLMAAsAgoOHCARLgUAaIfPFSAQtwVNVEbjAQgAkhV4R0nQ6SzkqtAAwdevde94BuAQLFVWADNcLBAIAGU4REIAFQgQAA8CCgATKAgAVQoAawpYEQpI-BAbYnUAc3Aiglo+KxsYn2EoKrBi2kIAOUhE5LSIACV4pJBU9LB4FDCAfgg42gAFXHgwMVxgWM4VQeG+9KZYWFBcCh51YQpe6uR4DIAHFOL4Wm6IAB9DIzaauctrCG6vDU8mgknIwBSAxAd2AIRA01oAw89QgADdYGIUk1oCgmFAcGIWLVgGIYKNWLUYO05ts7ohUWJyFBrrQYn8bOjMezDNzOVj1CkQDB4LAAJ61XkY-nsXxgUV3MJZXL5QrfMKYWBMABWICwYB8AkFWFhiDIFBg2DwPzKwliu1G-SVeQK122AFlYOleg70pRlQUANoAXSsEGDZxB-gkaxJAC8KREomYKNH1mJ4-BYttg+7PZwPLQ3WhioRIiZNrmoJxsXbvWNMjlnarijxeB0PV6RvW-U3g6Hw3J5FHVmmM7Qy4hzKm46ws2Gg5X88JuUW1aWkxWIB3qwJ2Aa5QrEyYpyOZ5m4nXHY2Va6t6TK9sMqGYjZMNxB8ZEIXiwR1+WX+6pJbOoAAkIGfiAq74Iy5qFqSAjchklgAGQQDaEBgRBUGrLA5oIluno+HusrymE2FEOE7pVqGbqdns87qLM4TqIWdE+mG-iBKwW7bAAdPxnHcW68BQEGjFbj+RAAOrwBIFHulsW4iUuEC0JcqKsAaRommE9jmpAmgyB06gXl2V7+s2BDAhA0jWW6dxsd216BiGmA5uobpYI5-Q9iqfZuWJEYQLgtC2eow4xumCYTsm05Ree2b8bx9lQNspAhQCsDmFxmbSDY7ALluXn5tyMUnpFGZzil2bpS8WUUDlhB5RwfH8Z5ok7n4ATwMEoRHog2EIFApbbM16GGdItAyXJI1ocFoUcDuMqGrqOlmhahlaMZwimfRTo3mq1laHZDn2k5Fn+fOdleWdPnORQl3BgIQVMLQx3hSmp7xeOSblaOs4xIlbUOWlECvZl2Wzlo+WFe1JUomVn0VQD7og2GpDg509WNdDLUQEl7VBp1-gED1cL9ZBP5DbNY3qJtU2yb+lFzeD0iLURJGHgA4uAc4ANLbAAYqG6msK+EB87d6TZCAoqwIEcwQLMMQBnzIa0IL+6kRA02-oDmTi0gDASHOMtywrGQ2KhGRa4eWOQ+e4sBjEQbO1LEC+QUSsQC4akgBp8BBj4QA) is what I modified. Please let me know what would be the pitfall.
devanshj: > Now, it made me think this works. But, it doesn't work in practice in our case???

> So, this is what I modified. Please let me know what would be the pitfall.


It's good to see you could come up with the simplified version. But it won't work as soon as you have a `T` anywhere (at a covariant) in the mutation. A real example of that is persist: https://tsplay.dev/NaYzyN Note that it would work if you make `T` concrete by passing it as type-parameter: https://tsplay.dev/WzLjRN But that still is not a good solution as the code should at least compile without having to pass type parameters.

Your suggestion is same as the "Ordered writes" section in #710. I suggest you to read from "Dependent Mutations" section onwards because this the place where I begin to explain how I solve this problem and how we end up with `StoreMutators`...

> ## Dependent mutations
>
> At any point if our mutation has a `T` on a covariant position, it no longer compiles. Basically if the middleware (mutation) "gives" or "returns" the state it means `T` is on a covariant position (like in case of `persist`, `onHydrate` and `onFinishHydration` "give" the state), it won't work. Here's a simple example...

---

> btw, I'm so amazed to see ^? in the playground.

Haha yeah it's cool, I think it'll eventually come to vscode too.

> Where can I find such tips in docs?

The docs probably don't mention it, I learnt when someone on the typescript team used it in an issue. [Orta](https://twitter.com/orta) usually implements features of the playground (like this one) so I guess one way to learn more tips is to follow him or twitter.

> Also, how do you create the tsplay.dev shorten url?

One way to do it is first copy the long usual url, then go to tsplay.dev, paste it there and you'd get a shorter version.
Another way to install the plugin in the playground, go to the "Plugins" tab, there should be a "Link Shortener" plugin, tick the checkbox next to it to install it. Then you should be able to use it directly from the playground:

![image](https://user-images.githubusercontent.com/30295578/148680916-e0250475-aac0-467a-a928-6adb45bf7276.png)




dai-shi: > Your suggestion is same as the "Ordered writes" section in #710. I suggest you to read from "Dependent Mutations" section onwards

Glad to see I'm following your path.

> The thing for some reason typescript (this could be a bug)

So, is this true that, if TS fixes this bug or behavior, we could have a simplified version.

Now, I understand why "Higher kinded mutations" is introduced. Let me see if I can read the code...

> there should be a "Link Shortener" plugin

Ah, I somehow missed it when I looked for it in the plugin list. 😓 

devanshj: > So, is this true that, if TS fixes this bug or behavior, we could have a simplified version.

Yes, although it's very unlikely it'd be a bug.

> Now, I understand why "Higher kinded mutations" is introduced. Let me see if I can read the code...

Yep, the idea is to make `T` somehow disappear from the mutation. So instead of a mutation `{ a: A, t: T }` we use a higher kinded mutation `(T, A) => { a: A, t: T }` and then eventually use a higher kinded mutator `(S, A) => S extends { getState: () => infer T } ? Write<S, { a: A, t: T }> : never` to simplify the things a bit.
dai-shi: btw, if I were to design the zustand api from scratch, I'd never introduce type complexity. 😄 
devanshj: > btw, if I were to design the zustand api from scratch, I'd never introduce type complexity. 😄

Haha yeah I'm aware you had mentioned that you too don't like the mutating api
devanshj: > @devanshj I hope you will be back soon!

Sorry somehow I missed this notification 😅 and was also working on my blog
dai-shi: @devanshj Hi, do you have a chance to continue this and #772? I hope to release v3.7 in a few weeks or before that. I think we need to agree on #725 first to complete #722.
devanshj: As I said before this is 99% complete, could you review it all at once? I don't have a lot time to invest in this, so making small changes multiple times is not something that's possible, I hope you'd understand.

Once you review it completely, I'd make the changes you'd request and then this would be good to go.
dai-shi: No, I definitely didn't know your status. I thought you have been busy and I have waited for a long period.

Okay, I'm not confident if this will finish in one cycle. (Maybe not technical ones. If you don't have time I can take this PR over for non technical ones.)
But anyway, I will try to finish reviewing soon (I think mostly fine too, but there will be many small ones.)
Meanwhile, can you check the lint and the test? That was the major reason why I didn't take it seriously. Are you sure if they are all trivial things?
dai-shi: @devanshj Looking it again, this is exciting. Please invest time on this to finished it up. docs will be crucial (can be a separate PR).
devanshj: > No, I definitely didn't know your status. I thought you have been busy and I have waited for a long period.
> 
> Okay, I'm not confident if this will finish in one cycle. (Maybe not technical ones. If you don't have time I can take this PR over for non technical ones.) But anyway, I will try to finish reviewing soon (I think mostly fine too, but there will be many small ones.) Meanwhile, can you check the lint and the test? That was the major reason why I didn't take it seriously. Are you sure if they are all trivial things?

Ah looks like you must have missed [this comment](https://github.com/pmndrs/zustand/pull/725#issuecomment-1008102464), it's okay I should have made it clear that I'm waiting for review.

I don't mind if it takes multiple cycles as long as we're doing things in one streak. I think I missed looking at the CI but they're trivial, I'll fix it.

dai-shi: Please "Re-request revew", when it's ready.
devanshj: There's one test failing I'm not sure why 🤔 Could you give it a look?
devanshj: > breaking change notes (so that I can just copy&paste in release note)?

Yep will do that soon. Also don't merge this yet I have to make some checks regarding types.
dai-shi: One more thing: please check #812, which may affect v3.7.x and then this.
devanshj: (Just updating that I remember I'm yet to get back to this, I'll be back in a day or two)
devanshj: Looks like the some tests running with react 18 are failing, not sure why, @dai-shi could you check?
dai-shi: @devanshj I guess you failed to merge v4 properly. the important change is this: https://github.com/pmndrs/zustand/commit/d4b72002b9d1368975b294fc81835cc41c6e8851
devanshj: No I think I've got it right: https://github.com/devanshj/zustand/commit/d4b72002b9d1368975b294fc81835cc41c6e8851

You can even check the files:
https://github.com/devanshj/zustand/blob/b653b4b62bdd05efe05d751be6dd0f39d70ec50e/package.json#L185-L187
https://github.com/devanshj/zustand/blob/b653b4b62bdd05efe05d751be6dd0f39d70ec50e/yarn.lock#L2195-L2198


dai-shi: Oh, my bad then...
devanshj: I've updated the readme. Will write the breaking changes in the description soon.
devanshj: Also do you really want me to write the breaking changes? I mean because almost everything is a breaking change. For example even something as trivial as the `redux` middleware: Imagine if someone is passing the type parameters, then the `Cms` will get fixated to `[]` (as that's the default) instead of getting inferred from contextual inference. Writing all breaking changes will be tedious for me.

I think it's possible to make a migration guide from the readme I just made and looking at the source. The migration guide should be less of "If you're doing X then now do Y" and more of "Just do Y" because doing Y is really simple.
dai-shi: @devanshj Thanks for working on this. We are almost there. (I'll review diffs later.)

> Writing all breaking changes will be tedious for me.

Hmm, that's unfortunate. Then, please write down all changes in terms of JS code (after tsc emitting).
For type changes, let's do it in b) `./v4-migration.md` as you suggest (and we'll keep improving it).

> a) Precise breaking changes of this PR (let's write them in the PR description)
> b) Migration guide (let's create a new md file like ./v4-migration.md.)
> c) Create a new "TypeScript" section in readme.md and show our recommended usage (combining existing typescript-related descriptions.
devanshj: > Then, please write down all changes in terms of JS code (after tsc emitting).

There are no breaking changes on runtime level. Only two additions...
- `create` (from `zustand`, `zustand/vanilla` and `zustand/react` modules) now also takes zero arguments and returns the "actual" `create` that takes the store initializer. (This is a workaround for [microsoft/TypeScript#10571](https://github.com/microsoft/TypeScript/issues/10571))
- `zustand/middleware` module now also exports a `immer` middleware. This is to make the migration of users that have copy-pasted the immer middleware from `middlewareTypes.test.tsx` easier.

> For type changes, let's do it in b) ./v4-migration.md as you suggest (and we'll keep improving it).

Sure (and yes `./docs/v4-migration.md` as you suggested in the other comment)
dai-shi: > There are no breaking changes on runtime level. Only two additions...

That's good. Would you be able to confirm by building v4 branch and this branch?

> Will continue tomorrow.

Looking forward to it!
devanshj: Did some updates, it's 90% complete now, will address the remaining things tomorrow.
dai-shi: > Would you be able to confirm by building v4 branch and this branch?

I just did this.

<details>
<summary>diff -ur dist-v4/esm dist/esm (js only)</summary>

```diff
diff -ur dist-v4/esm/context.js dist/esm/context.js
--- dist-v4/esm/context.js	2022-04-05 15:48:43.000000000 +0900
+++ dist/esm/context.js	2022-04-05 15:51:55.000000000 +0900
@@ -1,5 +1,11 @@
-import { createContext as createContext$1, useRef, createElement, useContext, useMemo } from 'react';
-import { useStore } from 'zustand';
+import { useDebugValue, createContext as createContext$1, useRef, createElement, useContext, useMemo } from 'react';
+import { useSyncExternalStoreWithSelector } from 'use-sync-external-store/shim/with-selector.js';
+
+function useStore(api, selector = api.getState, equalityFn) {
+  const slice = useSyncExternalStoreWithSelector(api.subscribe, api.getState, null, selector, equalityFn);
+  useDebugValue(slice);
+  return slice;
+}
 
 function createContext() {
   const ZustandContext = createContext$1(void 0);
diff -ur dist-v4/esm/index.js dist/esm/index.js
--- dist-v4/esm/index.js	2022-04-05 15:48:48.000000000 +0900
+++ dist/esm/index.js	2022-04-05 15:52:00.000000000 +0900
@@ -1,7 +1,7 @@
 import { useDebugValue } from 'react';
 import { useSyncExternalStoreWithSelector } from 'use-sync-external-store/shim/with-selector.js';
 
-function createStore(createState) {
+const createStoreImpl = (createState) => {
   let state;
   const listeners = /* @__PURE__ */ new Set();
   const setState = (partial, replace) => {
@@ -21,18 +21,22 @@
   const api = { setState, getState, subscribe, destroy };
   state = createState(setState, getState, api);
   return api;
-}
+};
+const createStore = (createState) => createState ? createStoreImpl(createState) : createStoreImpl;
+var createStore$1 = createStore;
 
 function useStore(api, selector = api.getState, equalityFn) {
   const slice = useSyncExternalStoreWithSelector(api.subscribe, api.getState, null, selector, equalityFn);
   useDebugValue(slice);
   return slice;
 }
-function create(createState) {
-  const api = typeof createState === "function" ? createStore(createState) : createState;
+const createImpl = (createState) => {
+  const api = typeof createState === "function" ? createStore$1(createState) : createState;
   const useBoundStore = (selector, equalityFn) => useStore(api, selector, equalityFn);
   Object.assign(useBoundStore, api);
   return useBoundStore;
-}
+};
+const create = (createState) => createState ? createImpl(createState) : createImpl;
+var create$1 = create;
 
-export { create as default, useStore };
+export { create$1 as default, useStore };
diff -ur dist-v4/esm/middleware.js dist/esm/middleware.js
--- dist-v4/esm/middleware.js	2022-04-05 15:48:44.000000000 +0900
+++ dist/esm/middleware.js	2022-04-05 15:51:56.000000000 +0900
@@ -1,3 +1,5 @@
+import { produce } from 'immer';
+
 var __defProp$1 = Object.defineProperty;
 var __getOwnPropSymbols$1 = Object.getOwnPropertySymbols;
 var __hasOwnProp$1 = Object.prototype.hasOwnProperty;
@@ -14,7 +16,7 @@
     }
   return a;
 };
-const redux = (reducer, initial) => (set, get, api) => {
+const reduxImpl = (reducer, initial) => (set, get, api) => {
   api.dispatch = (action) => {
     set((state) => reducer(state, action), false, action);
     return action;
@@ -22,140 +24,102 @@
   api.dispatchFromDevtools = true;
   return __spreadValues$1({ dispatch: (...a) => api.dispatch(...a) }, initial);
 };
+const redux = reduxImpl;
 
-function devtools(fn, options) {
-  return (set, get, api) => {
-    let didWarnAboutNameDeprecation = false;
-    if (typeof options === "string" && !didWarnAboutNameDeprecation) {
-      console.warn("[zustand devtools middleware]: passing `name` as directly will be not allowed in next majorpass the `name` in an object `{ name: ... }` instead");
-      didWarnAboutNameDeprecation = true;
-    }
-    const devtoolsOptions = options === void 0 ? {} : typeof options === "string" ? { name: options } : options;
-    let extensionConnector;
-    try {
-      extensionConnector = window.__REDUX_DEVTOOLS_EXTENSION__;
-    } catch {
-    }
-    if (!extensionConnector) {
-      if ((import.meta.env && import.meta.env.MODE) !== "production" && typeof window !== "undefined") {
-        console.warn("[zustand devtools middleware] Please install/enable Redux devtools extension");
-      }
-      return fn(set, get, api);
+const devtoolsImpl = (fn, options) => (set, get, api) => {
+  const devtoolsOptions = options === void 0 ? {} : typeof options === "string" ? { name: options } : options;
+  let extensionConnector;
+  try {
+    extensionConnector = window.__REDUX_DEVTOOLS_EXTENSION__;
+  } catch {
+  }
+  if (!extensionConnector) {
+    if ((import.meta.env && import.meta.env.MODE) !== "production" && typeof window !== "undefined") {
+      console.warn("[zustand devtools middleware] Please install/enable Redux devtools extension");
     }
-    let extension = Object.create(extensionConnector.connect(devtoolsOptions));
-    let didWarnAboutDevtools = false;
-    Object.defineProperty(api, "devtools", {
-      get: () => {
-        if (!didWarnAboutDevtools) {
-          console.warn("[zustand devtools middleware] `devtools` property on the store is deprecated it will be removed in the next major.\nYou shouldn't interact with the extension directly. But in case you still want to you can patch `window.__REDUX_DEVTOOLS_EXTENSION__` directly");
-          didWarnAboutDevtools = true;
-        }
-        return extension;
-      },
-      set: (value) => {
-        if (!didWarnAboutDevtools) {
-          console.warn("[zustand devtools middleware] `api.devtools` is deprecated, it will be removed in the next major.\nYou shouldn't interact with the extension directly. But in case you still want to you can patch `window.__REDUX_DEVTOOLS_EXTENSION__` directly");
-          didWarnAboutDevtools = true;
-        }
-        extension = value;
-      }
-    });
-    let didWarnAboutPrefix = false;
-    Object.defineProperty(extension, "prefix", {
-      get: () => {
-        if (!didWarnAboutPrefix) {
-          console.warn("[zustand devtools middleware] along with `api.devtools`, `api.devtools.prefix` is deprecated.\nWe no longer prefix the actions/names" + devtoolsOptions.name === void 0 ? ", pass the `name` option to create a separate instance of devtools for each store." : ", because the `name` option already creates a separate instance of devtools for each store.");
-          didWarnAboutPrefix = true;
-        }
-        return "";
-      },
-      set: () => {
-        if (!didWarnAboutPrefix) {
-          console.warn("[zustand devtools middleware] along with `api.devtools`, `api.devtools.prefix` is deprecated.\nWe no longer prefix the actions/names" + devtoolsOptions.name === void 0 ? ", pass the `name` option to create a separate instance of devtools for each store." : ", because the `name` option already creates a separate instance of devtools for each store.");
-          didWarnAboutPrefix = true;
-        }
-      }
-    });
-    let isRecording = true;
-    api.setState = (state, replace, nameOrAction) => {
-      set(state, replace);
-      if (!isRecording)
-        return;
-      extension.send(nameOrAction === void 0 ? { type: devtoolsOptions.anonymousActionType || "anonymous" } : typeof nameOrAction === "string" ? { type: nameOrAction } : nameOrAction, get());
-    };
-    const setStateFromDevtools = (...a) => {
-      const originalIsRecording = isRecording;
-      isRecording = false;
-      set(...a);
-      isRecording = originalIsRecording;
-    };
-    const initialState = fn(api.setState, get, api);
-    extension.init(initialState);
-    if (api.dispatchFromDevtools && typeof api.dispatch === "function") {
-      let didWarnAboutReservedActionType = false;
-      const originalDispatch = api.dispatch;
-      api.dispatch = (...a) => {
-        if (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;
+    return fn(set, get, api);
+  }
+  const extension = extensionConnector.connect(devtoolsOptions);
+  let isRecording = true;
+  api.setState = (state, replace, nameOrAction) => {
+    set(state, replace);
+    if (!isRecording)
+      return;
+    extension.send(nameOrAction === void 0 ? { type: devtoolsOptions.anonymousActionType || "anonymous" } : typeof nameOrAction === "string" ? { type: nameOrAction } : nameOrAction, get());
+  };
+  const setStateFromDevtools = (...a) => {
+    const originalIsRecording = isRecording;
+    isRecording = false;
+    set(...a);
+    isRecording = originalIsRecording;
+  };
+  const initialState = fn(api.setState, get, api);
+  extension.init(initialState);
+  extension.subscribe((message) => {
+    var _a;
+    switch (message.type) {
+      case "ACTION":
+        if (typeof message.payload !== "string") {
+          console.error("[zustand devtools middleware] Unsupported action format");
+          return;
         }
-        originalDispatch(...a);
-      };
-    }
-    extension.subscribe((message) => {
-      var _a;
-      switch (message.type) {
-        case "ACTION":
-          if (typeof message.payload !== "string") {
-            console.error("[zustand devtools middleware] Unsupported action format");
+        return parseJsonThen(message.payload, (action) => {
+          if (action.type === "__setState") {
+            setStateFromDevtools(action.state);
             return;
           }
-          return parseJsonThen(message.payload, (action) => {
-            if (action.type === "__setState") {
-              setStateFromDevtools(action.state);
-              return;
-            }
-            if (!api.dispatchFromDevtools)
-              return;
-            if (typeof api.dispatch !== "function")
-              return;
-            api.dispatch(action);
-          });
-        case "DISPATCH":
-          switch (message.payload.type) {
-            case "RESET":
-              setStateFromDevtools(initialState);
-              return extension.init(api.getState());
-            case "COMMIT":
-              return extension.init(api.getState());
-            case "ROLLBACK":
-              return parseJsonThen(message.state, (state) => {
-                setStateFromDevtools(state);
-                extension.init(api.getState());
-              });
-            case "JUMP_TO_STATE":
-            case "JUMP_TO_ACTION":
-              return parseJsonThen(message.state, (state) => {
-                setStateFromDevtools(state);
-              });
-            case "IMPORT_STATE": {
-              const { nextLiftedState } = message.payload;
-              const lastComputedState = (_a = nextLiftedState.computedStates.slice(-1)[0]) == null ? void 0 : _a.state;
-              if (!lastComputedState)
-                return;
-              setStateFromDevtools(lastComputedState);
-              extension.send(null, nextLiftedState);
+          if (!api.dispatchFromDevtools)
+            return;
+          if (typeof api.dispatch !== "function")
+            return;
+          api.dispatch(action);
+        });
+      case "DISPATCH":
+        switch (message.payload.type) {
+          case "RESET":
+            setStateFromDevtools(initialState);
+            return extension.init(api.getState());
+          case "COMMIT":
+            return extension.init(api.getState());
+          case "ROLLBACK":
+            return parseJsonThen(message.state, (state) => {
+              setStateFromDevtools(state);
+              extension.init(api.getState());
+            });
+          case "JUMP_TO_STATE":
+          case "JUMP_TO_ACTION":
+            return parseJsonThen(message.state, (state) => {
+              setStateFromDevtools(state);
+            });
+          case "IMPORT_STATE": {
+            const { nextLiftedState } = message.payload;
+            const lastComputedState = (_a = nextLiftedState.computedStates.slice(-1)[0]) == null ? void 0 : _a.state;
+            if (!lastComputedState)
               return;
-            }
-            case "PAUSE_RECORDING":
-              return isRecording = !isRecording;
+            setStateFromDevtools(lastComputedState);
+            extension.send(null, nextLiftedState);
+            return;
           }
-          return;
+          case "PAUSE_RECORDING":
+            return isRecording = !isRecording;
+        }
+        return;
+    }
+  });
+  if (api.dispatchFromDevtools && typeof api.dispatch === "function") {
+    let didWarnAboutReservedActionType = false;
+    const originalDispatch = api.dispatch;
+    api.dispatch = (...a) => {
+      if ((import.meta.env && import.meta.env.MODE) !== "production" && 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;
       }
-    });
-    return initialState;
-  };
-}
+      originalDispatch(...a);
+    };
+  }
+  return initialState;
+};
+const devtools = devtoolsImpl;
 const parseJsonThen = (stringified, f) => {
   let parsed;
   try {
@@ -167,7 +131,7 @@
     f(parsed);
 };
 
-const subscribeWithSelector = (fn) => (set, get, api) => {
+const subscribeWithSelectorImpl = (fn) => (set, get, api) => {
   const origSubscribe = api.subscribe;
   api.subscribe = (selector, optListener, options) => {
     let listener = selector;
@@ -190,8 +154,9 @@
   const initialState = fn(set, get, api);
   return initialState;
 };
+const subscribeWithSelector = subscribeWithSelectorImpl;
 
-const combine = (initialState, create) => (set, get, api) => Object.assign({}, initialState, create(set, get, api));
+const combine = (initialState, create) => (...a) => Object.assign({}, initialState, create(...a));
 
 var __defProp = Object.defineProperty;
 var __getOwnPropSymbols = Object.getOwnPropertySymbols;
@@ -234,14 +199,14 @@
     };
   }
 };
-const persist = (config, baseOptions) => (set, get, api) => {
+const persistImpl = (config, baseOptions) => (set, get, api) => {
   let options = __spreadValues({
     getStorage: () => localStorage,
     serialize: JSON.stringify,
     deserialize: JSON.parse,
     partialize: (state) => state,
     version: 0,
-    merge: (persistedState, currentState) => __spreadValues(__spreadValues({}, currentState), persistedState)
+    merge: (persistedState, currentState) => __spreadValues(__spreadValues({}, currentState), typeof persistedState === "object" ? persistedState : {})
   }, baseOptions);
   let hasHydrated = false;
   const hydrationListeners = /* @__PURE__ */ new Set();
@@ -341,5 +306,15 @@
   hydrate();
   return stateFromStorage || configResult;
 };
+const persist = persistImpl;
+
+const immerImpl = (initializer) => (set, get, store) => {
+  store.setState = (updater, replace) => {
+    const nextState = typeof updater === "function" ? produce(updater) : updater;
+    return set(nextState, replace);
+  };
+  return initializer(store.setState, get, store);
+};
+const immer = immerImpl;
 
-export { combine, devtools, persist, redux, subscribeWithSelector };
+export { combine, devtools, immer, persist, redux, subscribeWithSelector };
diff -ur dist-v4/esm/vanilla.js dist/esm/vanilla.js
--- dist-v4/esm/vanilla.js	2022-04-05 15:48:43.000000000 +0900
+++ dist/esm/vanilla.js	2022-04-05 15:51:55.000000000 +0900
@@ -1,4 +1,4 @@
-function createStore(createState) {
+const createStoreImpl = (createState) => {
   let state;
   const listeners = /* @__PURE__ */ new Set();
   const setState = (partial, replace) => {
@@ -18,6 +18,7 @@
   const api = { setState, getState, subscribe, destroy };
   state = createState(setState, getState, api);
   return api;
-}
+};
+const createStore = (createState) => createState ? createStoreImpl(createState) : createStoreImpl;
 
 export { createStore as default };
```

</details>

Oops, I found a weird diff in context.js.
devanshj: > @devanshj Would you check the diff of middleware/devtools.js and, if it's valid, list the changes in the PR description?

The diff looks fine but most of it shouldn't belong to this PR, so I opened #892. Once you merge it with v4, I'll update this branch and then the diff should be better.
dai-shi: Hi, v4 branch is updated (cherry-picked #892 commits). Please resolve conflicts.
dai-shi: <details>
<summary>diff -ur dist-v4/esm dist/esm (js only)</summary>

```diff
diff -urw dist-v4/esm/index.js dist/esm/index.js
--- dist-v4/esm/index.js	2022-04-07 10:42:20.000000000 +0900
+++ dist/esm/index.js	2022-04-07 10:43:16.000000000 +0900
@@ -1,7 +1,7 @@
 import { useDebugValue } from 'react';
 import { useSyncExternalStoreWithSelector } from 'use-sync-external-store/shim/with-selector.js';
 
-function createStore(createState) {
+const createStoreImpl = (createState) => {
   let state;
   const listeners = /* @__PURE__ */ new Set();
   const setState = (partial, replace) => {
@@ -21,18 +21,22 @@
   const api = { setState, getState, subscribe, destroy };
   state = createState(setState, getState, api);
   return api;
-}
+};
+const createStore = (createState) => createState ? createStoreImpl(createState) : createStoreImpl;
+var createStore$1 = createStore;
 
 function useStore(api, selector = api.getState, equalityFn) {
   const slice = useSyncExternalStoreWithSelector(api.subscribe, api.getState, api.getServerState || api.getState, selector, equalityFn);
   useDebugValue(slice);
   return slice;
 }
-function create(createState) {
-  const api = typeof createState === "function" ? createStore(createState) : createState;
+const createImpl = (createState) => {
+  const api = typeof createState === "function" ? createStore$1(createState) : createState;
   const useBoundStore = (selector, equalityFn) => useStore(api, selector, equalityFn);
   Object.assign(useBoundStore, api);
   return useBoundStore;
-}
+};
+const create = (createState) => createState ? createImpl(createState) : createImpl;
+var create$1 = create;
 
-export { create as default, useStore };
+export { create$1 as default, useStore };
diff -urw dist-v4/esm/middleware.js dist/esm/middleware.js
--- dist-v4/esm/middleware.js	2022-04-07 10:42:17.000000000 +0900
+++ dist/esm/middleware.js	2022-04-07 10:43:12.000000000 +0900
@@ -1,3 +1,5 @@
+import { produce } from 'immer';
+
 var __defProp$1 = Object.defineProperty;
 var __getOwnPropSymbols$1 = Object.getOwnPropertySymbols;
 var __hasOwnProp$1 = Object.prototype.hasOwnProperty;
@@ -14,7 +16,7 @@
     }
   return a;
 };
-const redux = (reducer, initial) => (set, get, api) => {
+const reduxImpl = (reducer, initial) => (set, get, api) => {
   api.dispatch = (action) => {
     set((state) => reducer(state, action), false, action);
     return action;
@@ -22,9 +24,9 @@
   api.dispatchFromDevtools = true;
   return __spreadValues$1({ dispatch: (...a) => api.dispatch(...a) }, initial);
 };
+const redux = reduxImpl;
 
-function devtools(fn, options) {
-  return (set, get, api) => {
+const devtoolsImpl = (fn, options) => (set, get, api) => {
     const devtoolsOptions = options === void 0 ? {} : typeof options === "string" ? { name: options } : options;
     let extensionConnector;
     try {
@@ -57,7 +59,7 @@
       let didWarnAboutReservedActionType = false;
       const originalDispatch = api.dispatch;
       api.dispatch = (...a) => {
-        if (a[0].type === "__setState" && !didWarnAboutReservedActionType) {
+      if ((import.meta.env && import.meta.env.MODE) !== "production" && 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;
         }
@@ -117,7 +119,7 @@
     });
     return initialState;
   };
-}
+const devtools = devtoolsImpl;
 const parseJsonThen = (stringified, f) => {
   let parsed;
   try {
@@ -129,7 +131,7 @@
     f(parsed);
 };
 
-const subscribeWithSelector = (fn) => (set, get, api) => {
+const subscribeWithSelectorImpl = (fn) => (set, get, api) => {
   const origSubscribe = api.subscribe;
   api.subscribe = (selector, optListener, options) => {
     let listener = selector;
@@ -152,8 +154,9 @@
   const initialState = fn(set, get, api);
   return initialState;
 };
+const subscribeWithSelector = subscribeWithSelectorImpl;
 
-const combine = (initialState, create) => (set, get, api) => Object.assign({}, initialState, create(set, get, api));
+const combine = (initialState, create) => (...a) => Object.assign({}, initialState, create(...a));
 
 var __defProp = Object.defineProperty;
 var __getOwnPropSymbols = Object.getOwnPropertySymbols;
@@ -196,14 +199,14 @@
     };
   }
 };
-const persist = (config, baseOptions) => (set, get, api) => {
+const persistImpl = (config, baseOptions) => (set, get, api) => {
   let options = __spreadValues({
     getStorage: () => localStorage,
     serialize: JSON.stringify,
     deserialize: JSON.parse,
     partialize: (state) => state,
     version: 0,
-    merge: (persistedState, currentState) => __spreadValues(__spreadValues({}, currentState), persistedState)
+    merge: (persistedState, currentState) => __spreadValues(__spreadValues({}, currentState), typeof persistedState === "object" ? persistedState : {})
   }, baseOptions);
   let hasHydrated = false;
   const hydrationListeners = /* @__PURE__ */ new Set();
@@ -303,5 +306,15 @@
   hydrate();
   return stateFromStorage || configResult;
 };
+const persist = persistImpl;
+
+const immerImpl = (initializer) => (set, get, store) => {
+  store.setState = (updater, replace) => {
+    const nextState = typeof updater === "function" ? produce(updater) : updater;
+    return set(nextState, replace);
+  };
+  return initializer(store.setState, get, store);
+};
+const immer = immerImpl;
 
-export { combine, devtools, persist, redux, subscribeWithSelector };
+export { combine, devtools, immer, persist, redux, subscribeWithSelector };
diff -urw dist-v4/esm/vanilla.js dist/esm/vanilla.js
--- dist-v4/esm/vanilla.js	2022-04-07 10:42:15.000000000 +0900
+++ dist/esm/vanilla.js	2022-04-07 10:43:11.000000000 +0900
@@ -1,4 +1,4 @@
-function createStore(createState) {
+const createStoreImpl = (createState) => {
   let state;
   const listeners = /* @__PURE__ */ new Set();
   const setState = (partial, replace) => {
@@ -18,6 +18,7 @@
   const api = { setState, getState, subscribe, destroy };
   state = createState(setState, getState, api);
   return api;
-}
+};
+const createStore = (createState) => createState ? createStoreImpl(createState) : createStoreImpl;
 
 export { createStore as default };
```

</details>

Looks much cleaner.
devanshj: > What's remaining on your end for this PR?

Adding a recipe for the slices pattern in the typescript md
dai-shi: Please change the base branch to `main`. Be most careful when resolving.
devanshj: There's one persist test failing not sure why
dai-shi: <details>
<summary>diff -urw dist-943af08/esm dist-ee512ec/esm</summary>

```diff
diff -urw dist-943af08/esm/index.js dist-ee512ec/esm/index.js
--- dist-943af08/esm/index.js	2022-04-09 08:04:30.000000000 +0900
+++ dist-ee512ec/esm/index.js	2022-04-09 08:04:00.000000000 +0900
@@ -1,7 +1,7 @@
 import { useDebugValue } from 'react';
 import { useSyncExternalStoreWithSelector } from 'use-sync-external-store/shim/with-selector.js';
 
-function createStore(createState) {
+const createStoreImpl = (createState) => {
   let state;
   const listeners = /* @__PURE__ */ new Set();
   const setState = (partial, replace) => {
@@ -21,18 +21,22 @@
   const api = { setState, getState, subscribe, destroy };
   state = createState(setState, getState, api);
   return api;
-}
+};
+const createStore = (createState) => createState ? createStoreImpl(createState) : createStoreImpl;
+var createStore$1 = createStore;
 
 function useStore(api, selector = api.getState, equalityFn) {
   const slice = useSyncExternalStoreWithSelector(api.subscribe, api.getState, api.getServerState || api.getState, selector, equalityFn);
   useDebugValue(slice);
   return slice;
 }
-function create(createState) {
-  const api = typeof createState === "function" ? createStore(createState) : createState;
+const createImpl = (createState) => {
+  const api = typeof createState === "function" ? createStore$1(createState) : createState;
   const useBoundStore = (selector, equalityFn) => useStore(api, selector, equalityFn);
   Object.assign(useBoundStore, api);
   return useBoundStore;
-}
+};
+const create = (createState) => createState ? createImpl(createState) : createImpl;
+var create$1 = create;
 
-export { create as default, useStore };
+export { create$1 as default, useStore };
diff -urw dist-943af08/esm/middleware.js dist-ee512ec/esm/middleware.js
--- dist-943af08/esm/middleware.js	2022-04-09 08:04:26.000000000 +0900
+++ dist-ee512ec/esm/middleware.js	2022-04-09 08:03:56.000000000 +0900
@@ -1,3 +1,5 @@
+import { produce } from 'immer';
+
 var __defProp$1 = Object.defineProperty;
 var __getOwnPropSymbols$1 = Object.getOwnPropertySymbols;
 var __hasOwnProp$1 = Object.prototype.hasOwnProperty;
@@ -14,7 +16,7 @@
     }
   return a;
 };
-const redux = (reducer, initial) => (set, get, api) => {
+const reduxImpl = (reducer, initial) => (set, get, api) => {
   api.dispatch = (action) => {
     set((state) => reducer(state, action), false, action);
     return action;
@@ -22,9 +24,9 @@
   api.dispatchFromDevtools = true;
   return __spreadValues$1({ dispatch: (...a) => api.dispatch(...a) }, initial);
 };
+const redux = reduxImpl;
 
-function devtools(fn, options) {
-  return (set, get, api) => {
+const devtoolsImpl = (fn, options) => (set, get, api) => {
     const devtoolsOptions = options === void 0 ? {} : typeof options === "string" ? { name: options } : options;
     let extensionConnector;
     try {
@@ -117,7 +119,7 @@
     });
     return initialState;
   };
-}
+const devtools = devtoolsImpl;
 const parseJsonThen = (stringified, f) => {
   let parsed;
   try {
@@ -129,7 +131,7 @@
     f(parsed);
 };
 
-const subscribeWithSelector = (fn) => (set, get, api) => {
+const subscribeWithSelectorImpl = (fn) => (set, get, api) => {
   const origSubscribe = api.subscribe;
   api.subscribe = (selector, optListener, options) => {
     let listener = selector;
@@ -152,8 +154,9 @@
   const initialState = fn(set, get, api);
   return initialState;
 };
+const subscribeWithSelector = subscribeWithSelectorImpl;
 
-const combine = (initialState, create) => (set, get, api) => Object.assign({}, initialState, create(set, get, api));
+const combine = (initialState, create) => (...a) => Object.assign({}, initialState, create(...a));
 
 var __defProp = Object.defineProperty;
 var __getOwnPropSymbols = Object.getOwnPropertySymbols;
@@ -196,7 +199,7 @@
     };
   }
 };
-const persist = (config, baseOptions) => (set, get, api) => {
+const persistImpl = (config, baseOptions) => (set, get, api) => {
   let options = __spreadValues({
     getStorage: () => localStorage,
     serialize: JSON.stringify,
@@ -264,8 +267,7 @@
         }
       }
     }).then((migratedState) => {
-      var _a2;
-      stateFromStorage = options.merge(migratedState, (_a2 = get()) != null ? _a2 : configResult);
+      stateFromStorage = options.merge(migratedState, configResult);
       set(stateFromStorage, true);
       return setItem();
     }).then(() => {
@@ -304,5 +306,15 @@
   hydrate();
   return stateFromStorage || configResult;
 };
+const persist = persistImpl;
+
+const immerImpl = (initializer) => (set, get, store) => {
+  store.setState = (updater, replace) => {
+    const nextState = typeof updater === "function" ? produce(updater) : updater;
+    return set(nextState, replace);
+  };
+  return initializer(store.setState, get, store);
+};
+const immer = immerImpl;
 
-export { combine, devtools, persist, redux, subscribeWithSelector };
+export { combine, devtools, immer, persist, redux, subscribeWithSelector };
diff -urw dist-943af08/esm/vanilla.js dist-ee512ec/esm/vanilla.js
--- dist-943af08/esm/vanilla.js	2022-04-09 08:04:25.000000000 +0900
+++ dist-ee512ec/esm/vanilla.js	2022-04-09 08:03:54.000000000 +0900
@@ -1,4 +1,4 @@
-function createStore(createState) {
+const createStoreImpl = (createState) => {
   let state;
   const listeners = /* @__PURE__ */ new Set();
   const setState = (partial, replace) => {
@@ -18,6 +18,7 @@
   const api = { setState, getState, subscribe, destroy };
   state = createState(setState, getState, api);
   return api;
-}
+};
+const createStore = (createState) => createState ? createStoreImpl(createState) : createStoreImpl;
 
 export { createStore as default };
```

</details>

> There's one persist test failing not sure why

You seem to fail merging the change in #894.
dai-shi: <details>
<summary>diff -urw dist-943af08/esm dist-160f11d/esm</summary>

```diff 
diff -urw dist-943af08/esm/index.js dist-160f11d/esm/index.js
--- dist-943af08/esm/index.js	2022-04-09 08:04:30.000000000 +0900
+++ dist-160f11d/esm/index.js	2022-04-10 09:15:04.000000000 +0900
@@ -1,7 +1,7 @@
 import { useDebugValue } from 'react';
 import { useSyncExternalStoreWithSelector } from 'use-sync-external-store/shim/with-selector.js';
 
-function createStore(createState) {
+const createStoreImpl = (createState) => {
   let state;
   const listeners = /* @__PURE__ */ new Set();
   const setState = (partial, replace) => {
@@ -21,18 +21,22 @@
   const api = { setState, getState, subscribe, destroy };
   state = createState(setState, getState, api);
   return api;
-}
+};
+const createStore = (createState) => createState ? createStoreImpl(createState) : createStoreImpl;
+var createStore$1 = createStore;
 
 function useStore(api, selector = api.getState, equalityFn) {
   const slice = useSyncExternalStoreWithSelector(api.subscribe, api.getState, api.getServerState || api.getState, selector, equalityFn);
   useDebugValue(slice);
   return slice;
 }
-function create(createState) {
-  const api = typeof createState === "function" ? createStore(createState) : createState;
+const createImpl = (createState) => {
+  const api = typeof createState === "function" ? createStore$1(createState) : createState;
   const useBoundStore = (selector, equalityFn) => useStore(api, selector, equalityFn);
   Object.assign(useBoundStore, api);
   return useBoundStore;
-}
+};
+const create = (createState) => createState ? createImpl(createState) : createImpl;
+var create$1 = create;
 
-export { create as default, useStore };
+export { create$1 as default, useStore };
diff -urw dist-943af08/esm/middleware.js dist-160f11d/esm/middleware.js
--- dist-943af08/esm/middleware.js	2022-04-09 08:04:26.000000000 +0900
+++ dist-160f11d/esm/middleware.js	2022-04-10 09:15:01.000000000 +0900
@@ -1,3 +1,5 @@
+import { produce } from 'immer';
+
 var __defProp$1 = Object.defineProperty;
 var __getOwnPropSymbols$1 = Object.getOwnPropertySymbols;
 var __hasOwnProp$1 = Object.prototype.hasOwnProperty;
@@ -14,7 +16,7 @@
     }
   return a;
 };
-const redux = (reducer, initial) => (set, get, api) => {
+const reduxImpl = (reducer, initial) => (set, get, api) => {
   api.dispatch = (action) => {
     set((state) => reducer(state, action), false, action);
     return action;
@@ -22,9 +24,9 @@
   api.dispatchFromDevtools = true;
   return __spreadValues$1({ dispatch: (...a) => api.dispatch(...a) }, initial);
 };
+const redux = reduxImpl;
 
-function devtools(fn, options) {
-  return (set, get, api) => {
+const devtoolsImpl = (fn, options) => (set, get, api) => {
     const devtoolsOptions = options === void 0 ? {} : typeof options === "string" ? { name: options } : options;
     let extensionConnector;
     try {
@@ -117,7 +119,7 @@
     });
     return initialState;
   };
-}
+const devtools = devtoolsImpl;
 const parseJsonThen = (stringified, f) => {
   let parsed;
   try {
@@ -129,7 +131,7 @@
     f(parsed);
 };
 
-const subscribeWithSelector = (fn) => (set, get, api) => {
+const subscribeWithSelectorImpl = (fn) => (set, get, api) => {
   const origSubscribe = api.subscribe;
   api.subscribe = (selector, optListener, options) => {
     let listener = selector;
@@ -152,8 +154,9 @@
   const initialState = fn(set, get, api);
   return initialState;
 };
+const subscribeWithSelector = subscribeWithSelectorImpl;
 
-const combine = (initialState, create) => (set, get, api) => Object.assign({}, initialState, create(set, get, api));
+const combine = (initialState, create) => (...a) => Object.assign({}, initialState, create(...a));
 
 var __defProp = Object.defineProperty;
 var __getOwnPropSymbols = Object.getOwnPropertySymbols;
@@ -196,7 +199,7 @@
     };
   }
 };
-const persist = (config, baseOptions) => (set, get, api) => {
+const persistImpl = (config, baseOptions) => (set, get, api) => {
   let options = __spreadValues({
     getStorage: () => localStorage,
     serialize: JSON.stringify,
@@ -304,5 +307,15 @@
   hydrate();
   return stateFromStorage || configResult;
 };
+const persist = persistImpl;
+
+const immerImpl = (initializer) => (set, get, store) => {
+  store.setState = (updater, replace) => {
+    const nextState = typeof updater === "function" ? produce(updater) : updater;
+    return set(nextState, replace);
+  };
+  return initializer(store.setState, get, store);
+};
+const immer = immerImpl;
 
-export { combine, devtools, persist, redux, subscribeWithSelector };
+export { combine, devtools, immer, persist, redux, subscribeWithSelector };
diff -urw dist-943af08/esm/vanilla.js dist-160f11d/esm/vanilla.js
--- dist-943af08/esm/vanilla.js	2022-04-09 08:04:25.000000000 +0900
+++ dist-160f11d/esm/vanilla.js	2022-04-10 09:15:00.000000000 +0900
@@ -1,4 +1,4 @@
-function createStore(createState) {
+const createStoreImpl = (createState) => {
   let state;
   const listeners = /* @__PURE__ */ new Set();
   const setState = (partial, replace) => {
@@ -18,6 +18,7 @@
   const api = { setState, getState, subscribe, destroy };
   state = createState(setState, getState, api);
   return api;
-}
+};
+const createStore = (createState) => createState ? createStoreImpl(createState) : createStoreImpl;
 
 export { createStore as default };
```

</details>

Looks perfect!
dai-shi: One remaining issue to be resolved: https://github.com/pmndrs/zustand/pull/725#discussion_r846546318

Anything else?
devanshj: > Anything else?

There are a few small type tweaks pending, I'll continue tomorrow.
devanshj: Okay so let me explain the recent commits a bit...

> f0d192a make immer more generic

This one is fairly simple, we're making less assumptions about the the "parent" `setState` so now we're forwarding the 2-n arguments. So for example previously `create(immer(devtools(() => ...))` would have not worked but now it'd work. 

The `SkipTwo` type is rather complex because a basic one like this...
```ts
type SkipTwoBasic<T> =
  T extends [] ? [] :
  T extends [unknown] ? [] :
  T extends [unknown, unknown, ...infer A] ? A :
  never
```
... won't work because the way subtyping works ie `["a"?] extends [unknown] ? true : false` is actually `false` instead of true. So `SkipTwoBasic<[{ count: number }, boolean?, string?]>` would return `never` instead of `[string?]`. So the actual `SkipTwo` that's committed gives `[string?]`.

> 58dc684 make devtools replace-friendly and more correctly typed

I don't know if you remember but this PR actually bring a new feature that whenever users are replacing the state, they'll have to pass the whole state, so now we have more type safety. I also added a test for this: https://github.com/pmndrs/zustand/pull/725/commits/061a057e8f050b596d167afbd23c806645258229.

But the problem is this feature only work if there are no higher kinded mutators that change `setState` and `devtools` does that...

```ts
declare module '../vanilla' {
  // eslint-disable-next-line @typescript-eslint/no-unused-vars
  interface StoreMutators<S, A> {
    'zustand/devtools': WithDevtools<S>
  }
}

type WithDevtools<S> =
  Write<Cast<S, object>, StoreSetStateWithAction<S>>

type StoreSetStateWithAction<S> =
  S extends { setState: (...a: infer A) => infer R }
    ? { setState: (...a: [...a: A, actionType?: string | { type: unknown }]) => R }
    : never

type Write<T extends object, U extends object> = Omit<T, keyof U> & U
type Cast<T, U> = T extends U ? T : U
```

Here the problem is that when you use `infer` to extract the arguments, it looses all type parameters (ie `Nt` and `R` on `setState` which make the replace feature work) so hence this replace feature wouldn't work if someone is using devtools middleware.

Unfortunately this is a TypeScript limitation, and there's no workaround to this. Hence what we do now is add special cases for the type of setStates we know, ie vanilla and immer. So now `StoreSetStateWithAction` would look like this...

```ts
type StoreSetStateWithAction<S> =
  S extends { getState: () => infer T, setState: (...a: infer A) => infer R }
    ? A extends [Partial<T> | ((state: T) => Partial<T>), (boolean | undefined)?]
        ? /* redefining vanilla setState but now with `actionType` */ :
      A extends [Partial<T> | ((state: Draft<T>) => void), (boolean | undefined)?]
        ? /* redefining immer setState but now with `actionType` */ :
      { setState: (...a: [...a: A, actionType?: string | { type: unknown }]) => R }
      // fallback, this would lose replace-aware feature
    : never
```

And now the replace feature would work even for devtools (with vanilla or immer). You can see I've added those tests for it in `types.test.ts` in below to the vanilla replace test.

Additionally, `[...a: A, actionType?: string | { type: unknown }]` is actually wrong because `A` can have more than two elements so now we do  `[...a: TakeTwo<A>, actionType?: string | { type: unknown }]`.

> 374a5f0 migrate tests

Ha we caught some bugs in the tests because of this additionally type safety. Previously as I said devtools wouldn't have the replace feature but now it does, so we caught passing partial state when we were replacing.

> 3063df1 make setState bivariant to make the state covariant

Okay now this one is a bit lengthy. And it's time to sleep so I'll explain tomorrow xD

I know the types have gotten a bit more complex but I think it's worth it, we've already come this far so a little more correctness is doesn't harm.




dai-shi: > I understand I'm being a perfectionist here, but that's who I am :P

well, I know, but keep in mind that if you make things too complicated and others can't follow, we will eventually remove that code in the future abandoning your work and effort. Non-understandable code is worse than non-perfect code. So, you should find the balance between perfectness and understandability. tbh, you already went too far.

> The SkipTwo type is rather complex because a basic one like this

SkipTwo is complex, but understandable because the goal is clear.

> may have some middleware which has a special setState that returns something

I actually like it. > `return r`

> Unfortunately this is a TypeScript limitation, and there's no workaround to this. Hence what we do now is add special cases for the types of setStates that we know of, ie vanilla and immer.

Can we give this up, and make things less complicated? Using `replace` is not a very common case, so putting more effort and making things complex don't give much benefit overall. I'd like to keep this limitation if there's a solution with manual typing (no type inference).


devanshj: I understand what you're getting at, but I don't think correctness of code and understandability of code are always mutually exclusive. It depends, here the assumption is the maintainers have medium-to-high typescript knowledge (and have my detailed documentations) and the type-level code requires zero maintenance (apart from future bugs which are unlikely and I'll mostly be available), so we can have both. Of course if these conditions are not met then it's perfectly reasonable to delete my code, I don't mind that nor do I think my work goes in vain.

> Can we give this up, and make things less complicated? Using replace is not a very common case, so putting more effort and making things complex don't give much benefit overall. I'd like to keep this limitation if there's a solution with manual typing (no type inference).

Keeping the limitation would not reduce the complexity by any significant margin, it'll only remove the two added branches (ie vanilla and immer) in `StoreSetStateWithAction`, that's it. I don't think it makes sense to come so far (with things like #710) and not go an inch more.

The alternative is that users will have to be careful themselves and pass the whole state when replacing, the compiler won't help them. I'd suggest not to compromise, but you can take the final call.
dai-shi: > here the assumption is the maintainers have medium-to-high typescript knowledge (and have my detailed documentations)

You are too optimistic about this. The maintainers skills are not high enough as you expect to modify those code (and understand your detailed documents). Moreover, other contributors/users can't read the code. How would they become confident with the code they don't understand? It's like a different programming language for someone.

> it'll only remove the two added branches (ie vanilla and immer) in StoreSetStateWithAction, that's it.

I know, and this means if someone creates a new middleware at their own, the type inference will break, right? Creating own middleware and using `replace` are both rare, though.

---

I don't have any solution. Just feels like it's too complicated, and if you have any idea to compromise, it's worth considering.
dai-shi: > 3063df1 make setState bivariant to make the state covariant

I don't fully understand invariant/covariant things, but the commit is readable. If there are tests to cover, I could follow it.
devanshj: > if someone creates a new middleware at their own, the type inference will break, right?

No, here's the scenario..

- **If** user creates a middleware (let's call it `myMiddleware`)...
- ...**and** the they mutate `setState` in it (like `devtools` and `immer` do)...
- ...**and** they use it with `devtools`...
- ...**and** they write `myMiddleware(devtools(...` instead of `devtools(myMiddleware(...`
- ...**and** they pass replace as `true`...
- ...**and** they forget to pass the whole state...
- ...**then** it'll still compile...
- ...**instead of** not compiling.

So it won't break anything, the type inference will still work, everything will be the same, just the types will become less strict.

Let me clarify the typescript limitation a bit...

```ts
type State = { count: number }

declare const setState:
  <R extends boolean | undefined>
    ( partial: R extends true ? State : Partial<State>
    , replace?: R
    ) =>
      void

setState({}, true)
// doesn't compile, great, it's strict and prevents runtime exceptions

declare const withDevtools:
  <F extends (...a: never[]) => unknown>
    (f: F) =>
      // code same as the fallback branch in `StoreSetStateWithAction`
      F extends (...a: infer A) => infer R
        ? (...a: [...a0: TakeTwo<A>, name?: string]) => R
        : never
    
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

const setStateWithDevtools = withDevtools(setState)
// all we wanted to do is add an argument at 2nd index
// we did that but now we also lost all generics
// (hover over `setStateWithDevtools`)

setStateWithDevtools({}, true)
// compiles, not that great as it's not strict and doesn't prevent runtime exceptions
```

---

> You are too optimistic about this. The maintainers skills are not high enough as you expect to modify those code (and understand your detailed documents). Moreover, other contributors/users can't read the code. How would they become confident with the code they don't understand? It's like a different programming language for someone.

> I don't have any solution. Just feels like it's too complicated, and if you have any idea to compromise, it's worth considering.

I just want to make sure you realize there's no pressure from my side, only suggestions. To quote from #662's description...

> You can keep as much code as you want.

My job is to offer you the most correct code possible. Your job is keep whatever code you like, which means I tailor the code according to your needs. I can at best suggest you, the decision will be yours.

So here given you find `StoreSetStateWithAction` complicated (I guess that's what you're referring to) then I can make it less complicated for you, the price is you'll have to give up the strict replace feature.

I just committed "fb624ac remove strict replace" you can check the change in code and tests.

You have two options here:
- Keep this commit
- Revert this commit

The first option takes away the strict replace feature for making the code less complicated. The second option makes the code more complicated for the strict replace feature.

My suggestion is still this...

> Keeping the limitation would not reduce the complexity by any significant margin, it'll only remove the two added branches (ie vanilla and immer) in `StoreSetStateWithAction`, that's it. I don't think it makes sense to come so far (with things like https://github.com/pmndrs/zustand/issues/710) and not go an inch more.

...and it's only a suggestion, really. You should do what is right according to you, because everyone has different OSS philosophies, different judgements of the risk-to-reward ratio, different responsibilities, etc. Any my apologies if my suggestions came across as some kind of pressure to accept them.

Also if the code is still complicated for you (given you tell me which part) then I can still make more compromises and make it less complicated, feel free to tell me.



dai-shi: > ...then it'll still compile...
> ...instead of not compiling.

I see. Yeah, this applies to `replace`. (Technically, `setState` can be mutated to break compiling?)

> ...and they write myMiddleware(devtools(... instead of devtools(myMiddleware(...

I missed this part. We can suggest using devtools at leaf level.

> So here given you find StoreSetStateWithAction complicated (I guess that's what you're referring to)

It was the start of this discussion, but in general, it's too complicated to be too good.

> Also if the code is still complicated for you (given you tell me which part)

I'll see the code again. But, probably more I read, more I get how it works, and "wow, it's good".

> You have two options here:

From the lib maintaining perspective, adding feature (including strict checks) is easier than removing it. So, I will keep fb624ac commit in this PR. Alternatively, as having an option is good, I will open a new PR for adding "strict replace" and I will consider merging after releasing v4 and getting more feedback. Or, maybe you want to open the PR, because it's clearer about your work. Would you?


devanshj: > Technically, `setState` can be mutated to break compiling?

No not if the middleware author writes correct types (just like we write correct types for `devtools` and `immer`).

> I missed this part. We can suggest using devtools at leaf level.

Yes but the other way round (as I wrote "`myMiddleware(devtools(...` instead of `devtools(myMiddleware(...`" so former is problematic).

> But, probably more I read, more I get how it works, and "wow, it's good".

Ha thanks!

> From the lib maintaining perspective, adding feature (including strict checks) is easier than removing it. So, I will keep https://github.com/pmndrs/zustand/commit/fb624ac4042f5cb84de4440128bcf7cebafeab03 commit in this PR. Alternatively, as having an option is good, I will open a new PR for adding "strict replace" and I will consider merging after releasing v4 and getting more feedback. Or, maybe you want to open the PR, because it's clearer about your work. Would you?

No you got it backwards. Adding a strict check is a breaking change, so it's easier to remove it than to add it. So if you publish v4 without strict replace then you can only add it in v5 as it's a breaking change. But if you publish v4 with strict replace then you can remove it in v4.x because it won't be a breaking change. Again, idm anything just correcting you.

And sure I can open a PR if you wish.
dai-shi: > Yes but the other way round

oops, "at the top level", i see.

> And sure I can open a PR if you wish.

Yes, please. I will (re-)consider about it after releasing v4-rc.0 with this PR before releasing final v4.
devanshj: LGTM
dai-shi: <details>
<summary>diff -urw dist-76eeacb/esm dist-9d6e7e2/esm</summary>

```diff
diff -urw dist-76eeacb/esm/index.js dist-9d6e7e2/esm/index.js
--- dist-76eeacb/esm/index.js	2022-04-18 08:56:01.000000000 +0900
+++ dist-9d6e7e2/esm/index.js	2022-04-18 08:57:05.000000000 +0900
@@ -8,11 +8,13 @@
   useDebugValue(slice);
   return slice;
 }
-function create(createState) {
+const createImpl = (createState) => {
   const api = typeof createState === "function" ? createStore__default(createState) : createState;
   const useBoundStore = (selector, equalityFn) => useStore(api, selector, equalityFn);
   Object.assign(useBoundStore, api);
   return useBoundStore;
-}
+};
+const create = (createState) => createState ? createImpl(createState) : createImpl;
+var create$1 = create;
 
-export { create as default, useStore };
+export { create$1 as default, useStore };
diff -urw dist-76eeacb/esm/middleware.js dist-9d6e7e2/esm/middleware.js
--- dist-76eeacb/esm/middleware.js	2022-04-18 08:55:57.000000000 +0900
+++ dist-9d6e7e2/esm/middleware.js	2022-04-18 08:57:01.000000000 +0900
@@ -1,3 +1,5 @@
+import { produce } from 'immer';
+
 var __defProp$1 = Object.defineProperty;
 var __getOwnPropSymbols$1 = Object.getOwnPropertySymbols;
 var __hasOwnProp$1 = Object.prototype.hasOwnProperty;
@@ -14,7 +16,7 @@
     }
   return a;
 };
-const redux = (reducer, initial) => (set, get, api) => {
+const reduxImpl = (reducer, initial) => (set, get, api) => {
   api.dispatch = (action) => {
     set((state) => reducer(state, action), false, action);
     return action;
@@ -22,9 +24,9 @@
   api.dispatchFromDevtools = true;
   return __spreadValues$1({ dispatch: (...a) => api.dispatch(...a) }, initial);
 };
+const redux = reduxImpl;
 
-function devtools(fn, options) {
-  return (set, get, api) => {
+const devtoolsImpl = (fn, options) => (set, get, api) => {
     const devtoolsOptions = options === void 0 ? {} : typeof options === "string" ? { name: options } : options;
     let extensionConnector;
     try {
@@ -40,10 +42,11 @@
     const extension = extensionConnector.connect(devtoolsOptions);
     let isRecording = true;
     api.setState = (state, replace, nameOrAction) => {
-      set(state, replace);
+    const r = set(state, replace);
       if (!isRecording)
-        return;
+      return r;
       extension.send(nameOrAction === void 0 ? { type: devtoolsOptions.anonymousActionType || "anonymous" } : typeof nameOrAction === "string" ? { type: nameOrAction } : nameOrAction, get());
+    return r;
     };
     const setStateFromDevtools = (...a) => {
       const originalIsRecording = isRecording;
@@ -117,7 +120,7 @@
     });
     return initialState;
   };
-}
+const devtools = devtoolsImpl;
 const parseJsonThen = (stringified, f) => {
   let parsed;
   try {
@@ -129,7 +132,7 @@
     f(parsed);
 };
 
-const subscribeWithSelector = (fn) => (set, get, api) => {
+const subscribeWithSelectorImpl = (fn) => (set, get, api) => {
   const origSubscribe = api.subscribe;
   api.subscribe = (selector, optListener, options) => {
     let listener = selector;
@@ -152,8 +155,9 @@
   const initialState = fn(set, get, api);
   return initialState;
 };
+const subscribeWithSelector = subscribeWithSelectorImpl;
 
-const combine = (initialState, create) => (set, get, api) => Object.assign({}, initialState, create(set, get, api));
+const combine = (initialState, create) => (...a) => Object.assign({}, initialState, create(...a));
 
 var __defProp = Object.defineProperty;
 var __getOwnPropSymbols = Object.getOwnPropertySymbols;
@@ -196,7 +200,7 @@
     };
   }
 };
-const persist = (config, baseOptions) => (set, get, api) => {
+const persistImpl = (config, baseOptions) => (set, get, api) => {
   let options = __spreadValues({
     getStorage: () => localStorage,
     serialize: JSON.stringify,
@@ -304,5 +308,15 @@
   hydrate();
   return stateFromStorage || configResult;
 };
+const persist = persistImpl;
+
+const immerImpl = (initializer) => (set, get, store) => {
+  store.setState = (updater, replace, ...a) => {
+    const nextState = typeof updater === "function" ? produce(updater) : updater;
+    return set(nextState, replace, ...a);
+  };
+  return initializer(store.setState, get, store);
+};
+const immer = immerImpl;
 
-export { combine, devtools, persist, redux, subscribeWithSelector };
+export { combine, devtools, immer, persist, redux, subscribeWithSelector };
diff -urw dist-76eeacb/esm/vanilla.js dist-9d6e7e2/esm/vanilla.js
--- dist-76eeacb/esm/vanilla.js	2022-04-18 08:55:56.000000000 +0900
+++ dist-9d6e7e2/esm/vanilla.js	2022-04-18 08:57:00.000000000 +0900
@@ -1,4 +1,4 @@
-function createStore(createState) {
+const createStoreImpl = (createState) => {
   let state;
   const listeners = /* @__PURE__ */ new Set();
   const setState = (partial, replace) => {
@@ -18,6 +18,7 @@
   const api = { setState, getState, subscribe, destroy };
   state = createState(setState, getState, api);
   return api;
-}
+};
+const createStore = (createState) => createState ? createStoreImpl(createState) : createStoreImpl;
 
 export { createStore as default };
```

</details>

JS change seems minimal, but realized that `zustand/middleware` shouldn't depend on `immer`. I will create a follow-up PR to separate it.

Patch:
diff --git a/docs/typescript.md b/docs/typescript.md
new file mode 100644
index 0000000..af599af
--- /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 0000000..be8c03a
--- /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 cfd7a22..7c8cf4c 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 f5266a4..44d0645 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 9064282..c1787ba 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 42899b5..8ca4ecb 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 72caf95..f19f93f 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 0000000..7c302ca
--- /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 49f1436..650da0a 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 8148b75..8bffda9 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 171acaf..70ba705 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 909b2fd..e5af6c7 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 3cdd969..4c7996f 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


Test patch:
diff --git a/tests/context.test.tsx b/tests/context.test.tsx
index 7427cf5..ff2f0af 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 6f9dadf..8805553 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 afe02f6..6e29d1c 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 3b888f1..30b2d7c 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 6ef1bb4..798c987 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 d088af2..69aaf23 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
+})