Strict replace types

src/middleware/devtools.ts

@@ -1,4 +1,5 @@
 import type {} from '@redux-devtools/extension'
+import type { Draft } from 'immer'
 import {
   PartialState,
   SetState,
@@ -47,13 +48,40 @@ type TakeTwo<T> = T extends []
 type WithDevtools<S> = Write<Cast<S, object>, StoreDevtools<S>>
 
 type StoreDevtools<S> = S extends {
+  getState: () => infer T
   setState: (...a: infer A) => infer Sr
 }
-  ? {
-      setState(
-        ...a: [...a: TakeTwo<A>, actionType?: string | { type: unknown }]
-      ): Sr
-    }
+  ? A extends [
+      T | Partial<T> | ((state: T) => T | Partial<T>),
+      (boolean | undefined)?
+    ]
+    ? {
+        setState<R extends boolean | undefined>(
+          partial:
+            | (R extends true ? T : T | Partial<T>)
+            | ((state: T) => R extends true ? T : T | Partial<T>),
+          replace?: R,
+          action?: string | { type: unknown }
+        ): Sr
+      }
+    : A extends [
+        T | Partial<T> | ((state: Draft<T>) => void),
+        (boolean | undefined)?
+      ]
+    ? {
+        setState<R extends boolean | undefined>(
+          nextStateOrUpdater:
+            | (R extends true ? T : T | Partial<T>)
+            | ((state: Draft<T>) => void),
+          shouldReplace?: R,
+          action?: string | { type: unknown }
+        ): Sr
+      }
+    : {
+        setState(
+          ...a: [...a: TakeTwo<A>, action?: string | { type: unknown }]
+        ): Sr
+      }
   : never
 
 interface DevtoolsOptions {

src/middleware/immer.ts

@@ -37,17 +37,17 @@ type WithImmer<S> = Write<Cast<S, object>, StoreImmer<S>>
 
 type StoreImmer<S> = S extends {
   getState: () => infer T
-  setState: infer SetState
+  setState: (...a: infer A) => infer Sr
 }
-  ? SetState extends (...a: infer A) => infer Sr
-    ? {
-        setState(
-          nextStateOrUpdater: T | Partial<T> | ((state: Draft<T>) => void),
-          shouldReplace?: boolean | undefined,
-          ...a: SkipTwo<A>
-        ): Sr
-      }
-    : never
+  ? {
+      setState<R extends boolean | undefined = false>(
+        nextStateOrUpdater:
+          | (R extends true ? T : T | Partial<T>)
+          | ((state: Draft<T>) => void),
+        shouldReplace?: R,
+        ...a: SkipTwo<A>
+      ): Sr
+    }
   : never
 
 type PopArgument<T extends (...a: never[]) => unknown> = T extends (

src/vanilla.ts

@@ -14,9 +14,11 @@ export type Subscribe<T extends State> = {
 }
 
 export type SetState<T extends State> = {
-  _(
-    partial: T | Partial<T> | ((state: T) => T | Partial<T>),
-    replace?: boolean | undefined
+  _<R extends boolean | undefined = false>(
+    partial:
+      | (R extends true ? T : T | Partial<T>)
+      | ((state: T) => R extends true ? T : T | Partial<T>),
+    replace?: R
   ): void
 }['_']
 export type GetState<T extends State> = () => T

tests/types.test.tsx

@@ -12,6 +12,8 @@ import create, {
   Subscribe,
   UseBoundStore,
 } from 'zustand'
+import { devtools } from 'zustand/middleware'
+import { immer } from 'zustand/middleware/immer'
 
 it('can use exposed types', () => {
   interface ExampleState {
@@ -188,3 +190,71 @@ it('state is covariant', () => {
     baz: string
   }> = store
 })
+
+it('setState with replace requires whole state', () => {
+  const store = create<{ count: number; something: string }>()(() => ({
+    count: 0,
+    something: 'foo',
+  }))
+
+  store.setState(
+    // @ts-expect-error missing `something`
+    { count: 1 },
+    true
+  )
+  store.setState({ count: 1 }, false)
+  store.setState({ count: 1 }, true as boolean)
+
+  const storeWithDevtools = create<{ count: number; something: string }>()(
+    devtools(() => ({
+      count: 0,
+      something: 'foo',
+    }))
+  )
+
+  storeWithDevtools.setState(
+    // @ts-expect-error missing `something`
+    { count: 1 },
+    true
+  )
+  storeWithDevtools.setState({ count: 1 }, false)
+  storeWithDevtools.setState({ count: 1 }, true as boolean)
+
+  storeWithDevtools.setState(
+    // @ts-expect-error missing `something`
+    { count: 1 },
+    true,
+    'FOO'
+  )
+  storeWithDevtools.setState({ count: 1 }, false, 'FOO')
+  storeWithDevtools.setState({ count: 1 }, true as boolean, 'FOO')
+
+  const storeWithDevtoolsAndImmer = create<{
+    count: number
+    something: string
+  }>()(
+    immer(
+      devtools(() => ({
+        count: 0,
+        something: 'foo',
+      }))
+    )
+  )
+
+  storeWithDevtoolsAndImmer.setState(
+    // @ts-expect-error missing `something`
+    { count: 1 },
+    true
+  )
+  storeWithDevtoolsAndImmer.setState({ count: 1 }, false)
+  storeWithDevtoolsAndImmer.setState({ count: 1 }, true as boolean)
+
+  storeWithDevtoolsAndImmer.setState(
+    // @ts-expect-error missing `something`
+    { count: 1 },
+    true,
+    'FOO'
+  )
+  storeWithDevtoolsAndImmer.setState({ count: 1 }, false, 'FOO')
+  storeWithDevtoolsAndImmer.setState({ count: 1 }, true as boolean, 'FOO')
+})