drop-in linq (generators)

linq.ts

export function* skip<T>(it: Iterable<T>, count: number): Iterable<T> {
  for (const $ of it) {
    if (count-- <= 0) {
      yield $;
    }
  }
}
export function* take<T>(it: Iterable<T>, count: number): Iterable<T> {
  for (const $ of it) {
    if (count-- <= 0) {
      break;
    }
    yield $;
  }
}
export function* takeWhile<T>(it: Iterable<T>, p: (x: T, i: number) => boolean): Iterable<T> {
  let i = 0;
  for (const $ of it) {
    if (!p($, i++)) {
      break;
    }
    yield $;
  }
}
export function* where<T>(it: Iterable<T>, p: (x: T, i: number) => boolean): Iterable<T> {
  let i = 0;
  for (const $ of it) if (p($, i++)) yield $;
}
export function* select<T, R>(it: Iterable<T>, s: (x: T, i: number) => R): Iterable<R> {
  let i = 0;
  for (const $ of it) yield s($, i++);
}
export function* selectMany<T, R>(it: Iterable<T>, s: (x: T) => Iterable<R>): Iterable<R> {
  for (const $ of it) {
    for (const $$ of s($)) {
      yield $$;
    }
  }
}
export function orderBy<T>(it: Iterable<T>, c: (a: T, b: T) => number): Iterable<T> {
  let sorted = toArray(it).sort(c);
  return sorted;
}
export function toArray<T>(it: Iterable<T>): T[] {
  let items = [], i = 0;
  for (const $ of it) items[i++] = $;
  return items;
}
export function count<T>(it: Iterable<T>): number {
  if (it["length"]) return it["length"];
  let sum = 0;
  for (const _ of it) { sum += 1; }
  return sum;
}

export function all<T>(it: Iterable<T>, p: (x: T) => boolean): boolean {
  let a = false;
  for (const $ of it) {
    if (!p($))
      return false;
    a = true;
  }
  return a;
}

export function any<T>(it: Iterable<T>, p?: (x: T) => boolean): boolean {
  if (p == undefined) {
    for (const $ of it) return true;
  } else {
    for (const $ of it) if (p($)) return true;
  }
  return false;
}

export function first<T>(it: Iterable<T>, p?: (x: T) => boolean): T | undefined {
  if (p == undefined) {
    for (const $ of it) return $;
  } else {
    for (const $ of it) if (p($)) return $;
  }
  return undefined;
}
export function last<T>(it: Iterable<T>, p?: (x: T) => boolean): T | undefined {
  let last: T;
  if (p == undefined) {
    for (const $ of it) last = $;
  } else {
    for (const $ of it) if (p($)) last = $;
  }
  return last;
}

export class Linq<T> implements Iterable<T> {
  constructor(private _it: Iterable<T>) { }
  [Symbol.iterator]() {
    return this._it[Symbol.iterator]();
  }
  skip(c: number): LinqIterable<T> {
    return new Linq(skip(this._it, c));
  }
  take(c: number): LinqIterable<T> {
    return new Linq(take(this._it, c));
  }
  takeWhile(p: (x: T, i: number) => boolean): LinqIterable<T> {
    return new Linq(takeWhile(this._it, p));
  }
  where(p: (x: T, i: number) => boolean): LinqIterable<T> {
    return new Linq(where(this._it, p));
  }
  select<R>(s: (x: T, i: number) => R): LinqIterable<R> {
    return new Linq(select(this._it, s));
  };
  selectMany<R>(s: (x: T) => LinqIterable<R>): LinqIterable<R> {
    return new Linq(selectMany(this._it, s));
  };
  orderBy(c: (a: T, b: T) => number): LinqIterable<T> {
    return new Linq(orderBy(this._it, c));
  };
  toArray(): T[] {
    return toArray(this._it);
  };
  count(): number {
    return count(this._it);
  }
  all(p: (x: T) => boolean): boolean {
    return all(this._it, p);
  }
  any(p?: (x: T) => boolean): boolean {
    return any(this._it, p);
  }
  first(p?: (x: T) => boolean): T | undefined {
    return first(this._it, p);
  }
  last(p?: (x: T) => boolean): T | undefined {
    return last(this._it, p);
  }
}

export type LinqIterable<T> = Iterable<T> & Linq<T>;

export function linq<T>(items: Iterable<T>): LinqIterable<T> {
  return new Linq(items);
}

declare global {
  interface Array<T> extends Linq<T> { }
}

Array.prototype.skip = function <T>(c: number): LinqIterable<T> {
  return new Linq(skip(<Array<T>>this, c));
};
Array.prototype.take = function <T>(c: number): LinqIterable<T> {
  return new Linq(take(<Array<T>>this, c));
};
Array.prototype.where = function <T>(p: (x: T, i: number) => boolean): LinqIterable<T> {
  return new Linq(where(<Array<T>>this, p));
};
Array.prototype.select = function <T, R>(s: (x: T, i: number) => R): LinqIterable<R> {
  return new Linq(select(<Array<T>>this, s));
};
Array.prototype.selectMany = function <T, R>(s: (x: T) => LinqIterable<R>): LinqIterable<R> {
  return new Linq(selectMany(<Array<T>>this, s));
};
Array.prototype.orderBy = function <T>(c: (a: T, b: T) => number): LinqIterable<T> {
  return new Linq(orderBy(<Array<T>>this, c));
};
Array.prototype.toArray = function <T>(): T[] {
  return toArray(<Array<T>>this);
};
Array.prototype.count = function <T>(): number {
  return count(<Array<T>>this);
};
Array.prototype.all = function <T>(p: (x: T) => boolean): boolean {
  return all(<Array<T>>this, p);
};
Array.prototype.any = function <T>(p: (x: T) => boolean): boolean {
  return any(<Array<T>>this, p);
};
Array.prototype.last = function <T>(p?: (x: T) => boolean): T {
  return first(<Array<T>>this, p);
};
Array.prototype.first = function <T>(p?: (x: T) => boolean): T {
  return last(<Array<T>>this, p);
};

import { linq } from './src/linq';

let items = {
  *[Symbol.iterator]() {
    yield 1;
    yield 2;
    yield 3;
    yield 4;
  }
};

let iterable = linq(items)
  .select((x, i) => { return { x: x, i: i }; })
  .orderBy((a, b) => b.i - a.i)
  .select((x) => x.x);

for (const item of iterable) {
  console.log(item);
}