fix: improve type inference for groupBy and indexBy functions (#205)

src/groupBy.ts

@@ -1,12 +1,13 @@
 import { AsyncFunctionException } from "./_internal/error";
 import { isAsyncIterable, isIterable, isPromise } from "./_internal/utils";
 import reduce from "./reduce";
+import type Cast from "./types/Cast";
+import type Equals from "./types/Equals";
+import type { GetKeyOf } from "./types/GetKeyOf";
 import type IterableInfer from "./types/IterableInfer";
 import type iterableInfer from "./types/IterableInfer";
 import type Key from "./types/Key";
-import type { Prettify } from "./types/Prettify";
 import type ReturnValueType from "./types/ReturnValueType";
-import type { GroupBy } from "./types/groupBy";
 
 /**
  * Splits Iterable/AsyncIterable into sets, grouped by the result of running each value through `f`.
@@ -37,22 +38,68 @@ import type { GroupBy } from "./types/groupBy";
  *
  * {@link https://codesandbox.io/s/fxts-groupby-v8q3b | Try It}
  */
+
+function groupBy<A extends Key>(
+  f: (a: A) => A,
+  iterable: Iterable<A>,
+): { [K in A]: K[] };
+
+function groupBy<A extends Key>(
+  f: (a: A) => A | Promise<A>,
+  iterable: AsyncIterable<A>,
+): Promise<{ [K in A]: K[] }>;
+
+function groupBy<A extends object, B extends Key & A[keyof A]>(
+  f: (a: A) => B,
+  iterable: Iterable<A>,
+): {
+  [K in B]: (A & { [K2 in GetKeyOf<A, B>]: K })[];
+};
+
+function groupBy<A extends object, B extends Key & A[keyof A]>(
+  f: (a: A) => B | Promise<B>,
+  iterable: AsyncIterable<A>,
+): Promise<{
+  [K in B]: (A & { [K2 in GetKeyOf<A, B>]: K })[];
+}>;
+
+function groupBy<
+  I extends Iterable<unknown> | AsyncIterable<unknown>,
+  F extends (a: IterableInfer<I>) => any,
+>(
+  f: F,
+): (iterable: I) => ReturnValueType<
+  I,
+  Equals<Awaited<ReturnType<F>>, IterableInfer<I>> extends 1
+    ? {
+        [key1 in Awaited<ReturnType<F>>]: key1[];
+      }
+    : {
+        [key1 in Awaited<ReturnType<F>>]: (IterableInfer<I> & {
+          [key2 in GetKeyOf<Cast<IterableInfer<I>, object>, key1>]: key1;
+        })[];
+      }
+>;
+
+function groupBy<A extends Key, B extends Iterable<A> | AsyncIterable<A>>(
+  f: (a: A) => A | Promise<A>,
+): (iterable: B) => ReturnValueType<B, { [K in A]: K[] }>;
 function groupBy<A, B extends Key>(
   f: (a: A) => B,
   iterable: Iterable<A>,
-): Prettify<GroupBy<A, B>>;
+): { [K in B]: A[] };
 
 function groupBy<A, B extends Key>(
   f: (a: A) => B | Promise<B>,
   iterable: AsyncIterable<A>,
-): Promise<Prettify<GroupBy<A, B>>>;
+): Promise<{ [K in B]: A[] }>;
 
 function groupBy<
   A extends Iterable<unknown> | AsyncIterable<unknown>,
   B extends Key,
 >(
   f: (a: IterableInfer<A>) => B | Promise<B>,
-): (iterable: A) => ReturnValueType<A, Prettify<GroupBy<IterableInfer<A>, B>>>;
+): (iterable: A) => ReturnValueType<A, { [K in B]: IterableInfer<A>[] }>;
 
 function groupBy<
   A extends Iterable<unknown> | AsyncIterable<unknown>,
@@ -61,14 +108,16 @@ function groupBy<
   f: (a: IterableInfer<A>) => B | Promise<B>,
   iterable?: A,
 ):
-  | Prettify<GroupBy<A, B>>
-  | Promise<Prettify<GroupBy<A, B>>>
-  | ((iterable: A) => ReturnValueType<A, Prettify<GroupBy<A, B>>>) {
+  | { [K in B]: IterableInfer<A>[] }
+  | Promise<{ [K in B]: IterableInfer<A>[] }>
+  | ((iterable: A) => ReturnValueType<A, { [K in B]: IterableInfer<A>[] }>) {
   if (iterable === undefined) {
-    return (iterable: A): ReturnValueType<A, GroupBy<A, B>> => {
-      return groupBy(f, iterable as any) as unknown as ReturnValueType<
+    return (
+      iterable: A,
+    ): ReturnValueType<A, { [K in B]: IterableInfer<A>[] }> => {
+      return groupBy(f, iterable as any) as ReturnValueType<
         A,
-        GroupBy<A, B>
+        { [K in B]: IterableInfer<A>[] }
       >;
     };
   }
@@ -85,19 +134,18 @@ function groupBy<
       },
       obj,
       iterable,
-    ) as unknown as GroupBy<A, B>;
+    );
   }
 
   if (isAsyncIterable<iterableInfer<A>>(iterable)) {
-    const reulst = reduce(
+    return reduce(
       async (group, a) => {
         const key = await f(a);
         return (group[key] || (group[key] = [])).push(a), group;
       },
       obj,
       iterable,
     );
-    return reulst as unknown as Promise<GroupBy<A, B>>;
   }
 
   throw new TypeError("'iterable' must be type of Iterable or AsyncIterable");

src/indexBy.ts

@@ -1,6 +1,9 @@
 import { AsyncFunctionException } from "./_internal/error";
 import { isAsyncIterable, isIterable, isPromise } from "./_internal/utils";
 import reduce from "./reduce";
+import type Cast from "./types/Cast";
+import type Equals from "./types/Equals";
+import type { GetKeyOf } from "./types/GetKeyOf";
 import type IterableInfer from "./types/IterableInfer";
 import type Key from "./types/Key";
 import type ReturnValueType from "./types/ReturnValueType";
@@ -28,27 +31,58 @@ import type ReturnValueType from "./types/ReturnValueType";
  *
  * {@link https://codesandbox.io/s/fxts-indexby-zpeok | Try It}
  */
-function indexBy<A, B extends Key>(
+
+function indexBy<A extends Key>(
+  f: (a: A) => A,
+  iterable: Iterable<A>,
+): { [K in A]: K };
+
+function indexBy<A extends Key>(
+  f: (a: A) => A | Promise<A>,
+  iterable: AsyncIterable<A>,
+): Promise<{ [K in A]: K }>;
+
+function indexBy<A extends object, B extends Key & A[keyof A]>(
   f: (a: A) => B,
   iterable: Iterable<A>,
 ): {
-  [K in B]: A extends object
-    ? {
-        [K2 in keyof A]: A[K2] extends B ? K : A[K2];
-      }
-    : A;
+  [K in B]: A & { [K2 in GetKeyOf<A, B>]: K };
 };
 
-function indexBy<A, B extends Key>(
+function indexBy<A extends object, B extends Key & A[keyof A]>(
   f: (a: A) => B | Promise<B>,
   iterable: AsyncIterable<A>,
 ): Promise<{
-  [K in B]: A extends object
+  [K in B]: A & { [K2 in GetKeyOf<A, B>]: K };
+}>;
+
+function indexBy<
+  I extends Iterable<unknown> | AsyncIterable<unknown>,
+  F extends (a: IterableInfer<I>) => any,
+>(
+  f: F,
+): (iterable: I) => ReturnValueType<
+  I,
+  Equals<Awaited<ReturnType<F>>, IterableInfer<I>> extends 1
     ? {
-        [K2 in keyof A]: A[K2] extends B ? K : A[K2];
+        [key1 in Awaited<ReturnType<F>>]: key1;
       }
-    : A;
-}>;
+    : {
+        [key1 in Awaited<ReturnType<F>>]: IterableInfer<I> & {
+          [key2 in GetKeyOf<Cast<IterableInfer<I>, object>, key1>]: key1;
+        };
+      }
+>;
+
+function indexBy<A, B extends Key>(
+  f: (a: A) => B,
+  iterable: Iterable<A>,
+): { [K in B]: A };
+
+function indexBy<A, B extends Key>(
+  f: (a: A) => B | Promise<B>,
+  iterable: AsyncIterable<A>,
+): Promise<{ [K in B]: A }>;
 
 function indexBy<
   A extends Iterable<unknown> | AsyncIterable<unknown>,

src/types/Equals.ts

@@ -1,7 +1,16 @@
-type Equals<A1, A2> = (<A>() => A extends A2 ? 1 : 0) extends <
+type _Equals1<A1, A2> = A1 extends A2 ? (A2 extends A1 ? 1 : 0) : 0;
+type _Equals2<A1, A2> = (<A>() => A extends A2 ? 1 : 0) extends <
   A,
 >() => A extends A1 ? 1 : 0
   ? 1
   : 0;
 
+// type Test1 = _Equals1<{ a: number; b: string }, { a: number } & { b: string }>;
+// type Test2 = _Equals2<{ a: number; b: string }, { a: number } & { b: string }>;
+//
+// const test1: Test1 = 1;
+// const test2: Test2 = 1;
+
+type Equals<A1, A2> = _Equals1<A1, A2> extends 1 ? 1 : _Equals2<A1, A2>;
+
 export default Equals;

src/types/ExcludeObject.ts

@@ -14,7 +14,7 @@ type LastOf<T> = UnionToIntersection<
 
 type Push<T extends any[], V> = [...T, V];
 
-type TuplifyUnion<
+export type TuplifyUnion<
   T,
   L = LastOf<T>,
   N = [T] extends [never] ? true : false,

src/types/GetKeyOf.ts

@@ -0,0 +1,30 @@
+import type { TuplifyUnion } from "./ExcludeObject";
+
+type IsSubtype<A, B> = A extends B ? (B extends A ? false : true) : false;
+
+type IsUnionKey<T> = IsSubtype<T, string> extends true
+  ? true
+  : IsSubtype<T, number> extends true
+  ? true
+  : IsSubtype<T, symbol> extends true
+  ? true
+  : false;
+
+export type GetKeyOf<
+  T extends object,
+  V extends T[keyof T],
+  R = Exclude<
+    {
+      [K in keyof T]: V extends T[K]
+        ? IsUnionKey<T[K]> extends true
+          ? K
+          : never
+        : never;
+    }[keyof T],
+    never
+  >,
+> = TuplifyUnion<R>["length"] extends 1 ? R : never;
+
+/*
+ * type Test = GetKeyOf<{a: 1, b: 2}, 1> // a
+ */

src/types/ReturnPartitionType.ts

@@ -1,14 +1,11 @@
+import type Awaited from "./Awaited";
 import type IterableInfer from "./IterableInfer";
 import type ReturnValueType from "./ReturnValueType";
-import type { GroupBy } from "./groupBy";
 
 type ReturnPartitionType<T extends Iterable<unknown> | AsyncIterable<unknown>> =
   ReturnValueType<
     T,
-    [
-      Awaited<GroupBy<IterableInfer<T>, "true" | "false">["true"]>,
-      Awaited<GroupBy<IterableInfer<T>, "true" | "false">["false"]>,
-    ]
+    [Awaited<IterableInfer<T>>[], Awaited<IterableInfer<T>>[]]
   >;
 
 export default ReturnPartitionType;

type-check/groupBy.test.ts

@@ -45,7 +45,8 @@ type Res5To8 = {
   }[];
 };
 
-type Res9To12 = { a: "a"[]; b: "b"[]; c: "c"[] };
+type NarrowedAlphabetListResult = { a: "a"[]; b: "b"[]; c: "c"[] };
+type NormalAlphabetListResult = { [p: string]: ("a" | "b" | "c")[] };
 
 const list1 = [
   { type: "a", value: 1 },
@@ -85,7 +86,46 @@ const res12 = pipe(
   list2,
   toAsync,
   groupBy((a) => a),
-); // Promise<Res>
+);
+
+const res13 = groupBy((a) => a + String(Math.random()), list2);
+
+const res14 = pipe(
+  list2,
+  groupBy((a) => a + String(Math.random())),
+);
+
+const res15 = groupBy((a) => a + String(Math.random()), toAsync(list2));
+const res16 = pipe(
+  list2,
+  toAsync,
+  groupBy((a) => a + String(Math.random())),
+);
+
+type Data2 = {
+  id: number;
+  value: "a" | "b" | "c";
+  name: "a" | "b" | "c";
+};
+
+const source2: Data2[] = [
+  { id: 1, value: "a", name: "c" },
+  { id: 2, value: "b", name: "c" },
+  { id: 3, value: "c", name: "b" },
+  { id: 4, value: "c", name: "a" },
+];
+
+const res17 = groupBy((a) => a.value, source2);
+const res18 = pipe(
+  source2,
+  groupBy((a) => a.value),
+);
+const res19 = groupBy((a) => a.value, toAsync(source2));
+const res20 = pipe(
+  source2,
+  toAsync,
+  groupBy((a) => a.value),
+);
 
 checks([
   check<typeof res1, { [p: string]: Data[] }, Test.Pass>(),
@@ -96,8 +136,24 @@ checks([
   check<typeof res6, Res5To8, Test.Pass>(),
   check<typeof res7, Promise<Res5To8>, Test.Pass>(),
   check<typeof res8, Promise<Res5To8>, Test.Pass>(),
-  check<typeof res9, Res9To12, Test.Pass>(),
-  check<typeof res10, Res9To12, Test.Pass>(),
-  check<typeof res11, Promise<Res9To12>, Test.Pass>(),
-  check<typeof res12, Promise<Res9To12>, Test.Pass>(),
+  check<typeof res9, NarrowedAlphabetListResult, Test.Pass>(),
+  check<typeof res10, NarrowedAlphabetListResult, Test.Pass>(),
+  check<typeof res11, Promise<NarrowedAlphabetListResult>, Test.Pass>(),
+  check<typeof res12, Promise<NarrowedAlphabetListResult>, Test.Pass>(),
+  check<typeof res13, NormalAlphabetListResult, Test.Pass>(),
+  check<typeof res14, NormalAlphabetListResult, Test.Pass>(),
+  check<typeof res15, Promise<NormalAlphabetListResult>, Test.Pass>(),
+  check<typeof res16, Promise<NormalAlphabetListResult>, Test.Pass>(),
+  check<typeof res17, { a: Data2[]; b: Data2[]; c: Data2[] }, Test.Pass>(),
+  check<typeof res18, { a: Data2[]; b: Data2[]; c: Data2[] }, Test.Pass>(),
+  check<
+    typeof res19,
+    Promise<{ a: Data2[]; b: Data2[]; c: Data2[] }>,
+    Test.Pass
+  >(),
+  check<
+    typeof res20,
+    Promise<{ a: Data2[]; b: Data2[]; c: Data2[] }>,
+    Test.Pass
+  >(),
 ]);