feat(async): add async/tee

async/README.md

@@ -83,3 +83,31 @@ for await (const value of results) {
   // ...
 }
 ```
+
+## tee
+
+Branches the given async iterable into the n branches.
+
+```typescript
+import { tee } from "https://deno.land/std/async/tee.ts";
+
+const gen = async function* gen() {
+  yield 1;
+  yield 2;
+  yield 3;
+};
+
+const [branch1, branch2] = tee(gen());
+
+(async () => {
+  for await (const n of branch1) {
+    console.log(n); // => 1, 2, 3
+  }
+})();
+
+(async () => {
+  for await (const n of branch2) {
+    console.log(n); // => 1, 2, 3
+  }
+})();
+```

async/mod.ts

@@ -3,3 +3,4 @@ export * from "./deferred.ts";
 export * from "./delay.ts";
 export * from "./mux_async_iterator.ts";
 export * from "./pool.ts";
+export * from "./tee.ts";

async/tee.ts

@@ -0,0 +1,102 @@
+// Copyright 2018-2021 the Deno authors. All rights reserved. MIT license.
+
+// Utility for representing n-tuple
+type Tuple<T, N extends number> = N extends N
+  ? number extends N ? T[] : TupleOf<T, N, []>
+  : never;
+type TupleOf<T, N extends number, R extends unknown[]> = R["length"] extends N
+  ? R
+  : TupleOf<T, N, [T, ...R]>;
+
+const noop = () => {};
+
+class AsyncIterableClone<T> implements AsyncIterable<T> {
+  currentPromise: Promise<IteratorResult<T>>;
+  resolveCurrent: (x: Promise<IteratorResult<T>>) => void = noop;
+  consumed: Promise<void>;
+  consume: () => void = noop;
+
+  constructor() {
+    this.currentPromise = new Promise<IteratorResult<T>>((resolve) => {
+      this.resolveCurrent = resolve;
+    });
+    this.consumed = new Promise<void>((resolve) => {
+      this.consume = resolve;
+    });
+  }
+
+  reset() {
+    this.currentPromise = new Promise<IteratorResult<T>>((resolve) => {
+      this.resolveCurrent = resolve;
+    });
+    this.consumed = new Promise<void>((resolve) => {
+      this.consume = resolve;
+    });
+  }
+
+  async next(): Promise<IteratorResult<T>> {
+    const res = await this.currentPromise;
+    this.consume();
+    this.reset();
+    return res;
+  }
+
+  async push(res: Promise<IteratorResult<T>>): Promise<void> {
+    this.resolveCurrent(res);
+    // Wait until current promise is consumed and next item is requested.
+    await this.consumed;
+  }
+
+  [Symbol.asyncIterator](): AsyncIterator<T> {
+    return this;
+  }
+}
+
+/**
+ * Branches the given async iterable into the n branches.
+ *
+ * Example:
+ *
+ *     const gen = async function* gen() {
+ *       yield 1;
+ *       yield 2;
+ *       yield 3;
+ *     }
+ *
+ *     const [branch1, branch2] = tee(gen());
+ *
+ *     (async () => {
+ *       for await (const n of branch1) {
+ *         console.log(n); // => 1, 2, 3
+ *       }
+ *     })();
+ *
+ *     (async () => {
+ *       for await (const n of branch2) {
+ *         console.log(n); // => 1, 2, 3
+ *       }
+ *     })();
+ */
+export function tee<T, N extends number = 2>(
+  src: AsyncIterable<T>,
+  n: N = 2 as N,
+): Tuple<AsyncIterable<T>, N> {
+  const clones: Tuple<AsyncIterableClone<T>, N> = Array.from({ length: n }).map(
+    () => new AsyncIterableClone(),
+    // deno-lint-ignore no-explicit-any
+  ) as any;
+  (async () => {
+    const iter = src[Symbol.asyncIterator]();
+    await Promise.resolve();
+    while (true) {
+      const res = iter.next();
+      await Promise.all(clones.map((c) => c.push(res)));
+      if ((await res).done) {
+        break;
+      }
+    }
+  })().catch((e) => {
+    console.error(e);
+  });
+  return clones;
+}

async/tee_test.ts

@@ -0,0 +1,62 @@
+// Copyright 2018-2021 the Deno authors. All rights reserved. MIT license.
+import { tee } from "./tee.ts";
+import { assertEquals } from "https://deno.land/[email protected]/testing/asserts.ts";
+
+/** An example async generator */
+const gen = async function* iter() {
+  yield 1;
+  yield 2;
+  yield 3;
+};
+
+/** Testing utility for accumulating the values in async iterable. */
+async function accumulate<T>(src: AsyncIterable<T>): Promise<T[]> {
+  const res: T[] = [];
+  for await (const item of src) {
+    res.push(item);
+  }
+  return res;
+}
+
+Deno.test("async/tee - 2 branches", async () => {
+  const iter = gen();
+  const [res0, res1] = tee(iter).map(accumulate);
+  assertEquals(
+    await Promise.all([res0, res1]),
+    [
+      [1, 2, 3],
+      [1, 2, 3],
+    ],
+  );
+});
+
+Deno.test("async/tee - 3 branches - immediate consumption", async () => {
+  const iter = gen();
+  const [res0, res1, res2] = tee(iter, 3).map(accumulate);
+  assertEquals(
+    await Promise.all([res0, res1, res2]),
+    [
+      [1, 2, 3],
+      [1, 2, 3],
+      [1, 2, 3],
+    ],
+  );
+});
+
+Deno.test("async/tee - 3 branches - delayed consumption", async () => {
+  const iter = gen();
+  const iters = tee(iter, 3);
+
+  await new Promise<void>((resolve) => {
+    setTimeout(() => resolve(), 20);
+  });
+
+  assertEquals(
+    await Promise.all(iters.map(accumulate)),
+    [
+      [1, 2, 3],
+      [1, 2, 3],
+      [1, 2, 3],
+    ],
+  );
+});