From d01cdfa1d929c86ea21d0d424f260057c90d8588 Mon Sep 17 00:00:00 2001
From: Patrick Roza <contact@patrickroza.com>
Date: Mon, 12 Aug 2024 05:51:49 +0200
Subject: [PATCH] return Exit values from `Fiber.awaitAll` (#3147)

Co-authored-by: Tim <hello@timsmart.co>
---
 .changeset/dry-badgers-boil.md               |  5 +++++
 packages/effect/dtslint/Fiber.ts             | 16 ++++++++++++++++
 packages/effect/src/Fiber.ts                 |  9 ++++++++-
 packages/effect/src/internal/fiberRuntime.ts | 10 ++++++++--
 packages/effect/test/Fiber.test.ts           |  4 +++-
 5 files changed, 40 insertions(+), 4 deletions(-)
 create mode 100644 .changeset/dry-badgers-boil.md
 create mode 100644 packages/effect/dtslint/Fiber.ts

diff --git a/.changeset/dry-badgers-boil.md b/.changeset/dry-badgers-boil.md
new file mode 100644
index 00000000000..f8faab1f3f9
--- /dev/null
+++ b/.changeset/dry-badgers-boil.md
@@ -0,0 +1,5 @@
+---
+"effect": minor
+---
+
+improve: type Fiber.awaitAll as Exit<A, E>[].
diff --git a/packages/effect/dtslint/Fiber.ts b/packages/effect/dtslint/Fiber.ts
new file mode 100644
index 00000000000..36195d251ed
--- /dev/null
+++ b/packages/effect/dtslint/Fiber.ts
@@ -0,0 +1,16 @@
+import { Fiber } from "effect"
+
+declare const string: Fiber.Fiber<string>
+declare const number: Fiber.Fiber<number>
+declare const array: Array<Fiber.Fiber<string> | Fiber.Fiber<number>>
+
+// awaitAll
+
+// $ExpectType Effect<[Exit<string, never>, Exit<number, never>], never, never>
+Fiber.awaitAll([string, number])
+
+// $ExpectType Effect<(Exit<string, never> | Exit<number, never>)[], never, never>
+Fiber.awaitAll(new Set([string, number]))
+
+// $ExpectType Effect<(Exit<string, never> | Exit<number, never>)[], never, never>
+Fiber.awaitAll(array)
diff --git a/packages/effect/src/Fiber.ts b/packages/effect/src/Fiber.ts
index d7df138b35f..15f8623a3f9 100644
--- a/packages/effect/src/Fiber.ts
+++ b/packages/effect/src/Fiber.ts
@@ -275,7 +275,14 @@ export {
  * @since 2.0.0
  * @category destructors
  */
-export const awaitAll: (fibers: Iterable<Fiber<any, any>>) => Effect.Effect<void> = fiberRuntime.fiberAwaitAll
+export const awaitAll: <const T extends Iterable<Fiber<any, any>>>(
+  fibers: T
+) => Effect.Effect<
+  [T] extends [ReadonlyArray<infer U>]
+    ? number extends T["length"] ? Array<U extends Fiber<infer A, infer E> ? Exit.Exit<A, E> : never>
+    : { -readonly [K in keyof T]: T[K] extends Fiber<infer A, infer E> ? Exit.Exit<A, E> : never }
+    : Array<T extends Iterable<infer U> ? U extends Fiber<infer A, infer E> ? Exit.Exit<A, E> : never : never>
+> = fiberRuntime.fiberAwaitAll
 
 /**
  * Retrieves the immediate children of the fiber.
diff --git a/packages/effect/src/internal/fiberRuntime.ts b/packages/effect/src/internal/fiberRuntime.ts
index ace835a650b..a649df30388 100644
--- a/packages/effect/src/internal/fiberRuntime.ts
+++ b/packages/effect/src/internal/fiberRuntime.ts
@@ -3288,8 +3288,14 @@ export const currentSupervisor: FiberRef.FiberRef<Supervisor.Supervisor<any>> =
 // circular with Fiber
 
 /* @internal */
-export const fiberAwaitAll = (fibers: Iterable<Fiber.Fiber<any, any>>): Effect.Effect<void> =>
-  core.asVoid(internalFiber._await(fiberAll(fibers)))
+export const fiberAwaitAll = <const T extends Iterable<Fiber.Fiber<any, any>>>(
+  fibers: T
+): Effect.Effect<
+  [T] extends [ReadonlyArray<infer U>]
+    ? number extends T["length"] ? Array<U extends Fiber.Fiber<infer A, infer E> ? Exit.Exit<A, E> : never>
+    : { -readonly [K in keyof T]: T[K] extends Fiber.Fiber<infer A, infer E> ? Exit.Exit<A, E> : never }
+    : Array<T extends Iterable<infer U> ? U extends Fiber.Fiber<infer A, infer E> ? Exit.Exit<A, E> : never : never>
+> => forEach(fibers, internalFiber._await) as any
 
 /** @internal */
 export const fiberAll = <A, E>(fibers: Iterable<Fiber.Fiber<A, E>>): Fiber.Fiber<Array<A>, E> => ({
diff --git a/packages/effect/test/Fiber.test.ts b/packages/effect/test/Fiber.test.ts
index e76e8c0c8bd..1bd813ccae8 100644
--- a/packages/effect/test/Fiber.test.ts
+++ b/packages/effect/test/Fiber.test.ts
@@ -210,7 +210,9 @@ describe("Fiber", () => {
   it.effect("awaitAll - stack safety", () =>
     Effect.gen(function*($) {
       const result = yield* $(Fiber.awaitAll(fibers))
-      assert.isUndefined(result)
+      assert.isArray(result)
+      assert(result.length === fibers.length)
+      result.forEach((_) => assert.isTrue(Exit.isSuccess(_) && _.value === undefined))
     }), 10000)
   it.effect("joinAll - stack safety", () =>
     Effect.gen(function*($) {