multiplex.ts

/**
 * Multiplexes multiple async iterators into a single async iterable iterator.
 * @param args An array of async iterators to be multiplexed.
 * @returns An async iterable iterator that yields values immediately from iterators.
 */
export async function* multiplex<T, N>(
  args: readonly AsyncIterator<T, void, N>[],
): AsyncIterableIterator<T> {
  const promises = new Map<AsyncIterator<T, unknown, N>, ReturnType<typeof step<T, N>>>();
  args.forEach((iterator) => promises.set(iterator, step(iterator)));

  while (promises.size > 0) {
    const { iterator, result } = await Promise.race([...promises.values()]);
    if (result.done) {
      promises.delete(iterator);
      continue;
    }

    const ret = yield result.value;
    promises.set(iterator, step(iterator, ret));
  }
}

async function step<T, N>(iterator: AsyncIterator<T, unknown, N>, arg?: N) {
  const result = await (arg === undefined ? iterator.next() : iterator.next(arg));
  return { result, iterator };
}