From bbc2a6bcfc3574d582c51d5957b33f57076b2bdd Mon Sep 17 00:00:00 2001
From: Dor Shtaif <dorshtaif@gmail.com>
Date: Sun, 29 Dec 2024 13:24:44 +0200
Subject: [PATCH] make iterators of the `useAsyncIterState` hook's iterable
 individually closable

---
 spec/tests/useAsyncIterState.spec.tsx    | 41 +++++++++++++++++++++++-
 spec/utils/checkPromiseState.ts          | 19 +++++++++++
 src/useAsyncIterState/IterableChannel.ts | 22 ++++++++++---
 src/useAsyncIterState/index.ts           |  6 ++--
 4 files changed, 79 insertions(+), 9 deletions(-)
 create mode 100644 spec/utils/checkPromiseState.ts

diff --git a/spec/tests/useAsyncIterState.spec.tsx b/spec/tests/useAsyncIterState.spec.tsx
index 819fc1c..c9141b9 100644
--- a/spec/tests/useAsyncIterState.spec.tsx
+++ b/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;
diff --git a/spec/utils/checkPromiseState.ts b/spec/utils/checkPromiseState.ts
new file mode 100644
index 0000000..682537a
--- /dev/null
+++ b/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 };
diff --git a/src/useAsyncIterState/IterableChannel.ts b/src/useAsyncIterState/IterableChannel.ts
index 996867b..6362c4d 100644
--- a/src/useAsyncIterState/IterableChannel.ts
+++ b/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>;
 }
diff --git a/src/useAsyncIterState/index.ts b/src/useAsyncIterState/index.ts
index df90142..c489e6c 100644
--- a/src/useAsyncIterState/index.ts
+++ b/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;