fix: delete root keys with immer middleware

src/middleware/immer.ts

@@ -42,7 +42,6 @@ type StoreImmer<S> = S extends {
     ? {
         setState(
           nextStateOrUpdater: T | Partial<T> | ((state: Draft<T>) => void),
-          shouldReplace?: boolean | undefined,
           ...a: SkipTwo<A>
         ): Sr
       }
@@ -56,12 +55,10 @@ type ImmerImpl = <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>
+  store.setState = (updater, _, ...a) => {
+    const nextState = produce(updater as any) as (s: T) => T
 
-    return set(nextState as any, replace, ...a)
+    return set(nextState as any, /* replace */ true, ...a)
   }
 
   return initializer(store.setState, get, store)

tests/immerMiddleware.test.tsx

@@ -0,0 +1,48 @@
+import { describe, expect, it } from 'vitest'
+import { create } from 'zustand'
+import { immer } from 'zustand/middleware/immer'
+import { createStore } from 'zustand/vanilla'
+
+describe('immer middleware', () => {
+  it('typical case', () => {
+    type CounterState = {
+      count: number
+      inc: () => void
+    }
+
+    const useBoundStore = create<CounterState>()(
+      immer((set, get) => ({
+        count: 0,
+        inc: () =>
+          set((state) => {
+            state.count = get().count + 1
+          }),
+      })),
+    )
+
+    expect(useBoundStore.getState().count).toEqual(0)
+
+    useBoundStore.getState().inc()
+
+    expect(useBoundStore.getState().count).toEqual(1)
+  })
+
+  // Bug: https://github.com/pmndrs/zustand/discussions/2480
+  it('delete key at root of state', () => {
+    type State = Record<string, any>
+
+    const store = createStore<State>()(immer(() => ({})))
+
+    expect(store.getState()).toEqual({})
+
+    store.setState((state) => {
+      state.x = 3
+    })
+
+    expect(store.getState()).toEqual({ x: 3 })
+    store.setState((state) => {
+      delete state.x
+    })
+    expect(store.getState()).toEqual({})
+  })
+})