feat(vanilla): non-object state (pmndrs#1144)

* feat: non-object state

* add test

* prefer unknown

* fix types for persist and subscribeWithSelector

* no Cast in persist and subscribeWithSelector

* simplify immer type

* simplify devtools type

* simplify redux type

* simplify type with looser action

* fix StoreApi type parameter

* fix types

src/context.ts

@@ -8,7 +8,7 @@ import {
 } from 'react'
 import { StoreApi, useStore } from 'zustand'
 
-type UseContextStore<S extends StoreApi> = {
+type UseContextStore<S extends StoreApi<unknown>> = {
   (): ExtractState<S>
   <U>(
     selector: (state: ExtractState<S>) => U,
@@ -20,7 +20,7 @@ type ExtractState<S> = S extends { getState: () => infer T } ? T : never
 
 type WithoutCallSignature<T> = { [K in keyof T]: T[K] }
 
-function createContext<S extends StoreApi>() {
+function createContext<S extends StoreApi<unknown>>() {
   const ZustandContext = reactCreateContext<S | undefined>(undefined)
 
   const Provider = ({

src/middleware/devtools.ts

@@ -15,8 +15,7 @@ type Message = {
   state?: any
 }
 
-type Write<T extends object, U extends object> = Omit<T, keyof U> & U
-type Cast<T, U> = T extends U ? T : U
+type Write<T, U> = Omit<T, keyof U> & U
 type TakeTwo<T> = T extends []
   ? [undefined, undefined]
   : T extends [unknown]
@@ -37,7 +36,7 @@ type TakeTwo<T> = T extends []
   ? [A0?, A1?]
   : never
 
-type WithDevtools<S> = Write<Cast<S, object>, StoreDevtools<S>>
+type WithDevtools<S> = Write<S, StoreDevtools<S>>
 
 type StoreDevtools<S> = S extends {
   setState: (...a: infer Sa) => infer Sr
@@ -49,6 +48,9 @@ type StoreDevtools<S> = S extends {
     }
   : never
 
+const isObjectWithTypeProperty = (x: unknown): x is { type: unknown } =>
+  x !== null && typeof x === 'object' && 'type' in x
+
 export interface DevtoolsOptions {
   enabled?: boolean
   anonymousActionType?: string
@@ -69,7 +71,7 @@ export interface DevtoolsOptions {
 }
 
 type Devtools = <
-  T extends object,
+  T,
   Mps extends [StoreMutatorIdentifier, unknown][] = [],
   Mcs extends [StoreMutatorIdentifier, unknown][] = []
 >(
@@ -84,7 +86,7 @@ declare module '../vanilla' {
   }
 }
 
-type DevtoolsImpl = <T extends object>(
+type DevtoolsImpl = <T>(
   storeInitializer: PopArgument<StateCreator<T, [], []>>,
   devtoolsOptions?: DevtoolsOptions
 ) => PopArgument<StateCreator<T, [], []>>
@@ -95,7 +97,7 @@ type PopArgument<T extends (...a: never[]) => unknown> = T extends (
   ? (...a: A) => R
   : never
 
-export type NamedSet<T extends object> = WithDevtools<StoreApi<T>>['setState']
+export type NamedSet<T> = WithDevtools<StoreApi<T>>['setState']
 
 const devtoolsImpl: DevtoolsImpl =
   (fn, devtoolsOptions = {}) =>
@@ -132,9 +134,9 @@ const devtoolsImpl: DevtoolsImpl =
       extension.send(
         nameOrAction === undefined
           ? { type: anonymousActionType || 'anonymous' }
-          : typeof nameOrAction === 'string'
-          ? { type: nameOrAction }
-          : nameOrAction,
+          : isObjectWithTypeProperty(nameOrAction)
+          ? nameOrAction
+          : { type: nameOrAction },
         get()
       )
       return r

src/middleware/immer.ts

@@ -3,7 +3,7 @@ import { Draft, produce } from 'immer'
 import { StateCreator, StoreMutatorIdentifier } from '../vanilla'
 
 type Immer = <
-  T extends object,
+  T,
   Mps extends [StoreMutatorIdentifier, unknown][] = [],
   Mcs extends [StoreMutatorIdentifier, unknown][] = []
 >(
@@ -17,8 +17,7 @@ declare module '../vanilla' {
   }
 }
 
-type Write<T extends object, U extends object> = Omit<T, keyof U> & U
-type Cast<T, U> = T extends U ? T : U
+type Write<T, U> = Omit<T, keyof U> & U
 type SkipTwo<T> = T extends []
   ? []
   : T extends [unknown]
@@ -33,7 +32,7 @@ type SkipTwo<T> = T extends []
   ? A
   : never
 
-type WithImmer<S> = Write<Cast<S, object>, StoreImmer<S>>
+type WithImmer<S> = Write<S, StoreImmer<S>>
 
 type StoreImmer<S> = S extends {
   getState: () => infer T
@@ -56,7 +55,7 @@ type PopArgument<T extends (...a: never[]) => unknown> = T extends (
   ? (...a: A) => R
   : never
 
-type ImmerImpl = <T extends object>(
+type ImmerImpl = <T>(
   storeInitializer: PopArgument<StateCreator<T, [], []>>
 ) => PopArgument<StateCreator<T, [], []>>
 

src/middleware/persist.ts

@@ -72,7 +72,7 @@ export interface PersistOptions<S, PersistedState = S> {
 
 type PersistListener<S> = (state: S) => void
 
-type StorePersist<S extends object, Ps> = {
+type StorePersist<S, Ps> = {
   persist: {
     setOptions: (options: Partial<PersistOptions<S, Ps>>) => void
     clearStorage: () => void
@@ -297,7 +297,7 @@ const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
 }
 
 type Persist = <
-  T extends object,
+  T,
   Mps extends [StoreMutatorIdentifier, unknown][] = [],
   Mcs extends [StoreMutatorIdentifier, unknown][] = [],
   U = T
@@ -312,14 +312,13 @@ declare module '../vanilla' {
   }
 }
 
-type Write<T extends object, U extends object> = Omit<T, keyof U> & U
-type Cast<T, U> = T extends U ? T : U
+type Write<T, U> = Omit<T, keyof U> & U
 
 type WithPersist<S, A> = S extends { getState: () => infer T }
-  ? Write<S, StorePersist<Cast<T, object>, A>>
+  ? Write<S, StorePersist<T, A>>
   : never
 
-type PersistImpl = <T extends object>(
+type PersistImpl = <T>(
   storeInitializer: PopArgument<StateCreator<T, [], []>>,
   options: PersistOptions<T, T>
 ) => PopArgument<StateCreator<T, [], []>>

src/middleware/redux.ts

@@ -1,29 +1,20 @@
 import { StateCreator, StoreMutatorIdentifier } from '../vanilla'
 import { NamedSet } from './devtools'
 
-type Write<T extends object, U extends object> = Omit<T, keyof U> & U
-type Cast<T, U> = T extends U ? T : U
+type Write<T, U> = Omit<T, keyof U> & U
 
-type Action = {
-  type: unknown
-}
-
-type ReduxState<A extends Action> = {
+type ReduxState<A> = {
   dispatch: StoreRedux<A>['dispatch']
 }
 
-type StoreRedux<A extends Action> = {
+type StoreRedux<A> = {
   dispatch: (a: A) => A
   dispatchFromDevtools: true
 }
 
-type WithRedux<S, A> = Write<Cast<S, object>, StoreRedux<Cast<A, Action>>>
+type WithRedux<S, A> = Write<S, StoreRedux<A>>
 
-type Redux = <
-  T extends object,
-  A extends Action,
-  Cms extends [StoreMutatorIdentifier, unknown][] = []
->(
+type Redux = <T, A, Cms extends [StoreMutatorIdentifier, unknown][] = []>(
   reducer: (state: T, action: A) => T,
   initialState: T
 ) => StateCreator<Write<T, ReduxState<A>>, Cms, [['zustand/redux', A]]>
@@ -40,16 +31,23 @@ type PopArgument<T extends (...a: never[]) => unknown> = T extends (
   ? (...a: A) => R
   : never
 
-type ReduxImpl = <T extends object, A extends Action>(
+type ReduxImpl = <T, A>(
   reducer: (state: T, action: A) => T,
   initialState: T
 ) => PopArgument<StateCreator<T & ReduxState<A>, [], []>>
 
+const isObjectWithTypeProperty = (x: unknown): x is { type: unknown } =>
+  x !== null && typeof x === 'object' && 'type' in x
+
 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)
+    ;(set as NamedSet<S>)(
+      (state: S) => reducer(state, action),
+      false,
+      isObjectWithTypeProperty(action) ? action : { type: action }
+    )
     return action
   }
   ;(api as any).dispatchFromDevtools = true

src/middleware/subscribeWithSelector.ts

@@ -1,7 +1,7 @@
 import { StateCreator, StoreMutatorIdentifier } from '../vanilla'
 
 type SubscribeWithSelector = <
-  T extends object,
+  T,
   Mps extends [StoreMutatorIdentifier, unknown][] = [],
   Mcs extends [StoreMutatorIdentifier, unknown][] = []
 >(
@@ -12,11 +12,10 @@ type SubscribeWithSelector = <
   >
 ) => 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 Write<T, U> = Omit<T, keyof U> & U
 
 type WithSelectorSubscribe<S> = S extends { getState: () => infer T }
-  ? Write<S, StoreSubscribeWithSelector<Cast<T, object>>>
+  ? Write<S, StoreSubscribeWithSelector<T>>
   : never
 
 declare module '../vanilla' {
@@ -26,7 +25,7 @@ declare module '../vanilla' {
   }
 }
 
-type StoreSubscribeWithSelector<T extends object> = {
+type StoreSubscribeWithSelector<T> = {
   subscribe: {
     (listener: (selectedState: T, previousSelectedState: T) => void): () => void
     <U>(

src/react.ts

@@ -15,19 +15,21 @@ const { useSyncExternalStoreWithSelector } = useSyncExternalStoreExports
 
 type ExtractState<S> = S extends { getState: () => infer T } ? T : never
 
-type WithReact<S extends StoreApi> = S & {
+type WithReact<S extends StoreApi<unknown>> = S & {
   getServerState?: () => ExtractState<S>
 }
 
-export function useStore<S extends WithReact<StoreApi>>(api: S): ExtractState<S>
+export function useStore<S extends WithReact<StoreApi<unknown>>>(
+  api: S
+): ExtractState<S>
 
-export function useStore<S extends WithReact<StoreApi>, U>(
+export function useStore<S extends WithReact<StoreApi<unknown>>, U>(
   api: S,
   selector: (state: ExtractState<S>) => U,
   equalityFn?: (a: U, b: U) => boolean
 ): U
 
-export function useStore<TState extends object, StateSlice>(
+export function useStore<TState, StateSlice>(
   api: WithReact<StoreApi<TState>>,
   selector: (state: TState) => StateSlice = api.getState as any,
   equalityFn?: (a: StateSlice, b: StateSlice) => boolean
@@ -43,7 +45,7 @@ export function useStore<TState extends object, StateSlice>(
   return slice
 }
 
-export type UseBoundStore<S extends WithReact<StoreApi>> = {
+export type UseBoundStore<S extends WithReact<StoreApi<unknown>>> = {
   (): ExtractState<S>
   <U>(
     selector: (state: ExtractState<S>) => U,
@@ -52,16 +54,16 @@ export type UseBoundStore<S extends WithReact<StoreApi>> = {
 } & S
 
 type Create = {
-  <T extends object, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
+  <T, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
     initializer: StateCreator<T, [], Mos>
   ): UseBoundStore<Mutate<StoreApi<T>, Mos>>
-  <T extends object>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
+  <T>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
     initializer: StateCreator<T, [], Mos>
   ) => UseBoundStore<Mutate<StoreApi<T>, Mos>>
-  <S extends StoreApi>(store: S): UseBoundStore<S>
+  <S extends StoreApi<unknown>>(store: S): UseBoundStore<S>
 }
 
-const createImpl = <T extends object>(createState: StateCreator<T, [], []>) => {
+const createImpl = <T>(createState: StateCreator<T, [], []>) => {
   const api =
     typeof createState === 'function' ? createStore(createState) : createState
 
@@ -73,8 +75,7 @@ const createImpl = <T extends object>(createState: StateCreator<T, [], []>) => {
   return useBoundStore
 }
 
-const create = (<T extends object>(
-  createState: StateCreator<T, [], []> | undefined
-) => (createState ? createImpl(createState) : createImpl)) as Create
+const create = (<T>(createState: StateCreator<T, [], []> | undefined) =>
+  createState ? createImpl(createState) : createImpl) as Create
 
 export default create

src/vanilla.ts

@@ -5,7 +5,7 @@ type SetStateInternal<T> = {
   ): void
 }['_']
 
-export interface StoreApi<T extends object = object> {
+export interface StoreApi<T> {
   setState: SetStateInternal<T>
   getState: () => T
   subscribe: (listener: (state: T, prevState: T) => void) => () => void
@@ -21,7 +21,7 @@ export type Mutate<S, Ms> = Ms extends []
   : never
 
 export type StateCreator<
-  T extends object,
+  T,
   Mis extends [StoreMutatorIdentifier, unknown][] = [],
   Mos extends [StoreMutatorIdentifier, unknown][] = [],
   U = T
@@ -37,17 +37,17 @@ export interface StoreMutators<S, A> {}
 export type StoreMutatorIdentifier = keyof StoreMutators<unknown, unknown>
 
 type CreateStore = {
-  <T extends object, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
+  <T, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
     initializer: StateCreator<T, [], Mos>
   ): Mutate<StoreApi<T>, Mos>
 
-  <T extends object>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
+  <T>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
     initializer: StateCreator<T, [], Mos>
   ) => Mutate<StoreApi<T>, Mos>
 }
 
 type CreateStoreImpl = <
-  T extends object,
+  T,
   Mos extends [StoreMutatorIdentifier, unknown][] = []
 >(
   initializer: StateCreator<T, [], Mos>
@@ -74,9 +74,10 @@ const createStoreImpl: CreateStoreImpl = (createState) => {
         : partial
     if (nextState !== state) {
       const previousState = state
-      state = replace
-        ? (nextState as TState)
-        : Object.assign({}, state, nextState)
+      state =
+        replace ?? typeof nextState !== 'object'
+          ? (nextState as TState)
+          : Object.assign({}, state, nextState)
       listeners.forEach((listener) => listener(state, previousState))
     }
   }

tests/basic.test.tsx

@@ -576,3 +576,25 @@ it('ensures a subscriber is not mistakenly overwritten', async () => {
   expect((await findAllByText('count1: 1')).length).toBe(2)
   expect((await findAllByText('count2: 1')).length).toBe(1)
 })
+
+it('works with non-object state', async () => {
+  const useCount = create(() => 1)
+  const inc = () => useCount.setState((c) => c + 1)
+
+  const Counter = () => {
+    const count = useCount()
+    return (
+      <>
+        <div>count: {count}</div>
+        <button onClick={inc}>button</button>
+      </>
+    )
+  }
+
+  const { getByText, findByText } = render(<Counter />)
+
+  await findByText('count: 1')
+
+  fireEvent.click(getByText('button'))
+  await findByText('count: 2')
+})