Repository: marpple/FxTS (ts)

Issue:
`flat`, `flatMap` is inferred incorrectly when iterable is nested
# Bug Report


The result of `flatMap` inferred incorrectly when iterable is type of `AsyncIterableIterator[]`.


### 💻 Code


```ts
import { pipe, flatMap, toArray, toAsync, flat } from "@fxts/core";

async function* asyncGen(seed: number) {
  let i = 0;
  yield seed + i++;
  yield seed + i++;
  yield seed + i++;
}

(async () => {
  const asyncA = asyncGen(1);
  const asyncB = asyncGen(10);
  const asyncC = asyncGen(100);

  const flatMapResult = pipe(
    [asyncA, asyncB, asyncC],
    flatMap((v) => v),
    toArray
  );
  // flatMapResult is inferred as number[]
  // but actually flatMapResult is AsyncGenerator<number>[]
  console.log("flatMapResult", flatMapResult);
  console.log(
    "awaited first element of flatMapResult",
    await ((flatMapResult[0] as unknown) as AsyncIterator<number>).next()
  );

  const asyncFlatMapResult = pipe(
    [asyncA, asyncB, asyncC],
    toAsync,
    flatMap((v) => v),
    toArray
  );
  // asyncFlatMapResult is inferred as Promise<number[]>
  // but actually asyncFlatMapResult is Promise<AsyncGenerator<number>[]>
  console.log("asyncFlatMapResult", asyncFlatMapResult);
  console.log("awaited asyncFlatMapResult", await asyncFlatMapResult);

  const flatIterator = flat([asyncA, asyncB, asyncC]);
  const flatResultArray = [...flatIterator];
  console.log("flatResultArray", flatResultArray); // expected: number[] ,actual result: AsyncGenerator<number>[]
})();

```

you can run code at https://codesandbox.io/s/fxts-flatmap-issue-n0sw9?file=/src/index.ts

Pull request description:
fix: fix incorrect type inference of flat and flatMap functions
Fixes #142



Patch:
diff --git a/src/Lazy/dropRight.ts b/src/Lazy/dropRight.ts
new file mode 100644
index 00000000..5ba104e1
--- /dev/null
+++ b/src/Lazy/dropRight.ts
@@ -0,0 +1,152 @@
+import isArray from "../isArray";
+import isString from "../isString";
+import toArray from "../toArray";
+import ReturnIterableIteratorType from "../types/ReturnIterableIteratorType";
+import { isAsyncIterable, isIterable } from "../_internal/utils";
+import concurrent, { isConcurrent } from "./concurrent";
+
+function* stringDropRight(length: number, str: string) {
+  const arr = Array.from(str);
+  for (let i = 0; i < str.length - length; i++) {
+    yield arr[i];
+  }
+}
+
+function* arrayDropRight<T>(length: number, arr: Array<T>) {
+  for (let i = 0; i < arr.length - length; i++) {
+    yield arr[i];
+  }
+}
+
+function* sync<T>(length: number, iterable: Iterable<T>) {
+  const arr = toArray(iterable);
+  for (let i = 0; i < arr.length - length; i++) {
+    yield arr[i];
+  }
+}
+
+async function* asyncSequential<T>(length: number, iterable: AsyncIterable<T>) {
+  const arr = await toArray(iterable);
+  for (let i = 0; i < arr.length - length; i++) {
+    yield arr[i];
+  }
+}
+
+function async<A>(
+  length: number,
+  iterable: AsyncIterable<A>,
+): AsyncIterableIterator<A> {
+  let iterator: AsyncIterator<A>;
+  return {
+    [Symbol.asyncIterator]() {
+      return this;
+    },
+
+    next(_concurrent: any) {
+      if (iterator === undefined) {
+        iterator = isConcurrent(_concurrent)
+          ? asyncSequential(length, concurrent(_concurrent.length, iterable))
+          : asyncSequential(length, iterable);
+      }
+
+      return iterator.next(_concurrent);
+    },
+  };
+}
+
+/**
+ * Returns all but the last `length` elements of the given iterable.
+ *
+ * @example
+ * ```ts
+ * const iter = dropRight(2, [1, 2, 3, 4]);
+ * iter.next() // {done:false, value: 3}
+ * iter.next() // {done:false, value: 4}
+ * iter.next() // {done:true, value: undefined}
+ *
+ * // with pipe
+ * pipe(
+ *  [1, 2, 3, 4],
+ *  dropRight(2),
+ *  toArray,
+ * ); // [1, 2]
+ *
+ * pipe(
+ *  "abcde",
+ *  dropRight(2),
+ *  toArray,
+ * ); // ["a", "b", "c"]
+ *
+ * await pipe(
+ *  Promise.resolve([1, 2, 3, 4]),
+ *  dropRight(2),
+ *  toArray,
+ * ); // [1, 2]
+ *
+ * // with toAsync
+ * await pipe(
+ *  [Promise.resolve(1), Promise.resolve(2), Promise.resolve(3), Promise.resolve(4)],
+ *  toAsync,
+ *  dropRight(2),
+ *  toArray,
+ * ); // [1, 2]
+ * ```
+ *
+ * see {@link https://fxts.dev/docs/pipe | pipe}, {@link https://fxts.dev/docs/toAsync | toAsync},
+ * {@link https://fxts.dev/docs/toArray | toArray}
+ */
+// prettier-ignore
+function dropRight<A>(
+  length: number,
+  iterable: Iterable<A>
+): IterableIterator<A>;
+
+function dropRight<A>(
+  length: number,
+  iterable: AsyncIterable<A>,
+): AsyncIterableIterator<A>;
+
+function dropRight<A extends Iterable<unknown> | AsyncIterable<unknown>>(
+  length: number,
+): (iterable: A) => ReturnIterableIteratorType<A>;
+
+function dropRight<A extends Iterable<unknown> | AsyncIterable<unknown>>(
+  length: number,
+  iterable?: A,
+):
+  | IterableIterator<A>
+  | AsyncIterableIterator<A>
+  | ((iterable: A) => ReturnIterableIteratorType<A>) {
+  if (iterable === undefined) {
+    return (iterable: A): ReturnIterableIteratorType<A> => {
+      return dropRight(
+        length,
+        iterable as any,
+      ) as ReturnIterableIteratorType<A>;
+    };
+  }
+
+  if (length < 1) {
+    throw new TypeError("'length' must be greater than 1");
+  }
+
+  if (isArray(iterable)) {
+    return arrayDropRight<A>(length, iterable);
+  }
+
+  if (isString(iterable)) {
+    return stringDropRight(length, iterable) as any;
+  }
+
+  if (isIterable<A>(iterable)) {
+    return sync(length, iterable);
+  }
+
+  if (isAsyncIterable<A>(iterable)) {
+    return async(length, iterable);
+  }
+
+  throw new TypeError("'iterable' must be type of Iterable or AsyncIterable");
+}
+
+export default dropRight;
diff --git a/src/Lazy/flat.ts b/src/Lazy/flat.ts
index ef879bb3..c6d70898 100644
--- a/src/Lazy/flat.ts
+++ b/src/Lazy/flat.ts
@@ -1,12 +1,21 @@
 import concurrent, { isConcurrent } from "./concurrent";
 import last from "../last";
 import IterableInfer from "../types/IterableInfer";
-import ReturnIterableIteratorType from "../types/ReturnIterableIteratorType";
-import DeepFlat from "../types/DeepFlat";
+import { type DeepFlat, DeepFlatSync } from "../types/DeepFlat";
 import { empty, isAsyncIterable, isIterable } from "../_internal/utils";
 import { Reject, Resolve } from "../types/Utils";
 import append from "./append";
 import concat from "./concat";
+import Awaited from "../types/Awaited";
+
+type ReturnFlatType<
+  A extends Iterable<unknown> | AsyncIterable<unknown>,
+  B extends number = 1,
+> = A extends Iterable<unknown>
+  ? IterableIterator<DeepFlatSync<IterableInfer<A>, B>>
+  : A extends AsyncIterable<unknown>
+  ? AsyncIterableIterator<DeepFlat<Awaited<IterableInfer<A>>, B>>
+  : never;
 
 const isFlatAble = (a: unknown) => typeof a !== "string" && isIterable(a);
 
@@ -221,7 +230,7 @@ function flat<
 >(
   iterator: A,
   depth?: B
-): ReturnIterableIteratorType<A, DeepFlat<IterableInfer<A>, B>>;
+): ReturnFlatType<A, B>;
 
 function flat<A extends Iterable<unknown> | AsyncIterable<unknown>>(
   iterable: Iterable<A> | AsyncIterable<A>,
diff --git a/src/Lazy/flatMap.ts b/src/Lazy/flatMap.ts
index 76dee4f6..1a8e61ff 100644
--- a/src/Lazy/flatMap.ts
+++ b/src/Lazy/flatMap.ts
@@ -2,14 +2,17 @@ import map from "./map";
 import flat from "./flat";
 import { isAsyncIterable, isIterable } from "../_internal/utils";
 import Awaited from "../types/Awaited";
-import ReturnIterableIteratorType from "../types/ReturnIterableIteratorType";
 import IterableInfer from "../types/IterableInfer";
-import DeepFlat from "../types/DeepFlat";
+import { type DeepFlat, type DeepFlatSync } from "../types/DeepFlat";
 
-type FlatMapValue<
-  T,
-  I extends Iterable<unknown> | AsyncIterable<unknown>,
-> = I extends AsyncIterable<any> ? Awaited<T> : T;
+type ReturnFlatMapType<
+  A extends Iterable<unknown> | AsyncIterable<unknown>,
+  B = unknown,
+> = A extends Iterable<unknown>
+  ? IterableIterator<DeepFlatSync<B, 1>>
+  : A extends AsyncIterable<unknown>
+  ? AsyncIterableIterator<DeepFlat<Awaited<B>, 1>>
+  : never;
 
 /**
  * Returns flattened Iterable/AsyncIterable of values by running each element
@@ -55,22 +58,25 @@ type FlatMapValue<
 function flatMap<A, B = unknown>(
   f: (a: A) => B,
   iterable: Iterable<A>,
-): IterableIterator<DeepFlat<B, 1>>;
+): IterableIterator<DeepFlatSync<B, 1>>;
 
 function flatMap<A, B = unknown>(
   f: (a: A) => B,
   iterable: AsyncIterable<A>,
 ): AsyncIterableIterator<DeepFlat<Awaited<B>, 1>>;
 
+function flatMap<
+  A extends Iterable<unknown> | AsyncIterable<unknown>,
+  B = unknown,
+>(f: (a: IterableInfer<A>) => B): (iterable: A) => ReturnFlatMapType<A, B>;
+
 function flatMap<
   A extends Iterable<unknown> | AsyncIterable<unknown>,
   B = unknown,
 >(
   f: (a: IterableInfer<A>) => B,
   iterable?: A,
-): (
-  iterable: A,
-) => ReturnIterableIteratorType<A, DeepFlat<FlatMapValue<B, A>, 1>>;
+): (iterable: A) => ReturnFlatMapType<A, B>;
 
 function flatMap<
   A extends Iterable<unknown> | AsyncIterable<unknown>,
@@ -78,27 +84,19 @@ function flatMap<
 >(
   f: (a: IterableInfer<A>) => B,
   iterable?: A,
-):
-  | IterableIterator<DeepFlat<FlatMapValue<B, A>, 1>>
-  | AsyncIterableIterator<DeepFlat<Awaited<B>, 1>>
-  | ((iterable: A) => ReturnIterableIteratorType<A, DeepFlat<B, 1>>) {
+): ReturnFlatMapType<A, B> | ((iterable: A) => ReturnFlatMapType<A, B>) {
   if (iterable === undefined) {
     return (iterable: A) => {
-      return flat(map(f, iterable as any)) as ReturnIterableIteratorType<
-        A,
-        DeepFlat<B, 1>
-      >;
+      return flat(map(f, iterable as any)) as any;
     };
   }
 
   if (isIterable<IterableInfer<A>>(iterable)) {
-    return flat(map(f, iterable)) as IterableIterator<
-      DeepFlat<FlatMapValue<B, A>, 1>
-    >;
+    return flat(map(f, iterable as any)) as any;
   }
 
   if (isAsyncIterable<Awaited<IterableInfer<A>>>(iterable)) {
-    return flat(map(f, iterable));
+    return flat(map(f, iterable as any)) as any;
   }
 
   throw new TypeError("'iterable' must be type of Iterable or AsyncIterable");
diff --git a/src/Lazy/index.ts b/src/Lazy/index.ts
index b1635c0c..f9a03f67 100644
--- a/src/Lazy/index.ts
+++ b/src/Lazy/index.ts
@@ -8,6 +8,7 @@ import cycle from "./cycle";
 import difference from "./difference";
 import differenceBy from "./differenceBy";
 import drop from "./drop";
+import dropRight from "./dropRight";
 import dropUntil from "./dropUntil";
 import dropWhile from "./dropWhile";
 import entries from "./entries";
@@ -50,6 +51,7 @@ export {
   difference,
   differenceBy,
   drop,
+  dropRight,
   dropUntil,
   dropWhile,
   entries,
diff --git a/src/types/DeepFlat.ts b/src/types/DeepFlat.ts
index ba639c55..1da1f51b 100644
--- a/src/types/DeepFlat.ts
+++ b/src/types/DeepFlat.ts
@@ -8,4 +8,10 @@ type DeepFlat<A, B extends number = 0> = A extends
     : DeepFlat<U, Sub<B, 1>>
   : A;
 
-export default DeepFlat;
+type DeepFlatSync<A, B extends number = 0> = A extends Iterable<infer U>
+  ? B extends 0
+    ? A
+    : DeepFlat<U, Sub<B, 1>>
+  : A;
+
+export { DeepFlat, DeepFlatSync };
diff --git a/website/function.json b/website/function.json
index 357f99c7..a2e925af 100644
--- a/website/function.json
+++ b/website/function.json
@@ -8,6 +8,7 @@
     "concurrent",
     "cycle",
     "drop",
+    "dropRight",
     "difference",
     "differenceBy",
     "dropUntil",


Test patch:
diff --git a/test/Lazy/dropRight.spec.ts b/test/Lazy/dropRight.spec.ts
new file mode 100644
index 00000000..709c40c9
--- /dev/null
+++ b/test/Lazy/dropRight.spec.ts
@@ -0,0 +1,134 @@
+import {
+  concurrent,
+  delay,
+  dropRight,
+  filter,
+  map,
+  pipe,
+  toArray,
+  toAsync,
+} from "../../src/index";
+import { Concurrent } from "../../src/Lazy/concurrent";
+
+import { callFuncAfterTime, generatorMock } from "../utils";
+
+describe("dropRightRight", function () {
+  describe("sync", function () {
+    it("should be discarded elements by length", function () {
+      const acc = [];
+      for (const a of dropRight(2, [1, 2, 3, 4, 5])) {
+        acc.push(a);
+      }
+      expect(acc).toEqual([1, 2, 3]);
+    });
+
+    it("should be discarded string by length", function () {
+      const acc = [];
+      for (const a of dropRight(2, "abcde")) {
+        acc.push(a);
+      }
+      expect(acc).toEqual(["a", "b", "c"]);
+    });
+
+    it("should be able to be used as a curried function in the pipeline", function () {
+      const res = pipe(
+        [1, 2, 3, 4, 5, 6, 7, 8],
+        map((a) => a + 10),
+        filter((a) => a % 2 === 0),
+        dropRight(2),
+        toArray,
+      );
+      expect(res).toEqual([12, 14]);
+    });
+  });
+
+  describe("async", function () {
+    it("should be discarded elements by length", async function () {
+      const acc = [];
+      for await (const a of dropRight(2, toAsync([1, 2, 3, 4, 5]))) {
+        acc.push(a);
+      }
+      expect(acc).toEqual([1, 2, 3]);
+    });
+
+    it("should be able to be used as a curried function in the pipeline", async function () {
+      const res = await pipe(
+        toAsync([1, 2, 3, 4, 5, 6, 7, 8]),
+        map((a) => a + 10),
+        filter((a) => a % 2 === 0),
+        dropRight(2),
+        toArray,
+      );
+
+      expect(res).toEqual([12, 14]);
+    });
+
+    it("should be discarded elements by length concurrently", async function () {
+      const fn = jest.fn();
+      callFuncAfterTime(fn, 400);
+      const res = await pipe(
+        toAsync([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
+        map((a) => delay(100, a)),
+        filter((a) => a % 2 === 0),
+        dropRight(2),
+        concurrent(3),
+        toArray,
+      );
+      expect(fn).toBeCalled();
+      expect(res).toEqual([2, 4, 6]);
+    }, 450);
+
+    it("should be able to handle an error when asynchronous", async function () {
+      await expect(
+        pipe(
+          toAsync([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
+          filter(() => {
+            throw new Error("err");
+          }),
+          dropRight(2),
+          toArray,
+        ),
+      ).rejects.toThrow("err");
+    });
+
+    it("should be able to handle an error when working concurrent", async function () {
+      await expect(
+        pipe(
+          toAsync([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
+          filter((a) => {
+            if (a === 1) {
+              throw new Error("err");
+            } else return true;
+          }),
+          dropRight(2),
+          concurrent(3),
+          toArray,
+        ),
+      ).rejects.toThrow("err");
+    });
+
+    it("should be able to handle an error when working concurrent - Promise.reject", async function () {
+      await expect(
+        pipe(
+          toAsync([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
+          map((a) => {
+            if (a < 3) return Promise.reject(new Error("err"));
+            else return a;
+          }),
+          dropRight(2),
+          concurrent(3),
+          toArray,
+        ),
+      ).rejects.toThrow("err");
+    });
+
+    it("should be passed concurrent object when job works concurrently", async function () {
+      const mock = generatorMock();
+      const iter = dropRight(2, mock);
+      const concurrent = Concurrent.of(2) as any;
+
+      await iter.next(concurrent);
+      expect((mock as any).getConcurrent()).toEqual(concurrent);
+    });
+  });
+});
diff --git a/type-check/Lazy/dropRight.test.ts b/type-check/Lazy/dropRight.test.ts
new file mode 100644
index 00000000..4122693b
--- /dev/null
+++ b/type-check/Lazy/dropRight.test.ts
@@ -0,0 +1,20 @@
+import * as Test from "../../src/types/Test";
+import { toAsync, dropRight, pipe } from "../../src";
+
+const { checks, check } = Test;
+
+const res1 = dropRight(2, []);
+const res2 = dropRight(2, [1, 2, 3, 4]);
+const res3 = dropRight(2, "abcd");
+const res4 = dropRight(2, toAsync([1, 2, 3, 4]));
+const res5 = pipe([1, 2, 3, 4], dropRight(2));
+const res6 = pipe(toAsync([1, 2, 3, 4]), dropRight(2));
+
+checks([
+  check<typeof res1, IterableIterator<never>, Test.Pass>(),
+  check<typeof res2, IterableIterator<number>, Test.Pass>(),
+  check<typeof res3, IterableIterator<string>, Test.Pass>(),
+  check<typeof res4, AsyncIterableIterator<number>, Test.Pass>(),
+  check<typeof res5, IterableIterator<number>, Test.Pass>(),
+  check<typeof res6, AsyncIterableIterator<number>, Test.Pass>(),
+]);
diff --git a/type-check/Lazy/flat.test.ts b/type-check/Lazy/flat.test.ts
index 205dbc58..49361a85 100644
--- a/type-check/Lazy/flat.test.ts
+++ b/type-check/Lazy/flat.test.ts
@@ -1,5 +1,5 @@
 import * as Test from "../../src/types/Test";
-import { toAsync, flat, pipe } from "../../src";
+import { flat, pipe, range, toAsync } from "../../src";
 
 const { checks, check } = Test;
 
@@ -22,6 +22,13 @@ const res13 = pipe(toAsync([1, 2, 3]), flat);
 const res14 = pipe(toAsync([1, 2, 3, [4, 5]]), flat);
 const res15 = pipe(toAsync([1, 2, 3, [4, 5, [6]]]), flat);
 
+const res16 = flat(
+  (function* () {
+    yield* range(1, 3);
+    yield toAsync(range(3, 5));
+  })(),
+);
+
 checks([
   check<typeof res1, IterableIterator<never>, Test.Pass>(),
   check<typeof res2, IterableIterator<number>, Test.Pass>(),
@@ -41,4 +48,5 @@ checks([
   check<typeof res13, AsyncIterableIterator<number>, Test.Pass>(),
   check<typeof res14, AsyncIterableIterator<number>, Test.Pass>(),
   check<typeof res15, AsyncIterableIterator<number | number[]>, Test.Pass>(), // prettier-ignore
+  check<typeof res16, IterableIterator<number | AsyncIterableIterator<number>>, Test.Pass>(), // prettier-ignore
 ]);
diff --git a/type-check/Lazy/flatMap.test.ts b/type-check/Lazy/flatMap.test.ts
index dbe04bdf..7c6fd4c3 100644
--- a/type-check/Lazy/flatMap.test.ts
+++ b/type-check/Lazy/flatMap.test.ts
@@ -1,5 +1,5 @@
 import * as Test from "../../src/types/Test";
-import { toAsync, flatMap, pipe } from "../../src";
+import { flatMap, pipe, range, toArray, toAsync } from "../../src";
 
 const { checks, check } = Test;
 
@@ -24,6 +24,12 @@ const res7 = pipe(
   flatMap(async (s) => s.split(" ")),
 );
 
+const res8 = pipe(
+  [toAsync(range(0, 5)), toAsync(range(0, 5)), toAsync(range(0, 5))],
+  flatMap((a) => a),
+  toArray,
+);
+
 checks([
   check<typeof res1, IterableIterator<string>, Test.Pass>(),
   check<typeof res2, IterableIterator<Promise<string[]>>, Test.Pass>(),
@@ -33,4 +39,5 @@ checks([
   check<typeof res5, IterableIterator<string>, Test.Pass>(),
   check<typeof res6, IterableIterator<Promise<string[]>>, Test.Pass>(), // prettier-ignore
   check<typeof res7, AsyncIterableIterator<string>, Test.Pass>(), // prettier-ignore
+  check<typeof res8, AsyncIterableIterator<number>[], Test.Pass>(),
 ]);