LinqTypescript

enumerable.ts

export type Comparer<T> = (a: T, b: T) => number

export class Enumerable<T> implements Iterable<T> {
  protected readonly source: Iterable<T>

  constructor(source: Iterable<T> | (() => Generator<T, void, undefined>)) {
    this.source = typeof source == 'function' ? source() : source
  }
  static range(start: number, count: number, step: number = 1) {
    return new Enumerable(function* () {
      for (let i = start; i < count; i += step) {
        yield i
      }
    })
  }
  private pipe<R>(generator: (items: Enumerable<T>) => Generator<R>) {
    return new Enumerable(generator(this))
  }

  [Symbol.iterator](): Iterator<T> {
    return this.source[Symbol.iterator]()
  }

  take(count: number) {
    return this.pipe(function* (source) {
      let i = 0
      for (const item of source) {
        if (i++ >= count) {
          break
        }
        yield item
      }
    })
  }

  whereNotNull() {
    return this.where((x) => !!x) as Enumerable<NonNullable<T>>
  }
  where(predicate: (item: T, index: number) => boolean) {
    return this.pipe(function* (items: Enumerable<T>) {
      let index = 0
      for (const item of items) {
        if (predicate(item, index++)) {
          yield item
        }
      }
    })
  }

  any(predicate: (item: T) => boolean = (_) => true) {
    for (const item of this) {
      if (predicate(item)) {
        return true
      }
    }
    return false
  }

  all(predicate: (item: T) => boolean) {
    for (const item of this) {
      if (!predicate(item)) {
        return false
      }
    }
    return true
  }

  select<R>(fn: (item: T, index: number) => R) {
    return this.pipe(function* (source) {
      let i = 0
      for (const item of source) {
        yield fn(item, i++)
      }
    })
  }

  selectMany<R>(fn: (item: T) => Iterable<R>) {
    return this.pipe(function* (items) {
      for (const item of items) {
        yield* fn(item)
      }
    })
  }

  skip(count: number) {
    const source = this
    return new Enumerable<T>(function* () {
      let i = 0
      for (const item of source) {
        if (i++ >= count) {
          yield item
        }
      }
    })
  }

  groupBy<K>(keySelector: (item: T) => K) {
    const source = this.source
    return new Enumerable(function* () {
      const groups = new Map<K, T[]>()
      for (const item of source) {
        const key = keySelector(item)
        const group = groups.get(key)
        if (!group) {
          groups.set(key, [item])
        } else {
          group.push(item)
        }
      }
      for (let [key, items] of groups.entries()) {
        yield new Group<T, K>(key, items)
      }
    })
  }

  distinctBy<K>(keySelector: (item: T) => K) {
    const source = this.source
    return new Enumerable(function* () {
      const keys = new Set<K>([])
      for (const item of source) {
        const key = keySelector(item)
        if (!keys.has(key)) {
          keys.add(key)
          yield item
        }
      }
    })
  }

  orderBy<TKey>(keySelector: (item: T) => TKey, comparer?: Comparer<TKey>) {
    return OrderedEnumerable.from(this.source, keySelector, 'asc', comparer)
  }

  orderByDescending<TKey>(keySelector: (item: T) => TKey, comparer?: Comparer<TKey>) {
    return OrderedEnumerable.from(this.source, keySelector, 'desc', comparer)
  }

  concat(iterable: Iterable<T>) {
    const source = this.source
    return new Enumerable<T>(function* () {
      for (let sourceElement of source) {
        yield sourceElement
      }
      for (let sourceElement of iterable) {
        yield sourceElement
      }
    })
  }

  join<Outer, K, Result = [T, Outer]>(
    outer: Iterable<Outer>,
    innerSelector: (item: T, index: number) => K,
    outerSelector: (item: Outer, index: number) => K,
    mergeFn: (a: T, b: Outer) => Result = (a, b) => [a, b] as Result
  ): Enumerable<Result> {
    const inner = this.source
    return new Enumerable<Result>(function* () {
      const innerMap = new Map<K, T[]>()
      let index = 0
      for (const item of inner) {
        const key = innerSelector(item, index++)
        const items = innerMap.get(key)
        if (items) {
          items.push(item)
        } else {
          innerMap.set(key, [item])
        }
      }
      index = 0
      for (let item of outer) {
        const innerItems = innerMap.get(outerSelector(item, index++))
        if (innerItems) {
          for (const innerItem of innerItems) {
            yield mergeFn(innerItem, item)
          }
        }
      }
    })
  }

  each(action: (item: T, index: number) => void) {
    let index = 0
    for (const item of this) {
      action(item, index++)
    }
  }
  tap(action: (item: T, index: number) => void) {
    return this.pipe(function* (source) {
      let index = 0
      for (const item of source) {
        action(item, index++)
        yield item
      }
    })
  }

  // Collectors

  count(predicate: (item: T) => boolean = (_) => true) {
    let count = 0
    this.where(predicate).each(() => count++)
    return count
  }

  sumBy(fn: (item: T) => number) {
    let sum = 0
    this.each((q) => (sum += fn(q)))
    return sum
  }

  first(predicate: (item: T) => boolean = (_) => true) {
    const result = this.where(predicate)[Symbol.iterator]().next()
    return result.value as T | undefined
  }

  firstOr(predicate: (item: T) => boolean, defaultValue: T) {
    return this.first(predicate) ?? defaultValue
  }

  maxBy(fn: (item: T) => number) {
    let max: T | undefined = undefined
    let maxNumber = Number.MIN_VALUE
    for (const item of this) {
      const value = fn(item)
      if (value > maxNumber) {
        max = item
        maxNumber = value
      }
    }
    return max
  }

  minBy(fn: (item: T) => number) {
    let min: T | undefined = undefined
    let minNumber = Number.MAX_VALUE
    for (const item of this) {
      const value = fn(item)
      if (value < minNumber) {
        min = item
        minNumber = value
      }
    }
    return min
  }

  averageBy(fn: (item: T) => number) {
    let sum = 0
    let count = 0
    for (const item of this) {
      sum += fn(item)
      count++
    }
    return sum / count
  }

  toArray(): T[] {
    return Array.from(this)
  }

  toSet(): Set<T> {
    return new Set(this)
  }

  toMap<K, V = T>(
    keySelector: (item: T) => K,
    valueSelector: (item: T) => V = (x) => x as unknown as V
  ) {
    const map = new Map<K, V>()
    for (const item of this) map.set(keySelector(item), valueSelector(item))
    return map
  }

  toObject<TKey extends string | number | symbol, TValue>(
    keySelector: (item: T) => TKey,
    valueSelector: (item: T) => TValue
  ): Record<TKey, TValue> {
    const dict = {} as Record<TKey, TValue>
    for (const item of this) {
      dict[keySelector(item)] = valueSelector(item)
    }
    return dict
  }

  stringJoin(separator: string) {
    let result = ''
    let first = true
    for (const item of this) {
      if (first) first = false
      else result += separator
      result += item
    }
    return result
  }
}

export class Group<T, K> extends Enumerable<T> {
  constructor(
    public readonly key: K,
    source: Iterable<T>
  ) {
    super(source)
  }
}

class OrderedEnumerable<T> extends Enumerable<T> {
  constructor(
    source: Iterable<T>,
    private readonly comparers: {
      selector: (x: T) => any
      comparer: Comparer<any>
      direction: 'asc' | 'desc'
    }[]
  ) {
    super(source)
  }

  static from<T, K>(
    source: Iterable<T>,
    keySelector: (item: T) => K,
    direction: 'asc' | 'desc',
    comparer?: Comparer<K>
  ): OrderedEnumerable<T> {
    return new OrderedEnumerable<T>(source, [
      {
        selector: keySelector,
        comparer: comparer ?? this.defaultComparer,
        direction
      }
    ])
  }

  [Symbol.iterator](): Iterator<T> {
    const items = Array.from(this.source)
    items.sort((a, b) => {
      for (const comparer of this.comparers) {
        const result = comparer.comparer(comparer.selector(a), comparer.selector(b))
        if (result !== 0) {
          return comparer.direction === 'desc' ? -result : result
        }
      }
      return 0
    })
    return items[Symbol.iterator]()
  }

  thenBy<TKey>(keySelector: (item: T) => TKey, comparer?: Comparer<TKey>) {
    return new OrderedEnumerable<T>(this.source, [
      ...this.comparers,
      {
        selector: keySelector,
        direction: 'asc',
        comparer: comparer ?? OrderedEnumerable.defaultComparer
      }
    ])
  }

  thenByDescending<TKey>(keySelector: (item: T) => TKey, comparer?: Comparer<TKey>) {
    return new OrderedEnumerable(this.source, [
      ...this.comparers,
      {
        selector: keySelector,
        direction: 'desc',
        comparer: comparer ?? OrderedEnumerable.defaultComparer
      }
    ])
  }

  static defaultComparer<K>(a: K, b: K) {
    if (typeof a === 'string' && typeof b === 'string') {
      return a.localeCompare(b)
    }
    if (a instanceof Date && b instanceof Date) {
      return a.getTime() - b.getTime()
    }
    return a === b ? 0 : a < b ? -1 : 1
  }
}

export function enumerate<T>(source: Iterable<T>) {
  return new Enumerable(source)
}

export function enumerateObject<T extends {}>(source: T) {
  return new Enumerable(Object.entries(source)) as unknown as Enumerable<[keyof T, T[keyof T]]>
}

export function range(start: number, end: number, step = 1) {
  return new Enumerable(function* () {
    for (let i = start; i < end; i += step) {
      yield i
    }
  })
}