feat(useAsyncIter): allow initial value to be a function, called once on mount (#48)

* allow initial value to be a function, called once on mount

* introduce `MaybeFunction` type to clean up places accepting a value-or-function types of inputs

Author: shtaif

spec/tests/useAsyncIter.spec.ts

@@ -1,4 +1,4 @@
-import { it, describe, expect, afterEach } from 'vitest';
+import { it, describe, expect, afterEach, vi } from 'vitest';
 import { gray } from 'colorette';
 import { cleanup as cleanupMountedReactTrees, act, renderHook } from '@testing-library/react';
 import { useAsyncIter, iterateFormatted } from '../../src/index.js';
@@ -404,6 +404,32 @@ describe('`useAsyncIter` hook', () => {
     }
   );
 
+  it(
+    gray(
+      'When given an initial value as a function, calls it once on mount and uses its result as the initial value correctly'
+    ),
+    async () => {
+      const channel = new IteratorChannelTestHelper<string>();
+      const initValFn = vi.fn(() => '_');
+
+      const renderedHook = await act(() => renderHook(() => useAsyncIter(channel, initValFn)));
+      const results = [renderedHook.result.current];
+
+      await act(() => renderedHook.rerender());
+      results.push(renderedHook.result.current);
+
+      await act(() => channel.put('a'));
+      results.push(renderedHook.result.current);
+
+      expect(initValFn).toHaveBeenCalledOnce();
+      expect(results).toStrictEqual([
+        { value: '_', pendingFirst: true, done: false, error: undefined },
+        { value: '_', pendingFirst: true, done: false, error: undefined },
+        { value: 'a', pendingFirst: false, done: false, error: undefined },
+      ]);
+    }
+  );
+
   it(gray('When unmounted will close the last active iterator it held'), async () => {
     const channel = new IteratorChannelTestHelper<string>();
 

src/common/MaybeFunction.ts

@@ -0,0 +1,3 @@
+export { MaybeFunction };
+
+type MaybeFunction<T, TPossibleArgs extends unknown[] = []> = T | ((...args: TPossibleArgs) => T);

src/common/callOrReturn.ts

@@ -0,0 +1,7 @@
+import { type MaybeFunction } from './MaybeFunction.js';
+
+export { callOrReturn };
+
+function callOrReturn<T>(value: MaybeFunction<T>): T {
+  return typeof value !== 'function' ? value : (value as () => T)();
+}

src/useAsyncIter/index.ts

@@ -1,20 +1,21 @@
-import { useRef, useMemo, useEffect } from 'react';
+import { useMemo, useEffect } from 'react';
 import { useLatest } from '../common/hooks/useLatest.js';
 import { isAsyncIter } from '../common/isAsyncIter.js';
 import { useSimpleRerender } from '../common/hooks/useSimpleRerender.js';
+import { useRefWithInitialValue } from '../common/hooks/useRefWithInitialValue.js';
+import { type MaybeFunction } from '../common/MaybeFunction.js';
 import { type ExtractAsyncIterValue } from '../common/ExtractAsyncIterValue.js';
 import {
   reactAsyncIterSpecialInfoSymbol,
   type ReactAsyncIterSpecialInfo,
 } from '../common/ReactAsyncIterable.js';
 import { iterateAsyncIterWithCallbacks } from '../common/iterateAsyncIterWithCallbacks.js';
+import { callOrReturn } from '../common/callOrReturn.js';
 import { type Iterate } from '../Iterate/index.js'; // eslint-disable-line @typescript-eslint/no-unused-vars
 import { type iterateFormatted } from '../iterateFormatted/index.js'; // eslint-disable-line @typescript-eslint/no-unused-vars
 
 export { useAsyncIter, type IterationResult };
 
-// TODO: The initial values should be able to be given as functions, having them called once on mount
-
 /**
  * `useAsyncIter` hooks up a single async iterable value to your component and its lifecycle.
  *
@@ -62,7 +63,7 @@ export { useAsyncIter, type IterationResult };
  * @template TInitVal The type of the initial value, defaults to `undefined`.
  *
  * @param input Any async iterable or plain value.
- * @param initialVal Any initial value for the hook to return prior to resolving the ___first emission___ of the ___first given___ async iterable, defaults to `undefined`.
+ * @param initialVal Any optional starting value for the hook to return prior to the ___first yield___ of the ___first given___ async iterable, defaults to `undefined`. You can pass an actual value, or a function that returns a value (which the hook will call once during mounting).
  *
  * @returns An object with properties reflecting the current state of the iterated async iterable or plain value provided via `input` (see {@link IterationResult `IterationResult`}).
  *
@@ -100,7 +101,10 @@ export { useAsyncIter, type IterationResult };
  */
 const useAsyncIter: {
   <TVal>(input: TVal, initialVal?: undefined): IterationResult<TVal>;
-  <TVal, TInitVal>(input: TVal, initialVal: TInitVal): IterationResult<TVal, TInitVal>;
+  <TVal, TInitVal>(
+    input: TVal,
+    initialVal: MaybeFunction<TInitVal>
+  ): IterationResult<TVal, TInitVal>;
 } = <
   TVal extends
     | undefined
@@ -112,19 +116,19 @@ const useAsyncIter: {
           ExtractAsyncIterValue<TVal>
         >;
       },
-  TInitVal = undefined,
+  TInitVal,
 >(
   input: TVal,
-  initialVal: TInitVal
+  initialVal: MaybeFunction<TInitVal>
 ): IterationResult<TVal, TInitVal> => {
   const rerender = useSimpleRerender();
 
-  const stateRef = useRef<IterationResult<TVal, TInitVal>>({
-    value: initialVal as any,
+  const stateRef = useRefWithInitialValue<IterationResult<TVal, TInitVal>>(() => ({
+    value: callOrReturn(initialVal) as any,
     pendingFirst: true,
     done: false,
     error: undefined,
-  });
+  }));
 
   const latestInputRef = useLatest(input);