Merge pull request #62 from chrisvander/fix-immer

fix: immer middleware throwing

Author: chrisvander

bun.lockb

@@ -17,6 +17,7 @@
     "@types/bun": "^1.1.13",
     "@types/node": "^22.9.0",
     "husky": "^9.1.6",
+    "immer": "^10.1.3",
     "react": "^18.3.1",
     "semver": "^7.6.3",
     "typescript": "^5.6.3",

package.json

@@ -1,5 +1,6 @@
 import { beforeEach, describe, expect, mock, test } from "bun:test"
 import { type StateCreator, create } from "zustand"
+import { immer } from "zustand/middleware/immer"
 import { type ComputedStateOpts, createComputed } from "./computed"
 
 type Store = {
@@ -196,3 +197,43 @@ describe("slices pattern", () => {
     expect(computeSliceMock).toHaveBeenCalledTimes(3)
   })
 })
+
+describe("immer middleware functions without throwing", () => {
+  type Store = {
+    count: number
+    inc: () => void
+    dec: () => void
+  }
+
+  type ComputedStore = {
+    countSq: number
+  }
+
+  const computed = createComputed(
+    (state: Store): ComputedStore => ({
+      countSq: state.count ** 2,
+    }),
+    { keys: ["count"] },
+  )
+
+  const useStore = create<Store>()(
+    immer(
+      computed((set) => ({
+        count: 1,
+        inc: () =>
+          set((state) => {
+            // example with Immer middleware
+            state.count += 1
+          }),
+        dec: () => set((state) => ({ count: state.count - 1 })),
+      })),
+    ),
+  )
+
+  expect(() => useStore.getState().inc()).not.toThrow()
+  expect(useStore.getState().count).toEqual(2)
+  expect(useStore.getState().countSq).toEqual(4)
+  expect(() => useStore.getState().dec()).not.toThrow()
+  expect(useStore.getState().count).toEqual(1)
+  expect(useStore.getState().countSq).toEqual(1)
+})

src/computed.test.ts

@@ -22,7 +22,7 @@ export type ComputedStateOpts<T> = (
        * should be *fast* - it determines whether or not you need to
        * recompute.
        */
-      shouldRecompute?: (state: T, nextState: T) => boolean
+      shouldRecompute?: (state: T, nextState: T | Partial<T>) => boolean
     }
 ) & {
   /**
@@ -76,18 +76,14 @@ type ComputedStateImpl = <T extends object, A extends object>(
   opts?: ComputedStateOpts<T>,
 ) => (f: StateCreator<T, [], []>) => StateCreator<T, [], [], T & A>
 
-type SetStateWithArgs = Parameters<ReturnType<ReturnType<ComputedStateImpl>>>[0] extends (...args: infer U) => void
-  ? (...args: [...U, ...unknown[]]) => void
-  : never
-
 const computedImpl: ComputedStateImpl = (compute, opts) => (f) => {
   type T = ReturnType<typeof f>
   type A = ReturnType<typeof compute>
 
   const optsKeys = !opts || !("keys" in opts) || opts.keys == null ? undefined : opts.keys
   const keysSet = optsKeys ? new Set(optsKeys as string[]) : undefined
 
-  function defaultShouldRecomputeFn<T>(_: T, nextState: T): boolean {
+  function defaultShouldRecomputeFn<T>(_: T, nextState: T | Partial<T>): boolean {
     if (!keysSet || nextState == null) return true
     return Object.keys(nextState).some((k) => keysSet.has(k))
   }
@@ -99,39 +95,49 @@ const computedImpl: ComputedStateImpl = (compute, opts) => (f) => {
   return (set, get, api) => {
     const equalityFn = opts?.equalityFn ?? shallow
 
-    const computeAndMerge = (state: T | (T & A)): T & A => {
+    function computeAndMerge(state: T | (T & A)): T & A {
       // Calculate the new computed state.
-      const computedState: A = compute({ ...state })
+      const computedState = compute(state)
 
       // If part of the computed state did not change according to the equalityFn,
       // then delete that key from the newly calculated computed state.
       for (const k of Object.keys(computedState) as (keyof A)[]) {
-        if (equalityFn(computedState[k], (state as T & A)[k])) {
+        if (k in state && equalityFn(computedState[k], (state as T & A)[k])) {
           delete computedState[k]
         }
       }
 
-      return { ...state, ...computedState }
+      return Object.assign(state, computedState)
     }
 
+    const _api = api as Mutate<StoreApi<T>, [["chrisvander/zustand-computed", A]]>
+
     /**
      * Higher level function to handle compute & compare overhead.
      */
-    const setWithComputed = (update: T | ((state: T) => T), replace?: boolean, ...args: unknown[]) => {
-      ;(set as SetStateWithArgs)(
-        (state: T): T & A => {
-          const updated = typeof update === "object" ? update : update(state)
-          if (!shouldRecomputeFn?.(state, updated)) return { ...state, ...updated } as T & A
-          return computeAndMerge({ ...state, ...updated })
-        },
-        replace,
-        ...args,
-      )
+    function setState(partial: T | Partial<T> | ((state: T) => T | Partial<T>), replace?: false): void
+    function setState(state: T | ((state: T) => T), replace: true): void
+    function setState(arg: T | Partial<T> | ((state: T) => T | Partial<T>), replace?: boolean): void {
+      if (replace === false || replace == null) {
+        // Merge the partial state with the current state.
+        set((state) => {
+          const newState = typeof arg === "function" ? arg(state) : arg
+          if (!shouldRecomputeFn(state, newState)) return newState
+          return computeAndMerge(Object.assign(state, newState))
+        }, replace)
+        return
+      }
+
+      set((state) => {
+        const newArg = arg as T | ((state: T) => T)
+        const newState: T = typeof newArg === "function" ? newArg(state) : newArg
+        if (!shouldRecomputeFn(state, newState)) return newState
+        return computeAndMerge(newState)
+      }, replace)
     }
 
-    const _api = api as Mutate<StoreApi<T>, [["chrisvander/zustand-computed", A]]>
-    _api.setState = setWithComputed
-    const st = f(setWithComputed, get, _api) as T & A
+    _api.setState = setState
+    const st = f(setState, get, _api)
     return Object.assign({}, st, compute(st))
   }
 }