From 343f742ab8d0c1a6cf1ff5021315d22ef95964cc Mon Sep 17 00:00:00 2001
From: unadlib <unadlib@gmail.com>
Date: Sun, 19 May 2024 05:12:04 +0800
Subject: [PATCH] feat(mutative): add options

---
 README.md          |  9 ++++++-
 src/index.ts       | 48 ++++++++++++++++++++++++------------
 test/index.test.ts | 61 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+), 16 deletions(-)

diff --git a/README.md b/README.md
index 2e87719..eb8980d 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,12 @@
 # zustand-mutative
-A Mutative middleware for Zustand enhances the efficiency of immutable state updates.
+
+![Node CI](https://github.com/mutativejs/zustand-mutative/workflows/Node%20CI/badge.svg)
+[![npm](https://img.shields.io/npm/v/zustand-mutative.svg)](https://www.npmjs.com/package/zustand-mutative)
+![license](https://img.shields.io/npm/l/zustand-mutative)
+
+A [Mutative](https://github.com/unadlib/mutative) middleware for Zustand enhances the efficiency of immutable state updates.
+
+`zustand-mutative` is 2-6x faster than zustand with spread operation, more than 10x faster than `zustand/middleware/immer`.[Read more about the performance comparison in Mutative](https://mutative.js.org/docs/getting-started/performance).
 
 ## Installation
 
diff --git a/src/index.ts b/src/index.ts
index 30def69..320893e 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,13 +1,24 @@
-import { create } from 'mutative';
+import {
+  create,
+  type Options as MutativeOptions,
+  PatchesOptions,
+} from 'mutative';
 import type { Draft } from 'mutative';
 import type { StateCreator, StoreMutatorIdentifier } from 'zustand';
 
+type Options<O extends PatchesOptions, F extends boolean> = Pick<
+  MutativeOptions<O, F>,
+  Exclude<keyof MutativeOptions<O, F>, 'enablePatches'>
+>;
+
 type Mutative = <
   T,
   Mps extends [StoreMutatorIdentifier, unknown][] = [],
-  Mcs extends [StoreMutatorIdentifier, unknown][] = []
+  Mcs extends [StoreMutatorIdentifier, unknown][] = [],
+  F extends boolean = false
 >(
-  initializer: StateCreator<T, [...Mps, ['zustand/mutative', never]], Mcs>
+  initializer: StateCreator<T, [...Mps, ['zustand/mutative', never]], Mcs>,
+  options?: Options<false, F>
 ) => StateCreator<T, Mps, [['zustand/mutative', never], ...Mcs]>;
 
 type Write<T, U> = Omit<T, keyof U> & U;
@@ -49,22 +60,29 @@ type StoreMutative<S> = S extends {
     : never
   : never;
 
-type MutativeImpl = <T>(
-  storeInitializer: StateCreator<T, [], []>
+type MutativeImpl = <T, F extends boolean = false>(
+  storeInitializer: StateCreator<T, [], []>,
+  options?: Options<false, F>
 ) => StateCreator<T, [], []>;
 
-const mutativeImpl: MutativeImpl = (initializer) => (set, get, store) => {
-  type T = ReturnType<typeof initializer>;
+const mutativeImpl: MutativeImpl =
+  (initializer, options) => (set, get, store) => {
+    type T = ReturnType<typeof initializer>;
 
-  store.setState = (updater, replace, ...a) => {
-    const nextState = (
-      typeof updater === 'function' ? create(updater as any) : updater
-    ) as ((s: T) => T) | T | Partial<T>;
+    store.setState = (updater, replace, ...a) => {
+      const nextState = (
+        typeof updater === 'function'
+          ? create(
+              updater as any,
+              options ? { ...options, enablePatches: false } : options
+            )
+          : updater
+      ) as ((s: T) => T) | T | Partial<T>;
 
-    return set(nextState as any, replace, ...a);
-  };
+      return set(nextState as any, replace, ...a);
+    };
 
-  return initializer(store.setState, get, store);
-};
+    return initializer(store.setState, get, store);
+  };
 
 export const mutative = mutativeImpl as unknown as Mutative;
diff --git a/test/index.test.ts b/test/index.test.ts
index 147680b..8b987e3 100644
--- a/test/index.test.ts
+++ b/test/index.test.ts
@@ -36,3 +36,64 @@ test('Updating simple states', () => {
   act(() => decrement(1));
   expect(result.current.count).toBe(0);
 });
+
+test('Updating simple states with options', () => {
+  class Counter {
+    value: number;
+    constructor(value: number) {
+      this.value = value;
+    }
+  }
+
+  type State = {
+    count: {
+      value: number;
+      instance: Counter;
+    };
+  };
+
+  type Actions = {
+    increment: (qty: number) => void;
+    decrement: (qty: number) => void;
+  };
+
+  const useCountStore = create<State & Actions>()(
+    mutative(
+      (set) => ({
+        count: {
+          value: 0,
+          instance: new Counter(0),
+        },
+        increment: (qty: number) =>
+          set((state) => {
+            state.count.value += qty;
+            state.count.instance.value += qty;
+          }),
+        decrement: (qty: number) =>
+          set((state) => {
+            state.count.value -= qty;
+            state.count.instance.value -= qty;
+          }),
+      }),
+      {
+        mark: () => 'immutable',
+      }
+    )
+  );
+
+  const { result } = renderHook(() => useCountStore());
+  const { increment, decrement } = result.current;
+
+  let currentCount = result.current.count;
+  act(() => increment(1));
+  expect(currentCount).not.toBe(result.current.count);
+  expect(currentCount.instance).not.toBe(result.current.count.instance);
+  expect(result.current.count.value).toBe(1);
+  expect(result.current.count.instance.value).toBe(1);
+
+  act(() => decrement(1));
+  expect(currentCount).not.toBe(result.current.count);
+  expect(currentCount.instance).not.toBe(result.current.count.instance);
+  expect(result.current.count.value).toBe(0);
+  expect(result.current.count.instance.value).toBe(0);
+});