Enumerable.ts

type Primitive = string | number | boolean;

module Iters {
  export function* where<T>(
    generator: () => IterableIterator<T>, 
    predicate: (value: T, index: number) => boolean
  ) {
    let index = 0;
    for (let item of generator()) {
      if (predicate(item, index)) {
        yield item;
      }
      index++;
    }
  }

  export function* select<TIn, TOut>(
    generator: () => IterableIterator<TIn>,
    selector: (value: TIn, index: number) => TOut
  ) {
    let index = 0;
    for (let item of generator()) {
      yield selector(item, index);
      index++;
    }
  }

  export function* selectMany<TSource, TResult>(
    generator: () => IterableIterator<TSource>, 
    selector: (value: TSource, index: number) => IterableIterator<TResult>
  ): IterableIterator<TResult> {
    let index = 0;
    for (let item of generator()) {
      yield* selector(item, index);
      index++;
    }
  }

  export function* join<TLeft, TRight, TKey extends Primitive, TResult>(
    leftGenerator: () => IterableIterator<TLeft>,
    rightGenerator: () => IterableIterator<TRight>,
    leftKeySelector: (left: TLeft) => TKey,
    rightKeySelector: (right: TRight) => TKey,
    resultSelector: (left: TLeft, right: TRight) => TResult
  ): IterableIterator<TResult> {
    let leftKey: TKey = null;
    let rightKey: TKey = null;
    for (let left of leftGenerator()) {
      leftKey = leftKeySelector(left);
      for (let right of rightGenerator()) {
        rightKey = rightKeySelector(right); 
        if (leftKey === rightKey) {
          yield resultSelector(left, right);
        }
      }
    }
  }

  export function* values<T>(array: T[]) {
    for (let x of array) {
      yield x;
    }
  }

  export function* range(start: number, count: number) {
    for (let x = 0; x < count; x++) {
      yield (start + x);
    }
  }

  export function* skip<T>(generator: () => IterableIterator<T>, count: number) {
    let index = 0;
    for (let item of generator()) {
      if (index >= count) {
        yield item;
      }
      index++;
    }
  }

  export function* concat<T>(
    leftGenerator: () => IterableIterator<T>,
    rightGenerator: () => IterableIterator<T>
  ) {
    for (let left of leftGenerator()) {
      yield left;
    }
    for (let right of rightGenerator()) {
      yield right;
    }
  }

  export function* lookup<T, TKey extends Primitive, TElement, TResult>(
    generator: () => IterableIterator<T>,
    keySelector: (value: T) => TKey,
    elementSelector: (value: T) => TElement,
    resultSelector: (key: TKey, elements: TElement[]) => TResult
  ): IterableIterator<[TKey, TElement[]]> {
    let map = new Map<TKey, TElement[]>();
    for (let x of generator()) {
      let key = keySelector(x);
      let group = map.get(key);
      if (!group) {
        group = [];
        map.set(key, group);
      }
      group.push(elementSelector(x));
    }
    yield* map.entries();
  }
}

class Enumerable<T> {
  public readonly generator: () => IterableIterator<T> = null;

  protected constructor(generator: () => IterableIterator<T>) {
    this.generator = generator;
  }

  public where(predicate: (value: T, index: number) => boolean): Enumerable<T> {
    let generator = () => Iters.where(this.generator, predicate);
    return Enumerable.fromGenerator(generator);
  }

  public select<TResult>(selector: (value: T, index: number) => TResult): Enumerable<TResult> {
    let generator = () => Iters.select(this.generator, selector);
    return Enumerable.fromGenerator(generator);
  }

  public selectMany<TResult>(selector: (value: T, index: number) => Enumerable<TResult>): Enumerable<TResult> {
    let generator = () => Iters.selectMany(this.generator, (value, index) => selector(value, index).generator());
    return Enumerable.fromGenerator(generator);
  }

  public skip(count: number): Enumerable<T> {
    let generator = () => Iters.skip(this.generator, count);
    return Enumerable.fromGenerator(generator);
  }

  public groupBy<TKey extends Primitive, TElement>(
    keySelector: (value: T) => TKey,
    elementSelector: (value: T) => TElement,
  ): Enumerable<Grouping<TKey, TElement>>;

  public groupBy<TKey extends Primitive, TElement, TResult>(
    keySelector: (value: T) => TKey,
    elementSelector: (value: T) => TElement,
    resultSelector: (key: TKey, elements: TElement[]) => TResult
  ): Enumerable<TResult>;

  public groupBy<TKey extends Primitive, TElement, TResult>(
    keySelector: (value: T) => TKey,
    elementSelector: (value: T) => TElement,
    resultSelector?: (key: TKey, elements: TElement[]) => TResult
  ): Enumerable<TResult> | Enumerable<Grouping<TKey, TElement>> {
    let lookup = () => Iters.lookup(this.generator, keySelector, elementSelector, resultSelector);
    if (resultSelector) {
      let generator = () => Iters.select(lookup, ([key, values]) => resultSelector(key, values));
      return Enumerable.fromGenerator(generator);
    } else {
      let generator = () => Iters.select(lookup, ([key, values]) => Grouping.fromKeyedArray(key, values));
      return Enumerable.fromGenerator(generator);
    }
  }

  public join<TRight, TKey extends Primitive, TResult>(
    right: Enumerable<TRight>,
    leftKeySelector: (value: T) => TKey,
    rightKeySelector: (value: TRight) => TKey,
    resultSelector: (left: T, right: TRight) => TResult
  ): Enumerable<TResult> {
    let generator = () => Iters.join(this.generator, right.generator, leftKeySelector, rightKeySelector, resultSelector);
    return Enumerable.fromGenerator(generator);
  }

  public concat(right: Enumerable<T>): Enumerable<T> {
    let generator = () => Iters.concat(this.generator, right.generator);
    return Enumerable.fromGenerator(generator);
  }

  public contains(value: T, comparer?: (left: T, right: T) => boolean): boolean {
    comparer = comparer ? comparer : (left, right) => left === right;
    for (let x of this.generator()) {
      if (comparer(x, value)) {
        return true;
      }
    }
    return false;
  }

  public toArray() {
    return Array.from(this.generator());
  }

  public static range(start: number, count: number): Enumerable<number> {
    let generator = () => Iters.range(start, count);
    return new Enumerable(generator);
  }

  public static empty<T>(): Enumerable<T> {
    let generator = () => Iters.values([]);
    return new Enumerable(generator);
  }

  public static fromGenerator<T>(generator: () => IterableIterator<T>): Enumerable<T> {
    return new Enumerable(generator);
  }

  public static fromArray<T>(array: T[]): Enumerable<T> {
    let generator = () => Iters.values(array);
    return Enumerable.fromGenerator(generator);
  }
}

class Grouping<TKey extends Primitive, TElement> extends Enumerable<TElement> {
  public readonly key: TKey;

  protected constructor(key: TKey, generator: () => IterableIterator<TElement>) {
    super(generator);
    this.key = key;
  }

  public static fromKeyedArray<TKey extends Primitive, TElement>(key: TKey, values: TElement[]) {
    return new Grouping(key, () => Iters.values(values));
  }
}