chain.ts

export async function* chain<T, U>(g: AsyncGenerator<T, void>, f: (x: T) => AsyncGenerator<U, void>): AsyncGenerator<U, void> {
  type R<I> = IteratorResult<I extends number ? U : T, void>;
  type IR<I> = [I, R<I>];
  const k = <I>(i: I) => (res: R<I>): IR<I> => [i, res];
  const running = new Set<AsyncGenerator<unknown, void>>([g]);
  let p = g.next().then(k(null));
  const ps: Promise<IR<number>>[] = [];
  const gs: AsyncGenerator<U, void>[] = [];
  const queue = new Set<Promise<IR<null> | IR<number>>>([p]);
  try {
    while (queue.size > 0) {
      const [i, res] = await Promise.race(queue);
      queue.delete(i === null ? p : ps[i]!);
      if (res.done) {
        running.delete(i === null ? g : gs[i]!);
      } else if (i === null) {
        const j = gs.length;
        const gsj = f(res.value);
        gs.push(gsj);
        running.add(gsj);
        const psj = gsj.next().then(k(j));
        ps.push(psj);
        queue.add(psj);
        p = g.next().then(k(null));
        queue.add(p);
      } else {
        const psi = ps[i] = gs[i]!.next().then(k(i));
        queue.add(psi);
        yield res.value;
      }
    }
  } finally {
    await Promise.all(Array.from(running, g => g.return()));
  }
}