From 15d7c3ee55f02d47af0cf53b2dbeebab94abb608 Mon Sep 17 00:00:00 2001
From: Alisue <lambdalisue@gmail.com>
Date: Thu, 8 Aug 2024 10:47:55 +0900
Subject: [PATCH] feat: support chaining malformed iterables

---
 async/chain.ts      | 28 ++++++++++++++++++++++++----
 async/chain_test.ts | 13 +++++++++++++
 chain.ts            | 21 +++++++++++++++++++--
 chain_test.ts       | 21 ++++++++++++++++-----
 4 files changed, 72 insertions(+), 11 deletions(-)

diff --git a/async/chain.ts b/async/chain.ts
index 362e477..cf9e8ad 100644
--- a/async/chain.ts
+++ b/async/chain.ts
@@ -12,13 +12,33 @@
  * const iter = chain([1, 2], [3, 4]);
  * console.log(await toArray(iter)); // [1, 2, 3, 4]
  * ```
+ *
+ * It supports chaining malformed iterables.
+ *
+ * @example
+ * ```ts
+ * import { toArray } from "@core/iterutil/async/to-array";
+ * import { chain } from "@core/iterutil/async/chain";
+ *
+ * const iter = chain([1, 2], ["a", "b"], [true]);
+ * console.log(await toArray(iter)); // [1, 2, "a", "b", true]
+ * ```
  */
-export async function* chain<T>(
-  ...iterables: (Iterable<T> | AsyncIterable<T>)[]
-): AsyncIterable<T> {
+export async function* chain<
+  T extends (Iterable<unknown> | AsyncIterable<unknown>)[],
+>(
+  ...iterables: T
+): AsyncIterable<Chain<T>> {
   for await (const iterable of iterables) {
     for await (const value of iterable) {
-      yield value;
+      yield value as Chain<T>;
     }
   }
 }
+
+export type Chain<T> = T extends readonly [] ? never
+  : T extends readonly [Iterable<infer U>] ? U
+  : T extends readonly [AsyncIterable<infer U>] ? U
+  : T extends readonly [Iterable<infer U>, ...infer R] ? U | Chain<R>
+  : T extends readonly [AsyncIterable<infer U>, ...infer R] ? U | Chain<R>
+  : never;
diff --git a/async/chain_test.ts b/async/chain_test.ts
index 898baa7..516ec54 100644
--- a/async/chain_test.ts
+++ b/async/chain_test.ts
@@ -33,4 +33,17 @@ Deno.test("chain", async (t) => {
     assertEquals(await toArray(result), expected);
     assertType<IsExact<typeof result, AsyncIterable<number>>>(true);
   });
+
+  await t.step("with malform iterable", async () => {
+    const result = chain(
+      toAsyncIterable([1, 2]),
+      ["a", "b"],
+      toAsyncIterable([true]),
+    );
+    const expected = [1, 2, "a", "b", true];
+    assertEquals(await toArray(result), expected);
+    assertType<
+      IsExact<typeof result, AsyncIterable<number | string | boolean>>
+    >(true);
+  });
 });
diff --git a/chain.ts b/chain.ts
index af73bac..2b601a1 100644
--- a/chain.ts
+++ b/chain.ts
@@ -11,9 +11,26 @@
  * const iter = chain([1, 2], [3, 4]);
  * console.log([...iter]); // [1, 2, 3, 4]
  * ```
+ *
+ * It supports chaining malformed iterables.
+ *
+ * @example
+ * ```ts
+ * import { chain } from "@core/iterutil/chain";
+ *
+ * const iter = chain([1, 2], ["a", "b"], [true]);
+ * console.log([...iter]); // [1, 2, "a", "b", true]
+ * ```
  */
-export function* chain<T>(...iterables: Iterable<T>[]): Iterable<T> {
+export function* chain<T extends Iterable<unknown>[]>(
+  ...iterables: T
+): Iterable<Chain<T>> {
   for (const iterable of iterables) {
-    yield* iterable;
+    yield* iterable as Iterable<Chain<T>>;
   }
 }
+
+export type Chain<T> = T extends readonly [] ? never
+  : T extends readonly [Iterable<infer U>] ? U
+  : T extends readonly [Iterable<infer U>, ...infer R] ? U | Chain<R>
+  : never;
diff --git a/chain_test.ts b/chain_test.ts
index 03df554..ed74e5b 100644
--- a/chain_test.ts
+++ b/chain_test.ts
@@ -2,9 +2,20 @@ import { assertEquals } from "@std/assert";
 import { assertType, type IsExact } from "@std/testing/types";
 import { chain } from "./chain.ts";
 
-Deno.test("chain", () => {
-  const result = chain([1, 2], [3, 4], [5]);
-  const expected = [1, 2, 3, 4, 5];
-  assertEquals([...result], expected);
-  assertType<IsExact<typeof result, Iterable<number>>>(true);
+Deno.test("chain", async (t) => {
+  await t.step("uniform iterables", () => {
+    const result = chain([1, 2], [3, 4], [5]);
+    const expected = [1, 2, 3, 4, 5];
+    assertEquals([...result], expected);
+    assertType<IsExact<typeof result, Iterable<number>>>(true);
+  });
+
+  await t.step("malform iterables", () => {
+    const result = chain([1, 2], ["a", "b"], [true]);
+    const expected = [1, 2, "a", "b", true];
+    assertEquals([...result], expected);
+    assertType<IsExact<typeof result, Iterable<number | string | boolean>>>(
+      true,
+    );
+  });
 });