improve Do generator by not requiring an async generator for Tasks

cevr · May 12, 2023 · a54c565 · a54c565
1 parent 675ddc0
commit a54c565
Show file tree

Hide file tree

Showing 6 changed files with 192 additions and 74 deletions.
diff --git a/.changeset/early-goats-vanish.md b/.changeset/early-goats-vanish.md
@@ -0,0 +1,5 @@
+---
+"ftld": minor
+---
+
+improve `Do` by not requiring async generators for async computations
diff --git a/README.md b/README.md
@@ -732,6 +732,10 @@ const settle: SettledResult<SomeError | OtherError | Error, number>[] =
 
 `Do` is a utility that allows you to unwrap monadic values in a synchronous manner. It's useful for working with `Task` and `Result` types, but can be used with any monadic type. Provides the same benefits as async/await, albeit with a more cumbersome syntax.
 
+It handles `Task`, `Result`, `Option` and any `PromiseLike` types, and will short-circuit on the first `Err` value.
+
+If there are any `Task` or `PromiseLike` types, it will return a `Task`. Otherwise, it will return a `Result`.
+
 ```ts
 import { Do, Task, Result } from "ftld";
 // non Do
@@ -752,8 +756,8 @@ function doSomething() {
 }
 
 // Do
-function doSomething() {
-  return Do(async function* ($) {
+function doSomething(): Task<unknown, unknown> {
+  return Do(function* ($) {
     const a = yield* $(
       Task.from(() => {
         //...
@@ -778,9 +782,7 @@ function doSomething() {
   });
 }
 
-// non async Do
-
-function doSomething() {
+function doSomething(): Result<unknown, unknown> {
   return Do(function* ($) {
     const a = yield* $(
       Result.from(() => {

diff --git a/lib/do.test.ts b/lib/do.test.ts
@@ -27,28 +27,120 @@ describe("Do", () => {
         )
       );
 
-      return `${a + b}`;
+      const c = yield* $(Option.Some(1));
+
+      return `${a + b + c}`;
     });
 
     expectTypeOf(result).toMatchTypeOf<
-      Result<SomeError | OtherError, string>
+      Result<SomeError | OtherError | UnwrapNoneError, string>
     >();
 
-    expect(result).toEqual(Result.Ok("2"));
+    expect(result).toEqual(Result.Ok("3"));
   });
 
   it("works with Tasks", async () => {
-    const result = await Do(async function* ($) {
-      const a = yield* $(Task.Ok(1));
-      const b = yield* $(Task.Ok(2));
-      return a + b;
+    const result = Do(function* ($) {
+      const a = yield* $(
+        Task.from(
+          () => 1,
+          () => new OtherError()
+        )
+      );
+      const b = yield* $(
+        Task.from(
+          () => 1,
+          () => new SomeError()
+        )
+      );
+      const c = yield* $(Option.Some(1));
+      return a + b + c;
+    });
+
+    expectTypeOf(result).toMatchTypeOf<
+      Task<SomeError | OtherError | UnwrapNoneError, number>
+    >();
+
+    expect(await result).toEqual(Result.Ok(3));
+  });
+
+  it("returns a task if it contains any promises", async () => {
+    const result = Do(function* ($) {
+      const a = yield* $(
+        Result.from(
+          () => 1,
+          () => new OtherError()
+        )
+      );
+      const b = yield* $(
+        Result.from(
+          () => 1,
+          () => new SomeError()
+        )
+      );
+      const c = yield* $(Option.from(1 as number | null));
+      const d = yield* $(Promise.resolve(1));
+      return a + b + c + d;
+    });
+
+    expectTypeOf(result).toMatchTypeOf<Task<unknown, number>>();
+
+    expect(await result).toEqual(Result.Ok(4));
+  });
+
+  it("handles Task errors", async () => {
+    const result = Do(function* ($) {
+      const a = yield* $(
+        Task.from(
+          () => 1,
+          () => new OtherError()
+        )
+      );
+      const b = yield* $(
+        Task.from(
+          () => {
+            throw 1;
+            return 1;
+          },
+          () => new SomeError()
+        )
+      );
+      const c = yield* $(Option.Some(1));
+      return a + b + c;
     });
 
-    expectTypeOf(result).toEqualTypeOf<Result<never, number>>();
+    expectTypeOf(result).toMatchTypeOf<
+      Task<SomeError | OtherError | UnwrapNoneError, number>
+    >();
 
-    expect(result).toEqual(Result.Ok(3));
+    expect(await result).toEqual(Result.Err(new SomeError()));
   });
 
+  it('handles promise errors', async () => {
+    const result = Do(function* ($) {
+      const a = yield* $(
+        Task.from(
+          () => 1,
+          () => new OtherError()
+        )
+      );
+      const b = yield* $(
+        Task.from(
+          () => Promise.reject(1),
+          () => new SomeError()
+        )
+      );
+      const c = yield* $(Option.Some(1));
+      return a + b + c;
+    });
+
+    expectTypeOf(result).toMatchTypeOf<
+      Task<SomeError | OtherError | UnwrapNoneError, number>
+    >();
+
+    expect(await result).toEqual(Result.Err(new SomeError()));
+  })
+
   it("should error if any of the monads are errors", () => {
     const result = Do(function* ($) {
       const a = yield* $(Result.Ok(1));
@@ -58,7 +150,7 @@ describe("Do", () => {
         })
       );
 
-      return Result.Ok(a + b);
+      return a + b;
     });
 
     expectTypeOf(result).toMatchTypeOf<Result<unknown, number>>();
@@ -68,7 +160,7 @@ describe("Do", () => {
     const none = Do(function* ($) {
       const a = yield* $(Option.Some(1));
       const b = yield* $(Option.from(null as number | null));
-      return a + b
+      return a + b;
     });
 
     expectTypeOf(none).toMatchTypeOf<Result<UnwrapNoneError, number>>();

diff --git a/lib/do.ts b/lib/do.ts
@@ -1,8 +1,8 @@
+import type { Option, UnwrapNoneError } from "./option";
+import type { Monad } from "./internals";
+import { isPromiseLike } from "./internals";
 import { Task } from "./task";
-import { UnwrapNoneError } from "./option";
-import type { Option } from "./option";
 import { Result } from "./result";
-import { isPromiseLike } from "./internals";
 
 class Gen<T, A> implements Generator<T, A> {
   called = false;
@@ -72,14 +72,14 @@ class AsyncGen<T, A> implements AsyncGenerator<T, A> {
   }
 }
 
-class UnwrapGen<E, A> {
-  declare _E: E;
+class UnwrapGen<A> {
+  declare _E: UnwrapError<A>;
   constructor(readonly value: unknown) {}
   [Symbol.iterator]() {
-    return new Gen<this, A>(this);
+    return new Gen<this, UnwrapValue<A>>(this);
   }
   [Symbol.asyncIterator]() {
-    return new AsyncGen<this, A>(this);
+    return new AsyncGen<this, UnwrapValue<A>>(this);
   }
 }
 
@@ -107,59 +107,37 @@ type UnwrapError<A> = [A] extends [never]
   ? UnwrapError<A>
   : unknown;
 
-export type Unwrapper = <const A>(
-  a: A
-) => UnwrapGen<UnwrapError<A>, UnwrapValue<A>>;
-
-export function Do<Gen extends UnwrapGen<unknown, unknown>, T>(
-  f: ($: Unwrapper) => Generator<Gen, T, never>
-): Result<
-  [Gen] extends [never]
-    ? never
-    : [Gen] extends [UnwrapGen<infer E, any>]
-    ? E
-    : never,
-  UnwrapValue<T>
->;
-export function Do<Gen extends UnwrapGen<unknown, unknown>, T>(
-  f: ($: Unwrapper) => AsyncGenerator<Gen, T, never>
-): Task<
-  [Gen] extends [never]
-    ? never
-    : [Gen] extends [UnwrapGen<infer E, unknown>]
-    ? E
-    : never,
-  UnwrapValue<T>
->;
-export function Do<T, Gen extends UnwrapGen<unknown, unknown>>(
-  f: ($: Unwrapper) => Generator<Gen, T, any> | AsyncGenerator<Gen, T, any>
-) {
-  const iterator = f((x: unknown) => new UnwrapGen(x));
-
-  const state = iterator.next();
-  if (isPromiseLike(state)) {
-    // @ts-expect-error
-    const run = async (state: any) => {
-      if (state.done) {
-        const next = unwrap(getGeneratorValue(state.value));
-        return isPromiseLike(next) ? await next : next;
-      }
-      const next = unwrap(getGeneratorValue(state.value));
-      const value = iterator.next(isPromiseLike(next) ? await next : next);
-      return run(isPromiseLike(value) ? await value : value);
-    };
+export type Unwrapper = <A>(a: A) => UnwrapGen<A>;
 
-    return Task.from(async () => run(await state)) as any;
-  }
+export function Do<T, Gen extends UnwrapGen<unknown>>(
+  f: ($: Unwrapper) => Generator<Gen, T, any>
+): Collect<UnionToTuple<Gen>, UnwrapValue<T>> {
+  const iterator = f((x: unknown) => new UnwrapGen(x));
 
   // @ts-expect-error
   const run = (state: any) => {
-    return state.done
-      ? unwrap(getGeneratorValue(state.value))
-      : run(iterator.next(unwrap(getGeneratorValue(state.value))));
+    if (isPromiseLike(state)) {
+      return state.then(run);
+    }
+    if (state.done) {
+      return unwrap(getGeneratorValue(state.value));
+    }
+
+    const next = unwrap(getGeneratorValue(state.value));
+    if (isPromiseLike(next)) {
+      return next.then((value) => run(iterator.next(value)));
+    }
+    return run(iterator.next(next));
   };
 
-  return Result.from(() => run(state)) as any;
+  const res = Result.from(() => run(iterator.next()));
+  if (res.isOk()) {
+    const val = res.unwrap();
+    if (isPromiseLike(val)) {
+      return Task.from(val) as any;
+    }
+  }
+  return res as any;
 }
 
 function isUnwrapable(x: unknown): x is {
@@ -180,3 +158,39 @@ function unwrap(x: unknown): unknown {
 function getGeneratorValue(x: unknown): unknown {
   return x instanceof UnwrapGen ? x.value : x;
 }
+
+type CollectErrors<T extends any[]> = {
+  [K in keyof T]: T[K] extends UnwrapGen<infer Value>
+    ? UnwrapError<Value>
+    : never;
+}[number];
+
+// if the generator includes any Tasks, the return type will be a Task
+// otherwise it will be a Result
+type Collect<T, V> = T extends Array<
+  UnwrapGen<Exclude<Monad<unknown, unknown>, Task<unknown, unknown>>>
+>
+  ? Result<CollectErrors<T>, V>
+  : T extends Array<UnwrapGen<Monad<unknown, unknown> | PromiseLike<unknown>>>
+  ? Task<CollectErrors<T>, V>
+  : T extends Array<UnwrapGen<PromiseLike<unknown>>>
+  ? Task<CollectErrors<T>, V>
+  : never;
+
+type UnionToIntersection<U> = (U extends any ? (k: U) => void : never) extends (
+  k: infer I
+) => void
+  ? I
+  : never;
+
+type UnionToOvlds<U> = UnionToIntersection<
+  U extends any ? (f: U) => void : never
+>;
+
+type PopUnion<U> = UnionToOvlds<U> extends (a: infer A) => void ? A : never;
+
+type IsUnion<T> = [T] extends [UnionToIntersection<T>] ? false : true;
+
+type UnionToTuple<T, A extends unknown[] = []> = IsUnion<T> extends true
+  ? UnionToTuple<Exclude<T, PopUnion<T>>, [PopUnion<T>, ...A]>
+  : [T, ...A];
diff --git a/lib/internals.ts b/lib/internals.ts
@@ -1,5 +1,11 @@
+import type { Option } from "./option";
+import type { Result } from "./result";
+import type { Task } from "./task";
+
 export function isPromiseLike<T>(value: unknown): value is PromiseLike<T> {
   return typeof value === "object" && value !== null && "then" in value;
 }
 
-export declare const _tag: unique symbol;
+export declare const _tag: unique symbol;
+
+export type Monad<E, A> = Option<A> | Result<E, A> | Task<E, A>;
diff --git a/lib/task.ts b/lib/task.ts
@@ -1,4 +1,5 @@
-import { _tag, isPromiseLike } from "./internals";
+import type { _tag, Monad } from "./internals";
+import { isPromiseLike } from "./internals";
 import { identity, isOption, isResult } from "./utils";
 import { Result } from "./result";
 import type { Err, SettledResult } from "./result";
@@ -1055,5 +1056,3 @@ const maybeBoolToInt = (value: boolean | number) => {
   }
   return value;
 };
-
-type Monad<E, A> = Option<A> | Result<E, A> | Task<E, A>;