Make iterators of the useAsyncIterState's iterable individually closable

spec/tests/useAsyncIterState.spec.tsx

@@ -5,6 +5,7 @@ import { renderHook, cleanup as cleanupMountedReactTrees, act } from '@testing-l
 import { useAsyncIterState } from '../../src/index.js';
 import { asyncIterToArray } from '../utils/asyncIterToArray.js';
 import { asyncIterTake } from '../utils/asyncIterTake.js';
+import { checkPromiseState } from '../utils/checkPromiseState.js';
 import { pipe } from '../utils/pipe.js';
 
 afterEach(() => {
@@ -26,6 +27,44 @@ describe('`useAsyncIterState` hook', () => {
     expect(await collectPromise).toStrictEqual(['a', 'b', 'c']);
   });
 
+  it(
+    gray(
+      'Each iterator of the hook-returned iterable, upon getting manually closed, will immediately resolve all outstanding yieldings specifically pulled from it to "done'
+    ),
+    async () => {
+      const [values] = renderHook(() => useAsyncIterState<string>()).result.current;
+
+      const iterator1 = values[Symbol.asyncIterator]();
+      const iterator2 = values[Symbol.asyncIterator]();
+      const yieldPromise1 = iterator1.next();
+      const yieldPromise2 = iterator2.next();
+
+      await iterator1.return!();
+
+      {
+        const promiseStates = await Promise.all(
+          [yieldPromise1, yieldPromise2].map(checkPromiseState)
+        );
+        expect(promiseStates).toStrictEqual([
+          { state: 'FULFILLED', value: { done: true, value: undefined } },
+          { state: 'PENDING', value: undefined },
+        ]);
+      }
+
+      await iterator2.return!();
+
+      {
+        const promiseStates = await Promise.all(
+          [yieldPromise1, yieldPromise2].map(checkPromiseState)
+        );
+        expect(promiseStates).toStrictEqual([
+          { state: 'FULFILLED', value: { done: true, value: undefined } },
+          { state: 'FULFILLED', value: { done: true, value: undefined } },
+        ]);
+      }
+    }
+  );
+
   it(
     gray(
       'When hook is unmounted, all outstanding yieldings of the returned iterable resolve to "done"'
@@ -79,7 +118,7 @@ describe('`useAsyncIterState` hook', () => {
 
   it(
     gray(
-      "The returned iterable's values are each shared between all its parallel consumers so that each receives all the values from the start of consumption and onwards"
+      "The returned iterable's values are each shared between all its parallel consumers so that each receives all the values that will yield after the start of its consumption"
     ),
     async () => {
       const [values, setValue] = renderHook(() => useAsyncIterState<string>()).result.current;

spec/utils/checkPromiseState.ts

@@ -0,0 +1,19 @@
+export { checkPromiseState, type PromiseCurrentState };
+
+async function checkPromiseState<T>(p: Promise<T>): Promise<PromiseCurrentState<T>> {
+  let result: PromiseCurrentState<T> = { state: 'PENDING', value: undefined };
+
+  p.then(
+    val => (result = { state: 'FULFILLED', value: val }),
+    reason => (result = { state: 'REJECTED', value: reason })
+  );
+
+  await undefined;
+
+  return result;
+}
+
+type PromiseCurrentState<T> =
+  | { state: 'PENDING'; value: void }
+  | { state: 'FULFILLED'; value: T }
+  | { state: 'REJECTED'; value: unknown };

src/useAsyncIterState/IterableChannel.ts

@@ -5,11 +5,6 @@ export { IterableChannel };
 class IterableChannel<TVal> {
   #isClosed = false;
   #nextIteration = promiseWithResolvers<IteratorResult<TVal, void>>();
-  iterable = {
-    [Symbol.asyncIterator]: () => ({
-      next: () => this.#nextIteration.promise,
-    }),
-  };
 
   put(value: TVal): void {
     if (!this.#isClosed) {
@@ -22,4 +17,21 @@ class IterableChannel<TVal> {
     this.#isClosed = true;
     this.#nextIteration.resolve({ done: true, value: undefined });
   }
+
+  iterable = {
+    [Symbol.asyncIterator]: () => {
+      const whenIteratorClosed = promiseWithResolvers<IteratorReturnResult<undefined>>();
+
+      return {
+        next: () => {
+          return Promise.race([this.#nextIteration.promise, whenIteratorClosed.promise]);
+        },
+
+        return: async () => {
+          whenIteratorClosed.resolve({ done: true, value: undefined });
+          return { done: true as const, value: undefined };
+        },
+      };
+    },
+  } satisfies AsyncIterable<TVal, void, void>;
 }

src/useAsyncIterState/index.ts

@@ -54,13 +54,13 @@ function useAsyncIterState<TVal>(): AsyncIterStateResult<TVal> {
     result: AsyncIterStateResult<TVal>;
   }>();
 
-  if (!ref.current) {
+  ref.current ??= (() => {
     const channel = new IterableChannel<TVal>();
-    ref.current = {
+    return {
       channel,
       result: [channel.iterable, newVal => channel.put(newVal)],
     };
-  }
+  })();
 
   const { channel, result } = ref.current;