From 648af8150e721f4abaf247df6ed3f83bd2d37cba Mon Sep 17 00:00:00 2001
From: Mikhail Kireev <29187880+kireevmp@users.noreply.github.com>
Date: Wed, 17 Jan 2024 01:08:22 +0100
Subject: [PATCH] feat(condition): simplify types

---
 src/condition/index.ts    | 142 +++++++++++++++++++-------------------
 test-typings/condition.ts | 131 +++++++++++++++++++++++++++++++++--
 2 files changed, 198 insertions(+), 75 deletions(-)

diff --git a/src/condition/index.ts b/src/condition/index.ts
index f49a7f28..7c836cef 100644
--- a/src/condition/index.ts
+++ b/src/condition/index.ts
@@ -7,91 +7,91 @@ import {
   Store,
   UnitTargetable,
   split,
+  UnitValue,
   EventCallable,
+  EventCallableAsReturnType,
 } from 'effector';
 
-type NoInfer<T> = T & { [K in keyof T]: T[K] };
-type EventAsReturnType<Payload> = any extends Payload ? Event<Payload> : never;
+type NoInfer<T> = [T][T extends any ? 0 : never];
+type NonFalsy<T> = T extends null | undefined | false | 0 | 0n | '' ? never : T;
 
-export function condition<State>(options: {
-  source: Event<State>;
-  if: ((payload: State) => boolean) | Store<boolean> | State;
-  then: UnitTargetable<NoInfer<State> | void>;
-  else: UnitTargetable<NoInfer<State> | void>;
-}): EventAsReturnType<State>;
-export function condition<State>(options: {
-  source: Store<State>;
-  if: ((payload: State) => boolean) | Store<boolean> | State;
-  then: UnitTargetable<State | void>;
-  else: UnitTargetable<State | void>;
-}): Store<State>;
-export function condition<Params, Done, Fail>(options: {
-  source: Effect<Params, Done, Fail>;
-  if: ((payload: Params) => boolean) | Store<boolean> | Params;
-  then: UnitTargetable<NoInfer<Params> | void>;
-  else: UnitTargetable<NoInfer<Params> | void>;
-}): Effect<Params, Done, Fail>;
+type SourceUnit<T> = Store<T> | Event<T> | Effect<T, any, any>;
 
-export function condition<State>(options: {
-  source: Event<State>;
-  if: ((payload: State) => boolean) | Store<boolean> | State;
-  then: UnitTargetable<NoInfer<State> | void>;
-}): EventAsReturnType<State>;
-export function condition<State>(options: {
-  source: Store<State>;
-  if: ((payload: State) => boolean) | Store<boolean> | State;
-  then: UnitTargetable<NoInfer<State> | void>;
-}): Store<State>;
-export function condition<Params, Done, Fail>(options: {
-  source: Effect<Params, Done, Fail>;
-  if: ((payload: Params) => boolean) | Store<boolean> | Params;
-  then: UnitTargetable<NoInfer<Params> | void>;
-}): Effect<Params, Done, Fail>;
+// -- Without `source`, with type guard --
+export function condition<Payload, Then extends Payload = Payload>(options: {
+  source?: undefined;
+  if: ((payload: Payload) => payload is Then) | Then;
+  then?: UnitTargetable<NoInfer<Then> | void>;
+  else?: UnitTargetable<Exclude<NoInfer<Payload>, Then> | void>;
+}): EventCallableAsReturnType<Payload>;
 
-export function condition<State>(options: {
-  source: Event<State>;
-  if: ((payload: State) => boolean) | Store<boolean> | State;
-  else: UnitTargetable<NoInfer<State> | void>;
-}): EventAsReturnType<State>;
-export function condition<State>(options: {
-  source: Store<State>;
-  if: ((payload: State) => boolean) | Store<boolean> | State;
-  else: UnitTargetable<NoInfer<State> | void>;
-}): Store<State>;
-export function condition<Params, Done, Fail>(options: {
-  source: Effect<Params, Done, Fail>;
-  if: ((payload: Params) => boolean) | Store<boolean> | Params;
-  else: UnitTargetable<NoInfer<Params> | void>;
-}): Effect<Params, Done, Fail>;
+// -- Without `source`, with BooleanConstructor --
+export function condition<
+  Payload,
+  Then extends NonFalsy<Payload> = NonFalsy<Payload>,
+>(options: {
+  source?: undefined;
+  if: BooleanConstructor;
+  then?: UnitTargetable<NoInfer<Then> | void>;
+  else?: UnitTargetable<Exclude<NoInfer<Payload>, Then> | void>;
+}): EventCallableAsReturnType<Payload>;
 
-// Without `source`
+// -- Without `source` --
+export function condition<Payload>(options: {
+  source?: undefined;
+  if: ((payload: Payload) => boolean) | Store<boolean> | NoInfer<Payload>;
+  then?: UnitTargetable<NoInfer<Payload> | void>;
+  else?: UnitTargetable<NoInfer<Payload> | void>;
+}): EventCallableAsReturnType<Payload>;
 
-export function condition<State>(options: {
-  if: ((payload: State) => boolean) | Store<boolean> | State;
-  then: UnitTargetable<NoInfer<State> | void>;
-  else: UnitTargetable<NoInfer<State> | void>;
-}): EventCallable<State>;
-export function condition<State>(options: {
-  if: ((payload: State) => boolean) | Store<boolean> | State;
-  then: UnitTargetable<NoInfer<State> | void>;
-}): EventCallable<State>;
-export function condition<State>(options: {
-  if: ((payload: State) => boolean) | Store<boolean> | State;
-  else: UnitTargetable<NoInfer<State> | void>;
-}): EventCallable<State>;
-export function condition<State>({
+// -- With `source` and type guard --
+export function condition<
+  Payload extends UnitValue<Source>,
+  Then extends Payload = Payload,
+  Source extends SourceUnit<any> = SourceUnit<Payload>,
+>(options: {
+  source: Source;
+  if: ((payload: Payload) => payload is Then) | Then;
+  then?: UnitTargetable<NoInfer<Then>>;
+  else?: UnitTargetable<Exclude<NoInfer<Payload>, Then>>;
+}): Source;
+
+// -- With `source` and BooleanConstructor --
+export function condition<
+  Payload extends UnitValue<Source>,
+  Then extends NonFalsy<Payload> = NonFalsy<Payload>,
+  Source extends SourceUnit<any> = SourceUnit<Payload>,
+>(options: {
+  source: Source;
+  if: BooleanConstructor;
+  then?: UnitTargetable<NoInfer<Then> | void>;
+  else?: UnitTargetable<Exclude<NoInfer<Payload>, Then>>;
+}): EventCallable<Payload>;
+
+// -- With `source` --
+export function condition<
+  Payload extends UnitValue<Source>,
+  Source extends SourceUnit<any> = SourceUnit<Payload>,
+>(options: {
+  source: SourceUnit<Payload>;
+  if: ((payload: Payload) => boolean) | Store<boolean> | NoInfer<Payload>;
+  then?: UnitTargetable<NoInfer<Payload> | void>;
+  else?: UnitTargetable<NoInfer<Payload> | void>;
+}): Source;
+
+export function condition<Payload>({
+  source = createEvent<Payload>(),
   if: test,
   then: thenBranch,
   else: elseBranch,
-  source = createEvent<State>(),
 }: {
-  if: ((payload: State) => boolean) | Store<boolean> | State;
-  source?: Store<State> | Event<State> | Effect<State, any, any>;
-  then?: UnitTargetable<State | void>;
-  else?: UnitTargetable<State | void>;
+  source?: SourceUnit<Payload>;
+  if: ((payload: Payload) => boolean) | Store<boolean> | Payload;
+  then?: UnitTargetable<Payload | void>;
+  else?: UnitTargetable<Payload | void>;
 }) {
   const checker =
-    is.unit(test) || isFunction(test) ? test : (value: State) => value === test;
+    is.unit(test) || isFunction(test) ? test : (value: Payload) => value === test;
 
   if (thenBranch && elseBranch) {
     split({
diff --git a/test-typings/condition.ts b/test-typings/condition.ts
index f0c81ae2..6cc8c2de 100644
--- a/test-typings/condition.ts
+++ b/test-typings/condition.ts
@@ -6,6 +6,7 @@ import {
   createStore,
   Effect,
   Event,
+  EventCallable,
   Store,
 } from 'effector';
 import { condition } from '../dist/condition';
@@ -220,25 +221,25 @@ import { condition } from '../dist/condition';
   );
 }
 
-// Disallow pass invalid type to then/else
+// Disallow pass invalid type to then/else/if
 {
   condition({
-    // @ts-expect-error
     source: createStore(0),
     if: 0,
+    // @ts-expect-error 'string' is not assignable to 'number | void'
     then: createEvent<string>(),
   });
 
   condition({
-    // @ts-expect-error
     source: createStore<boolean>(false),
+    // @ts-expect-error 'Console' is not assignable to `if`
     if: console,
     then: createEvent(),
   });
 
   condition({
-    // @ts-expect-error
     source: createStore<string>(''),
+    // @ts-expect-error 'number' is not assignable to type 'boolean'
     if: (a) => 1,
     then: createEvent(),
   });
@@ -259,3 +260,125 @@ import { condition } from '../dist/condition';
     else: fxOtherVoid,
   });
 }
+
+// allows nesting conditions
+{
+  condition({
+    source: createEvent<'a' | 'b' | 1>(),
+    if: (value): value is 'a' | 'b' => typeof value === 'string',
+    then: condition<'a' | 'b'>({
+      if: () => true,
+      then: createEvent<'a'>(),
+      else: createEvent<'b'>(),
+    }),
+  });
+}
+
+// returns `typeof source` when source is provided
+{
+  const source = createEvent<string | number>();
+
+  expectType<typeof source>(
+    condition({
+      source,
+      if: 'string?',
+      then: createEvent<void>(),
+    }),
+  );
+}
+
+// Correctly passes type to `if`
+{
+  condition({
+    source: createEvent<string>(),
+    if: (payload) => (expectType<string>(payload), true),
+    then: createEvent<string>(),
+  });
+
+  condition({
+    source: createEvent<'complex' | 'type'>(),
+    if: (payload) => (expectType<'complex' | 'type'>(payload), true),
+    then: createEvent<void>(),
+  });
+
+  condition({
+    source: createEvent<string>(),
+    // @ts-expect-error 'string' is not assignable to type 'number'
+    if: (_: number) => true,
+    then: createEvent<void>(),
+  });
+}
+
+// `Boolean` as type guard: disallows invalid type in `then`
+{
+  condition({
+    source: createEvent<string | null>(),
+    if: Boolean,
+    // @ts-expect-error 'number' is not assignable to type 'string | void'
+    then: createEvent<number>(),
+  });
+}
+
+// `Boolean` as type guard: disallows invalid type in then/else
+{
+  condition({
+    source: createEvent<string | null>(),
+    if: Boolean,
+    // @ts-expect-error 'number' is not assignable to type 'string | void'
+    then: createEvent<number>(),
+  });
+
+  condition({
+    source: createEvent<string | null>(),
+    if: Boolean,
+    // @ts-expect-error 'number' is not assignable to type 'string | void'
+    else: createEvent<number>(),
+  });
+}
+
+// `Boolean` as type guard: works for all sources
+{
+  expectType<EventCallable<string | null>>(
+    condition({
+      source: createEvent<string | null>(),
+      if: Boolean,
+      then: createEvent<string>(),
+      else: createEvent<null>(),
+    }),
+  );
+
+  expectType<EventCallable<string | null>>(
+    condition({
+      source: createStore<string | null>(null),
+      if: Boolean,
+      then: createEvent<string>(),
+      else: createEvent<null>(),
+    }),
+  );
+
+  expectType<EventCallable<string | null>>(
+    condition({
+      source: createEffect<string | null, void>(),
+      if: Boolean,
+      then: createEvent<string>(),
+      else: createEvent<null>(),
+    }),
+  );
+}
+
+// `Boolean` as type guard: disallows invalid type in then/else
+{
+  condition({
+    source: createEvent<string | null>(),
+    if: Boolean,
+    // @ts-expect-error
+    then: createEvent<number>(),
+  });
+
+  condition({
+    source: createEvent<string | null>(),
+    if: Boolean,
+    // @ts-expect-error
+    else: createEvent<number>(),
+  });
+}