Repository: marpple/FxTS (ts)

Issue:
Improve the type inference of the `groupBy` and `indexBy` functions
## ⭐ Suggestion
It looks like it could be made better by improving the type inference of the `groupBy` function.

## 💻 Use Cases

```ts
const list = [
  { type: "a", value: 1 },
  { type: "b", value: 2 },
  { type: "c", value: 3 },
  { type: "a", value: 4 } ,
  { type: "b", value: 5 },
  { type: "c", value: 6 },
  { type: "a", value: 7 },
  { type: "b", value: 8 },
  { type: "c", value: 9 },
  { type: "a", value: 10 },
] as { type: "a" | "b" | "c"; value: number; }[]

const res = groupBy(a => a.type, list);

/*
AS-IS
typeof res:
{
    a: {
        type: "a" | "b" | "c";
        value: number;
    }[];
    b: {
        type: "a" | "b" | "c";
        value: number;
    }[];
    c: {
        type: "a" | "b" | "c";
        value: number;
    }[];
}
*/

/*
TO-BE
typeof res:
{
    a: {
        type: "a";
        value: number;
    }[];
    b: {
        type: "b";
        value: number;
    }[];
    c: {
        type: "c";
        value: number;
    }[];
}
*/
```

Pull request description:
fix: improve type inference for groupBy and indexBy functions
Fixes #197 

This change is a fix for the following issue: https://github.com/marpple/FxTS/pull/199#discussion_r1161289493


```ts
type Data = { id: number; state1: 'created' | 'deleted', state2: 'created' | 'deleted' };
const list: Data[] = [...];
const res = groupBy(a => a.state1, list);

// typeof res
// AS-IS: { created: { id: number; state1: 'created', state2: 'created'}[]; deleted: { id: number; state1: 'deleted', state2: 'deleted'}[]  }
// TO-BE: { created: Data[]; deleted: Data[] }
```


Patch:
diff --git a/src/groupBy.ts b/src/groupBy.ts
index 73225b56..47fa4c45 100644
--- a/src/groupBy.ts
+++ b/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,11 +134,11 @@ 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;
@@ -97,7 +146,6 @@ function groupBy<
       obj,
       iterable,
     );
-    return reulst as unknown as Promise<GroupBy<A, B>>;
   }
 
   throw new TypeError("'iterable' must be type of Iterable or AsyncIterable");
diff --git a/src/indexBy.ts b/src/indexBy.ts
index 94ab710c..48e17c4a 100644
--- a/src/indexBy.ts
+++ b/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>,
diff --git a/src/types/Equals.ts b/src/types/Equals.ts
index 23646c05..ca1438e2 100644
--- a/src/types/Equals.ts
+++ b/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;
diff --git a/src/types/ExcludeObject.ts b/src/types/ExcludeObject.ts
index 85997a9e..c3e72968 100644
--- a/src/types/ExcludeObject.ts
+++ b/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,
diff --git a/src/types/GetKeyOf.ts b/src/types/GetKeyOf.ts
new file mode 100644
index 00000000..858d408e
--- /dev/null
+++ b/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
+ */
diff --git a/src/types/ReturnPartitionType.ts b/src/types/ReturnPartitionType.ts
index 9b9b8149..38ed5824 100644
--- a/src/types/ReturnPartitionType.ts
+++ b/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;


Test patch:
diff --git a/type-check/groupBy.test.ts b/type-check/groupBy.test.ts
index 90aa3490..100d3e19 100644
--- a/type-check/groupBy.test.ts
+++ b/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
+  >(),
 ]);
diff --git a/type-check/indexBy.test.ts b/type-check/indexBy.test.ts
index 157d9ca9..188cf5c5 100644
--- a/type-check/indexBy.test.ts
+++ b/type-check/indexBy.test.ts
@@ -11,6 +11,13 @@ const source: Data[] = [
   { id: 3, value: "c" },
 ];
 
+const alphabetList = ["a", "b", "c", "d", "a", "b", "c", "d"] as Array<
+  "a" | "b" | "c" | "d"
+>;
+
+type NormalAlphabetListResult = { [p: string]: "a" | "b" | "c" | "d" };
+type NarrowedAlphabetListResult = { a: "a"; b: "b"; c: "c"; d: "d" };
+
 const res1 = indexBy((a) => a.id, source);
 const res2 = pipe(
   source,
@@ -42,6 +49,59 @@ const res5 = indexBy((a) => a.type, [
   { type: "a", value: 10 },
 ] as { type: "a" | "b" | "c"; value: number }[]);
 
+const res6 = indexBy((a) => a, alphabetList);
+
+const res7 = pipe(
+  alphabetList,
+  indexBy((a) => a),
+);
+
+const res8 = indexBy((a) => a, toAsync(alphabetList));
+const res9 = pipe(
+  alphabetList,
+  toAsync,
+  indexBy((a) => a),
+);
+
+const res10 = indexBy((a) => a + String(Math.random()), alphabetList);
+
+const res11 = pipe(
+  alphabetList,
+  indexBy((a) => a + String(Math.random())),
+);
+
+const res12 = indexBy((a) => a + String(Math.random()), toAsync(alphabetList));
+const res13 = pipe(
+  alphabetList,
+  toAsync,
+  indexBy((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 res15 = indexBy((a) => a.value, source2);
+const res16 = pipe(
+  source2,
+  indexBy((a) => a.value),
+);
+const res17 = indexBy((a) => a.value, toAsync(source2));
+const res18 = pipe(
+  source2,
+  toAsync,
+  indexBy((a) => a.value),
+);
+
 checks([
   check<typeof res1, { [index: number]: Data }, Test.Pass>(),
   check<typeof res2, { [index: number]: Data }, Test.Pass>(),
@@ -65,4 +125,16 @@ checks([
     },
     Test.Pass
   >(),
+  check<typeof res6, NarrowedAlphabetListResult, Test.Pass>(),
+  check<typeof res7, NarrowedAlphabetListResult, Test.Pass>(),
+  check<typeof res8, Promise<NarrowedAlphabetListResult>, Test.Pass>(),
+  check<typeof res9, Promise<NarrowedAlphabetListResult>, Test.Pass>(),
+  check<typeof res10, NormalAlphabetListResult, Test.Pass>(),
+  check<typeof res11, NormalAlphabetListResult, Test.Pass>(),
+  check<typeof res12, Promise<NormalAlphabetListResult>, Test.Pass>(),
+  check<typeof res13, Promise<NormalAlphabetListResult>, Test.Pass>(),
+  check<typeof res15, { a: Data2; b: Data2; c: Data2 }, Test.Pass>(),
+  check<typeof res16, { a: Data2; b: Data2; c: Data2 }, Test.Pass>(),
+  check<typeof res17, Promise<{ a: Data2; b: Data2; c: Data2 }>, Test.Pass>(),
+  check<typeof res18, Promise<{ a: Data2; b: Data2; c: Data2 }>, Test.Pass>(),
 ]);