From 99cd63073d342798fb95455d523d174984021b7a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ti=E1=BA=BFn=20Nguy=E1=BB=85n=20Kh=E1=BA=AFc?=
 <tien.nguyenkhac@icloud.com>
Date: Sun, 23 Jun 2024 23:29:45 +1200
Subject: [PATCH] fix: `withImmer` doesn't work with resettable atom

---
 src/withImmer.ts         | 18 +++++++++++++----
 tests/withImmer.spec.tsx | 42 +++++++++++++++++++++++++++++++++++++++-
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/src/withImmer.ts b/src/withImmer.ts
index e0fd53f..14bcaff 100644
--- a/src/withImmer.ts
+++ b/src/withImmer.ts
@@ -7,14 +7,24 @@ const cache1 = new WeakMap();
 const memo1 = <T>(create: () => T, dep1: object): T =>
   (cache1.has(dep1) ? cache1 : cache1.set(dep1, create())).get(dep1);
 
+export function withImmer<Value, Args extends unknown[], Result>(
+  anAtom: WritableAtom<Value, Args, Result>,
+): WritableAtom<
+  Value,
+  Args extends [Value | infer OtherValue]
+    ? [
+        | Value
+        | ((draft: Draft<Value>) => void)
+        | Exclude<OtherValue, (...args: never[]) => unknown>,
+      ]
+    : unknown[],
+  Result
+>;
+
 export function withImmer<Value>(
   anAtom: PrimitiveAtom<Value>,
 ): WritableAtom<Value, [Value | ((draft: Draft<Value>) => void)], void>;
 
-export function withImmer<Value, Result>(
-  anAtom: WritableAtom<Value, [Value], Result>,
-): WritableAtom<Value, [Value | ((draft: Draft<Value>) => void)], Result>;
-
 export function withImmer<Value, Result>(
   anAtom: WritableAtom<Value, [Value], Result>,
 ) {
diff --git a/tests/withImmer.spec.tsx b/tests/withImmer.spec.tsx
index c316a12..8f57e2c 100644
--- a/tests/withImmer.spec.tsx
+++ b/tests/withImmer.spec.tsx
@@ -2,6 +2,7 @@ import { afterEach, test } from 'vitest';
 import { StrictMode } from 'react';
 import { cleanup, fireEvent, render } from '@testing-library/react';
 import { useAtom } from 'jotai/react';
+import { atomWithDefault, useResetAtom } from 'jotai/utils';
 import { atom } from 'jotai/vanilla';
 import { withImmer } from 'jotai-immer';
 
@@ -40,7 +41,7 @@ test('withImmer derived atom with useAtom', async () => {
   await findByText('count: 0');
 });
 
-test('withImmer derived atom with WritableAtom<Value, Value> signature', async () => {
+test('withImmer derived atom with WritableAtom<Value, [Value]> signature', async () => {
   const regularCountAtom = atom(0);
 
   const Parent = () => {
@@ -68,3 +69,42 @@ test('withImmer derived atom with WritableAtom<Value, Value> signature', async (
   fireEvent.click(getByText('Decrease'));
   await findByText('count: 0');
 });
+
+test('withImmer derived atom with WritableAtom<Value, [Value | OtherValue]> signature', async () => {
+  const regularCountAtom = atomWithDefault(() => 0);
+  const immerCountAtom = withImmer(regularCountAtom);
+
+  const Parent = () => {
+    const [count, setCount] = useAtom(immerCountAtom);
+    const resetCount = useResetAtom(immerCountAtom);
+    return (
+      <>
+        <div>count: {count}</div>
+        <button onClick={() => setCount(count + 1)}>Increase</button>
+        <button onClick={() => setCount(count - 1)}>Decrease</button>
+        <button onClick={() => resetCount()}>Reset</button>
+      </>
+    );
+  };
+
+  const { findByText, getByText } = render(
+    <StrictMode>
+      <Parent />
+    </StrictMode>,
+  );
+
+  await findByText('count: 0');
+
+  fireEvent.click(getByText('Increase'));
+  await findByText('count: 1');
+
+  fireEvent.click(getByText('Decrease'));
+  await findByText('count: 0');
+
+  fireEvent.click(getByText('Increase'));
+  fireEvent.click(getByText('Increase'));
+  await findByText('count: 2');
+
+  fireEvent.click(getByText('Reset'));
+  await findByText('count: 0');
+});