Make useAsyncIterState's setter accept a function as to calculate next state

spec/tests/useAsyncIterState.spec.tsx

@@ -1,4 +1,4 @@
-import { it, describe, expect, afterEach } from 'vitest';
+import { it, describe, expect, afterEach, vi } from 'vitest';
 import { gray } from 'colorette';
 import { range } from 'lodash-es';
 import { renderHook, cleanup as cleanupMountedReactTrees, act } from '@testing-library/react';
@@ -14,43 +14,6 @@ afterEach(() => {
 });
 
 describe('`useAsyncIterState` hook', () => {
-  it(gray('Updating states iteratively with the returned setter works correctly'), async () => {
-    const [values, setValue] = renderHook(() => useAsyncIterState<number>()).result.current;
-
-    const rounds = 3;
-
-    const yieldsPromise = pipe(values, asyncIterTake(rounds), asyncIterToArray);
-    const currentValues = [values.value.current];
-
-    for (let i = 0; i < rounds; ++i) {
-      await act(() => {
-        setValue(i);
-        currentValues.push(values.value.current);
-      });
-    }
-
-    expect(currentValues).toStrictEqual([undefined, 0, 1, 2]);
-    expect(await yieldsPromise).toStrictEqual([0, 1, 2]);
-  });
-
-  it(
-    gray('Updating states as rapidly as possible with the returned setter works correctly'),
-    async () => {
-      const [values, setValue] = renderHook(() => useAsyncIterState<number>()).result.current;
-
-      const yieldPromise = pipe(values, asyncIterTakeFirst());
-      const currentValues = [values.value.current];
-
-      for (let i = 0; i < 3; ++i) {
-        setValue(i);
-        currentValues.push(values.value.current);
-      }
-
-      expect(currentValues).toStrictEqual([undefined, 0, 1, 2]);
-      expect(await yieldPromise).toStrictEqual(2);
-    }
-  );
-
   it(gray('The returned iterable can be async-iterated upon successfully'), async () => {
     const [values, setValue] = renderHook(() => useAsyncIterState<string>()).result.current;
 
@@ -108,6 +71,93 @@ describe('`useAsyncIterState` hook', () => {
     }
   );
 
+  it(gray('Updating states iteratively with the returned setter works correctly'), async () => {
+    const [values, setValue] = renderHook(() => useAsyncIterState<number>()).result.current;
+
+    const rounds = 3;
+
+    const yieldsPromise = pipe(values, asyncIterTake(rounds), asyncIterToArray);
+    const currentValues = [values.value.current];
+
+    for (let i = 0; i < rounds; ++i) {
+      await act(() => {
+        setValue(i);
+        currentValues.push(values.value.current);
+      });
+    }
+
+    expect(currentValues).toStrictEqual([undefined, 0, 1, 2]);
+    expect(await yieldsPromise).toStrictEqual([0, 1, 2]);
+  });
+
+  it(
+    gray('Updating states as rapidly as possible with the returned setter works correctly'),
+    async () => {
+      const [values, setValue] = renderHook(() => useAsyncIterState<number>()).result.current;
+
+      const yieldPromise = pipe(values, asyncIterTakeFirst());
+      const currentValues = [values.value.current];
+
+      for (let i = 0; i < 3; ++i) {
+        setValue(i);
+        currentValues.push(values.value.current);
+      }
+
+      expect(currentValues).toStrictEqual([undefined, 0, 1, 2]);
+      expect(await yieldPromise).toStrictEqual(2);
+    }
+  );
+
+  it(
+    gray(
+      'Updating states iteratively with the returned setter *in the functional form* works correctly'
+    ),
+    async () => {
+      const renderFn = vi.fn<(prevState: number | undefined) => number>();
+      const [values, setValue] = renderHook(() => useAsyncIterState<number>()).result.current;
+
+      const rounds = 3;
+
+      const yieldsPromise = pipe(values, asyncIterTake(rounds), asyncIterToArray);
+      const currentValues = [values.value.current];
+
+      for (let i = 0; i < rounds; ++i) {
+        await act(() => {
+          setValue(renderFn.mockImplementation(_prev => i));
+          currentValues.push(values.value.current);
+        });
+      }
+
+      expect(renderFn.mock.calls).toStrictEqual([[undefined], [0], [1]]);
+      expect(currentValues).toStrictEqual([undefined, 0, 1, 2]);
+      expect(await yieldsPromise).toStrictEqual([0, 1, 2]);
+    }
+  );
+
+  it(
+    gray(
+      'Updating states as rapidly as possible with the returned setter *in the functional form* works correctly'
+    ),
+    async () => {
+      const renderFn = vi.fn<(prevState: number | undefined) => number>();
+
+      const [values, setValue] = renderHook(() => useAsyncIterState<number>()).result.current;
+
+      const yieldPromise = pipe(values, asyncIterTakeFirst());
+
+      const currentValues = [values.value.current];
+
+      for (let i = 0; i < 3; ++i) {
+        setValue(renderFn.mockImplementation(_prev => i));
+        currentValues.push(values.value.current);
+      }
+
+      expect(renderFn.mock.calls).toStrictEqual([[undefined], [0], [1]]);
+      expect(currentValues).toStrictEqual([undefined, 0, 1, 2]);
+      expect(await yieldPromise).toStrictEqual(2);
+    }
+  );
+
   it(
     gray(
       'When hook is unmounted, all outstanding yieldings of the returned iterable resolve to "done"'
@@ -198,7 +248,7 @@ describe('`useAsyncIterState` hook', () => {
     const [values] = renderHook(() => useAsyncIterState<number>()).result.current;
 
     expect(() => {
-      (values.value as any).current = "can't do this...";
+      (values.value as any).current = `CAN'T DO THIS...`;
     }).toThrow(TypeError);
   });
 });

src/useAsyncIterState/IterableChannel.ts

@@ -8,8 +8,16 @@ class IterableChannel<T> {
   #nextIteration = promiseWithResolvers<IteratorResult<T, void>>();
   #currentValue: T | undefined;
 
-  put(value: T): void {
+  put(update: T | ((prevState: T | undefined) => T)): void {
     if (!this.#isClosed) {
+      const value =
+        typeof update !== 'function'
+          ? update
+          : (() => {
+              const updateFnTypePatched = update as (prevState: T | undefined) => T;
+              return updateFnTypePatched(this.#currentValue);
+            })();
+
       (async () => {
         this.#currentValue = value;
         await undefined; // Deferring to the next microtick so that an attempt to pull the a value before making multiple rapid synchronous calls to `put()` will make that pull ultimately yield only the last value that was put - instead of the first one as were if this otherwise wasn't deferred.

src/useAsyncIterState/index.ts

@@ -8,9 +8,9 @@ export { useAsyncIterState, type AsyncIterStateResult, type AsyncIterableSubject
  * Basically like {@link https://react.dev/reference/react/useState `React.useState`}, only that the value
  * is provided back __wrapped as an async iterable__.
  *
- * This hook allows a component to declare and manage a piece of state while easily letting it control
- * what area(s) specifically within the UI should be bound to it (should re-render in reaction to changes
- * in it) - combined for example with one or more {@link Iterate `<Iterate>`}s.
+ * This hook allows a component to declare and manage a piece of state while easily letting you control
+ * what specifically area(s) within the UI should be bound to it (should re-render in reaction to changes
+ * in it) - for example, if combined with one or more {@link Iterate `<Iterate>`}s.
  *
  * @example
  * ```tsx
@@ -36,14 +36,24 @@ export { useAsyncIterState, type AsyncIterStateResult, type AsyncIterableSubject
  *
  * ---
  *
- * This is unlike vanila `React.useState` which simply re-renders the entire component. Instead,
- * `useAsyncIterState` helps confine UI updates as well as facilitate layers of sub-components that pass
- * actual async iterables across one another as props, skipping typical cascading re-renderings down to
- * __only the inner-most leafs__ of the UI tree.
  *
- * The returned async iterable contains a `.current.value` property which shows the current up to date
- * state value at all times. Use this any case you just need to read the immediate current state rather
- * than directly rendering it, since for rendering you may simply async-iterate it.
+ *
+ * The returned async iterable can be passed over to any level down the component tree and rendered
+ * using `<Iterate>`, `useAsyncIter`, and so on. It also contains a `.current.value` property which shows
+ * the current up to date state value at all times. Use this any case you just need to read the immediate
+ * current state rather than directly rendering it, since for rendering you may simply async-iterate it.
+ *
+ * Returned also alongside the async iterable is a function for updating the state. Calling it with a new
+ * value will cause the paired iterable to yield the updated state value as well as immediately set the
+ * iterable's `.current.value` property to that new state. Just like
+ * [`React.useState`'s setter](https://react.dev/reference/react/useState#setstate), you can pass it
+ * the next state directly, or a function that calculates it from the previous state.
+ *
+ * Unlike vanila `React.useState`, which simply re-renders the entire component - `useAsyncIterState`
+ * helps confine UI updates by handing you an iterable which choose how and where in the component tree
+ * to render it. This work method can facilitate layers of sub-components that pass actual async iterables
+ * across one another as props, skipping typical cascading re-renderings down to __only the inner-most
+ * leafs__ of the UI tree.
  *
  * @example
  * ```tsx
@@ -107,7 +117,8 @@ function useAsyncIterState<TVal>(): AsyncIterStateResult<TVal> {
 }
 
 /**
- * A pair of stateful async iterable and a function which modifies the state and yields the updated value.
+ * A pair of stateful async iterable and a function which updates the state and making the paired
+ * async iterable yield the new value.
  * Returned from the {@link useAsyncIterState `useAsyncIterState`} hook.
  *
  * @see {@link useAsyncIterState `useAsyncIterState`}
@@ -125,8 +136,8 @@ type AsyncIterStateResult<TVal> = [
   values: AsyncIterableSubject<TVal>,
 
   /**
-   * A function which modifies the state, causing the paired async iterable to yield the updated state
+   * A function which updates the state, causing the paired async iterable to yield the updated state
    * value and immediately sets its `.current.value` property to the latest state.
    */
-  setValue: (newValue: TVal) => void,
+  setValue: (update: TVal | ((prevState: TVal | undefined) => TVal)) => void,
 ];