iterable-chunk.ts

function* chunk<T>(
  source: Iterable<T>,
  size = 1,
): Iterable<IterableIterator<T>> {
  const iterator = source[Symbol.iterator]();
  let done = false;

  function* take() {
    let count = 0;
    while (count++ < size) {
      const result = iterator.next();
      if (result.done) {
        done = true;
      } else {
        yield result.value;
      }
    }
  }

  while (!done) {
    const takeIterator = take();
    const first = takeIterator.next();
    if (!first.done) {
      const buffer: T[] = [];
      const chunkIterator = (function*() {
        if (buffer.length) {
          yield* buffer;
        } else {
          yield first.value;
          yield* takeIterator;
        }
      })();
      yield {
        next: chunkIterator.next,
        [Symbol.iterator]() {
          return chunkIterator;
        },
      };
      buffer.push(first.value, ...takeIterator);
    }
  }
}