Example from the docs not working

[a-rebets]

Description

I tried to onboard to this library after using the standard immer / zustand combo for quite some time. Apparently adapting my store was not straightforward. Maybe I'm just missing the correct way we should be dealing with mutative-like updates, but docs don't say much. I tried changing my types or playing with store init options but failed to find a fix. I then started trying examples from the docs and ended up with a type error even on simple stuff

Steps

1. Install zustand-x/zustand in a clean TS project
2. Paste the example from docs (comparing with Zustand section)
3. Observe type error

const store = createStore(
    { count: 0 },
    { name: "test", mutative: true })
	.extendSelectors(({ get }) => ({
		// Computed values are auto-memoized
		double: () => get("count") * 2,
	}))
	.extendActions(({ set }) => ({
		increment: () => set("count", (count) => count + 1),
	}));

type error:

Overload 1 of 3, '(key: "state", value: { count: number; } | ((state: DraftedObject<Partial<{ count: number; }>>) => void)): void', gave the following error.
    Argument of type '"count"' is not assignable to parameter of type '"state"'.
  Overload 2 of 3, '(key: "count", value: number): void', gave the following error.
    Argument of type '(count: any) => any' is not assignable to parameter of type 'number'.

so the overload handling a specific field update only accepts a new value, but not a function that depends on the current state

Expectation

It's expected that the set method destructured inside extendActions is typed correctly, i.e. accepts both a function for mutating a draft of the selected field and a raw value