fix: withImmer doesn't work with resettable atom

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>,
 ) {

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');
+});