From 1db7dd45fb352db8b63196e0b456f2003a85b7a2 Mon Sep 17 00:00:00 2001
From: Dor Shtaif <dorshtaif@gmail.com>
Date: Thu, 9 Jan 2025 13:26:14 +0200
Subject: [PATCH] feat(useAsyncIter): allow initial value to be a function,
 called once on mount

---
 spec/tests/useAsyncIter.spec.ts | 28 +++++++++++++++++++++++++++-
 src/common/callOrReturn.ts      |  5 +++++
 src/useAsyncIter/index.ts       | 23 +++++++++++++----------
 3 files changed, 45 insertions(+), 11 deletions(-)
 create mode 100644 src/common/callOrReturn.ts

diff --git a/spec/tests/useAsyncIter.spec.ts b/spec/tests/useAsyncIter.spec.ts
index e9c5e7e..2127d96 100644
--- a/spec/tests/useAsyncIter.spec.ts
+++ b/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>();
 
diff --git a/src/common/callOrReturn.ts b/src/common/callOrReturn.ts
new file mode 100644
index 0000000..0f31993
--- /dev/null
+++ b/src/common/callOrReturn.ts
@@ -0,0 +1,5 @@
+export { callOrReturn };
+
+function callOrReturn<T>(value: T | (() => T)): T {
+  return typeof value !== 'function' ? value : (value as () => T)();
+}
diff --git a/src/useAsyncIter/index.ts b/src/useAsyncIter/index.ts
index 91042bd..27d9238 100644
--- a/src/useAsyncIter/index.ts
+++ b/src/useAsyncIter/index.ts
@@ -1,20 +1,20 @@
-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 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 +62,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 +100,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: TInitVal | (() => TInitVal)
+  ): IterationResult<TVal, TInitVal>;
 } = <
   TVal extends
     | undefined
@@ -112,19 +115,19 @@ const useAsyncIter: {
           ExtractAsyncIterValue<TVal>
         >;
       },
-  TInitVal = undefined,
+  TInitVal,
 >(
   input: TVal,
-  initialVal: TInitVal
+  initialVal: TInitVal | (() => 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);