Extentions for ECMA262 6th iterator

ExIterable.ts

/**
 *  This provides extensions for ECMA262 2015's iterator.
 *
 *  This adds `map()`, `forEach()`, `filter()`, `flatMap()`, or others
 *  to `Iterable<T>`. This enables features looks like "lazy evaluation".
 *  The design refers RxJS v5's one.
 *
 *  See example:
 *  ```
 *
 *      const iter = ExIterable.create([1, 2, 3]);
 *      // Don't evaluate the result.
 *      const mapped = iter.map( (v) => v + 1 );
 *
 *      // At this, we start to consume the data source.
 *      mapped.forEach( (v) => console.log(v) );
 *
 *      // At this, we start to consume the data source _newly_.
 *      for (const i of mapped) { console.log(v); }
 *   ```
 */
export class ExIterable<T> implements Iterable<T> {

    protected _source: Iterable<any> | void; // cheat to drop type param `R`.
    protected _operator: Operator<any, T> | void; // cheat to drop type param `R`.

    protected constructor(source?: Iterable<T>) {
        this._source = source;
        this._operator = undefined;
    }

    static create<T>(source: Iterable<T>): ExIterable<T> {
        return new ExIterable<T>(source);
    }

    lift<U>(operator: Operator<T, U>): ExIterable<U> {
        const iterable = new ExIterable<U>();
        iterable._source = this;
        iterable._operator = operator;
        return iterable;
    }

    forEach(fn: (v: T, index: number) => void): void {
        const iter: Iterator<T> = this[Symbol.iterator]();
        let index = 0;
        let next: IteratorResult<T> = iter.next();
        while (!next.done) {
            fn(next.value, index++);
            next = iter.next();
        }
    }

    map<U>(selector: (this: undefined, value: T, index: number) => U): ExIterable<U> {
        const op = new MapOperator<T, U>(selector);
        const lifted = this.lift<U>(op);
        return lifted;
    }

    flatMap<U>(selector: (this: undefined, value: T, index: number) => Iterable<U>): ExIterable<U> {
        const op = new FlatMapOperator<T, U>(selector);
        const lifted = this.lift<U>(op);
        return lifted;
    }

    filter(filter: (this: undefined, value: T, index: number) => boolean): ExIterable<T> {
        const op = new FilterOperator<T>(filter);
        const lifted = this.lift<T>(op);
        return lifted;
    }

    do(action: (this: undefined, value: T, index: number) => void): ExIterable<T> {
        const op = new DoOperator<T>(action);
        const lifted = this.lift<T>(op);
        return lifted;
    }

    cache(): ExIterable<T> {
        const op = new CacheOperator<T>();
        return this.lift(op);
    }

    [Symbol.iterator](): Iterator<T> {
        // XXX:
        // There are still overhead to create "source" iterator
        // even if we call `CacheOperator`. To avoid the overhead,
        // we would need to change `Operator` interface.
        const source = this._source[Symbol.iterator]();
        if (this._operator === undefined) {
            return this._source[Symbol.iterator]();
        }
        const iter = this._operator.call(source);
        return iter;
    }
}

type MapFn<T, U> = (v: T, index: number) => U;

class MapOperator<S, T> implements Operator<S, T> {
    private _selector: MapFn<S, T>;

    constructor(selector: MapFn<S, T>) {
        this._selector = selector;
    }

    call(source: Iterator<S>): Iterator<T> {
        const iter = generateMapIterator<S, T>(source, this._selector);
        return iter;
    }
}

function* generateMapIterator<S, T>(source: Iterator<S>, selector: MapFn<S, T>): IterableIterator<T> {
    let index: number = 0;
    while (true) {
        const original: IteratorResult<S> = source.next();
        if (original.done) {
            return;
        }

        const result: T = selector(original.value, index++);
        yield result;
    }
}

interface Operator<S, T> {
    call(source: Iterator<S>): Iterator<T>;
}

type FilterFn<T> = (value: T, index: number) => boolean;

class FilterOperator<T> implements Operator<T, T> {

    private _filter: FilterFn<T>;

    constructor(filter: FilterFn<T>) {
        this._filter = filter;
    }

    call(source: Iterator<T>): Iterator<T> {
        const iter = generateFilterIterator<T>(source, this._filter);
        return iter;
    }
}

function* generateFilterIterator<T>(source: Iterator<T>, filter: FilterFn<T>): IterableIterator<T> {
    let index: number = 0;

    while (true) {
        let next: IteratorResult<T> = source.next();
        while (!next.done) {
            const ok: boolean = filter(next.value, index++);
            if (ok) {
                yield next.value;
            }

            next = source.next();
        }

        return;
    }
}

type FlatMapFn<T, U> = (v: T, index: number) => Iterable<U>;

class FlatMapOperator<S, T> implements Operator<S, T> {
    private _selector: FlatMapFn<S, T>;

    constructor(selector: FlatMapFn<S, T>) {
        this._selector = selector;
    }

    call(source: Iterator<S>): Iterator<T> {
        const iter = generateFlatMap<S, T>(source, this._selector);
        return iter;
    }
}

function* generateFlatMap<S, T>(source: Iterator<S>, selector: FlatMapFn<S, T>): IterableIterator<T> {
    let inner: Iterator<T> | void = undefined;
    let index: number = 0;

    while (true) {
        if (inner === undefined) {
            const outer: IteratorResult<S> = source.next();
            if (outer.done) {
                return;
            }

            const result: Iterable<T> = selector(outer.value, index++);
            inner = result[Symbol.iterator]();
            if (!inner) {
                throw new Error('selector cannot return a valid iterable.');
            }
        }
        else {
            const result: IteratorResult<T> = inner.next();
            if (result.done) {
                inner = undefined;
                continue;
            }
            else {
                yield result.value;
            }
        }
    }
}

type DoFn<T> = (value: T, index: number) => void;

class DoOperator<T> implements Operator<T, T> {

    private _action: DoFn<T>;

    constructor(action: DoFn<T>) {
        this._action = action;
    }

    call(source: Iterator<T>): Iterator<T> {
        const iter = generateDoIterator<T>(source, this._action);
        return iter;
    }
}

function* generateDoIterator<T>(source: Iterator<T>, action: DoFn<T>): IterableIterator<T> {
    let index: number = 0;

    while (true) {
        const next: IteratorResult<T> = source.next();
        if (next.done) {
            return;
        }
        else {
            const result: T = next.value;
            action(result, index++);
            yield result;
        }
    }
}

class CacheOperator<T> implements Operator<T, T> {
    private _cacheIterator: Iterator<T> | void;
    private _cacheResult: Array<T>;

    constructor() {
        this._cacheIterator = undefined;
        this._cacheResult = [];
    }

    call(source: Iterator<T>): Iterator<T> {
        if (this._cacheIterator === undefined) {
            this._cacheIterator = source;
        }
        const iter = generateCacheIterator<T>(this._cacheIterator, this._cacheResult);
        return iter;
    }
}

// XXX:
// This cache logic is just a concept. There may be a some potential leak
function* generateCacheIterator<T>(source: Iterator<T>, cache: Array<T>): IterableIterator<T> {
    let index: number = 0;
    while (true) {
        const current: number = index;
        ++index;

        // Even if the slot is filled with `undefined`,
        // it includes as the array's length after assignment a value.
        if (current <= (cache.length - 1)) {
            const value: T = cache[current];
            yield value;
        }
        else {
            const { done, value }: IteratorResult<T> = source.next();
            if (done) {
                return;
            }
            else {
                cache[current] = value;
                yield value;
            }
        }
    }
}

sample.js

const list = [
    [1, 2],
    [3, 4],
    [5, 6],
];

const iter = ExIterable.create(list)
    .flatMap( (v) => v )
    .filter( (v) => (v % 2) === 0 )
    .map( (v) => v * v )
    .do( (v) => console.log('do:' + v) ); // <-- don't evaluate in here

iter.forEach( (v, i) => console.log(i, v) );
// do:4
// 0 4
// do:16
// 1 16
// do:36
// 2 36

test.ts

import * as assert from 'assert';

import {ExIterable} from './ExIterable';

class HelperIterable<T> implements Iterable<T> {

    private _source: Iterable<T>;
    private _onNext: (v: IteratorResult<T>) => void;
    private _onAfterFinish: (() => void) | void;

    constructor(src: Iterable<T>, onNext: (v: IteratorResult<T>) => void, onAfterFinish: (() => void) | void = undefined) {
        this._source = src;
        this._onNext = onNext;
        this._onAfterFinish = onAfterFinish;
    }

    [Symbol.iterator](): Iterator<T> {
        const src = this._source[Symbol.iterator]();
        const iter = new HelperIterator(src, this._onNext, this._onAfterFinish);
        return iter;
    }
}

class HelperIterator<T> implements Iterator<T> {

    private _source: Iterator<T>;
    private _onNext: (v: IteratorResult<T>) => void;
    private _onAfterFinish: (() => void) | void;

    constructor(src: Iterator<T>, onNext: (v: IteratorResult<T>) => void, onAfterFinish: (() => void) | void = undefined)  {
        this._source = src;
        this._onNext = onNext;
        this._onAfterFinish = onAfterFinish;
    }

    next(): IteratorResult<T> {
        const result: IteratorResult<T> = this._source.next();
        this._onNext(result);

        if (result.done && (this._onAfterFinish !== undefined)) {
            this._onAfterFinish();
        }

        return result;
    }
}

describe('ExIterable', function () {
    describe('create()', function () {
        let isCalledNext = false;

        before(function () {
            const src = new HelperIterable([1, 2, 3], () => {
                isCalledNext = true;
            });
            const iter = ExIterable.create(src);
            iter;
        });

        it('don\'t evaluate immidiately on creating an instance', () => {
            assert.strictEqual(isCalledNext, false);
        });
    });

    describe('forEach()', function () {
        describe('simple iteration', function () {
            const src = [1, 2, 3];
            const result: Array<number> = [];

            before(function () {
                const iter = ExIterable.create(src);
                iter.forEach((v) => {
                    result.push(v);
                });
            });

            it('iterate all values in source', () => {
                assert.deepStrictEqual(result, src);
            });
        });

        describe('iterate from zero per iteration', function () {
            class Helper {
                private _i: number;
                constructor(seed: number) {
                    this._i = seed;
                }
                value() {
                    return this._i;
                }
                increment() {
                    this._i = this._i + 1;
                }
            }

            const firstSeq: Array<number> = [];
            const secondSeq: Array<number> = [];

            before(function(){
                const src = [0, 1, 2].map((i) => new Helper(i));
                const iter = ExIterable.create(src);
                iter.forEach((v) => {
                    v.increment();
                    firstSeq.push( v.value() );
                });
                iter.forEach((v) => {
                    v.increment();
                    secondSeq.push( v.value() );
                });
            });

            it('first iteration', function () {
                assert.deepStrictEqual(firstSeq, [1, 2, 3]);
            });

            it('second iteration', function () {
                assert.deepStrictEqual(secondSeq, [2, 3, 4]);
            });
        });
    });

    describe('for-of statement', function () {
        describe('simple iteration', function () {
            const src = [1, 2, 3];
            const result: Array<number> = [];

            before(function () {
                const iter = ExIterable.create(src);

                for (const v of iter) {
                    result.push(v);
                }
            });

            it('iterate all values in source', () => {
                assert.deepStrictEqual(result, src);
            });
        });

        describe('iterate from zero per iteration', function () {
            class Helper {
                private _i: number;
                constructor(seed: number) {
                    this._i = seed;
                }
                value() {
                    return this._i;
                }
                increment() {
                    this._i = this._i + 1;
                }
            }

            const firstSeq: Array<number> = [];
            const secondSeq: Array<number> = [];

            before(function(){
                const src = [0, 1, 2].map((i) => new Helper(i));
                const iter = ExIterable.create(src);
                for (const v of iter) {
                    v.increment();
                    firstSeq.push( v.value() );
                }

                for (const v of iter) {
                    v.increment();
                    secondSeq.push( v.value() );
                }
            });

            it('first iteration', function () {
                assert.deepStrictEqual(firstSeq, [1, 2, 3]);
            });

            it('second iteration', function () {
                assert.deepStrictEqual(secondSeq, [2, 3, 4]);
            });
        });
    });

    describe('map()', function () {
        const resultSeq: Array<number> = [];
        const indexSeq: Array<number> = [];

        before(function () {
            const iter = ExIterable.create([0, 1, 2])
                .map((v, i) => {
                    indexSeq.push(i);
                    return v + 1;
                });
            iter.forEach((v) => {
                resultSeq.push(v);
            });
        });

        it('expected result sequence', function () {
            assert.deepStrictEqual(resultSeq, [1, 2, 3]);
        });

        it('expected index sequence', function () {
            assert.deepStrictEqual(indexSeq, [0, 1, 2]);
        });
    });

    describe('filter()', function () {
        const resultSeq: Array<number> = [];
        const indexSeq: Array<number> = [];

        before(function () {
            const iter = ExIterable.create([0, 1, 2, 3, 4])
                .filter((v, i) => {
                    indexSeq.push(i);
                    return (v % 2 === 0);
                });
            iter.forEach((v) => {
                resultSeq.push(v);
            });
        });

        it('expected result sequence', function () {
            assert.deepStrictEqual(resultSeq, [0, 2, 4]);
        });

        it('expected index sequence', function () {
            assert.deepStrictEqual(indexSeq, [0, 1, 2, 3, 4]);
        });
    });

    describe('do()', function () {
        const resultSeq: Array<number> = [];
        const indexSeq: Array<number> = [];

        before(function () {
            const iter = ExIterable.create([0, 1, 2])
                .do((v, i) => {
                    indexSeq.push(i);
                    return String(v);
                });
            iter.forEach((v) => {
                resultSeq.push(v);
            });
        });

        it('expected result sequence', function () {
            assert.deepStrictEqual(resultSeq, [0, 1, 2]);
        });

        it('expected index sequence', function () {
            assert.deepStrictEqual(indexSeq, [0, 1, 2]);
        });
    });

    describe('cache()', function () {
        describe('simple case', function () {
            const resultSeq1: Array<number> = [];
            const resultSeq2: Array<number> = [];

            before(function () {
                const src = ExIterable.create([0, 1, 2])
                    .map(() => Math.random());
                const iter = ExIterable.create(src).cache();

                iter.forEach((v) => {
                    resultSeq1.push(v);
                });

                iter.forEach((v) => {
                    resultSeq2.push(v);
                });
            });

            it('expected result 1 & 2 sequence', function () {
                assert.deepStrictEqual(resultSeq1, resultSeq2);
            });
        });

        describe('call `next()` from some iterator by turns', function () {
            function getIterator<T>(s: Iterable<T>): Iterator<T> {
                return s[Symbol.iterator]();
            }
            function pushToArray<T>(i: Iterator<T>, target: Array<IteratorResult<T>>): void {
                const result = i.next();
                target.push(result);
            }

            let iter1: Iterator<number>;
            let iter2: Iterator<number>;
            let iter3: Iterator<number>;

            const seq1: Array<IteratorResult<number>> = [];
            const seq2: Array<IteratorResult<number>> = [];
            const seq3: Array<IteratorResult<number>> = [];

            before(function () {
                const src = ExIterable.create([0, 1, 2, 3, 4])
                    .map((v) => v + Math.random());
                const iterable = ExIterable.create(src).cache();
                iter1 = getIterator(iterable);
                iter2 = getIterator(iterable);
                iter3 = getIterator(iterable);

                iter1.next();
                iter3.next();
                iter2.next();

                pushToArray(iter1, seq1);
                pushToArray(iter2, seq2);
                pushToArray(iter3, seq3);

                pushToArray(iter3, seq3);
                pushToArray(iter2, seq2);
                pushToArray(iter1, seq1);

                pushToArray(iter1, seq1);
                pushToArray(iter3, seq3);
                pushToArray(iter2, seq2);

                pushToArray(iter2, seq2);
                pushToArray(iter1, seq1);
                pushToArray(iter3, seq3);

                pushToArray(iter2, seq2);
                pushToArray(iter3, seq3);
                pushToArray(iter1, seq1);

                pushToArray(iter3, seq3);
                pushToArray(iter1, seq1);
                pushToArray(iter2, seq2);
            });

            it('iter1 & iter2 are different', function () {
                assert.notStrictEqual(iter1, iter2);
            });

            it('iter1 & iter3 are different', function () {
                assert.notStrictEqual(iter1, iter3);
            });

            it('iter2 & iter3 are different', function () {
                assert.notStrictEqual(iter2, iter3);
            });

            it('seq1 & seq2 are same', function() {
                assert.deepStrictEqual(seq1, seq2);
            });

            it('seq1 & seq3 are same', function() {
                assert.deepStrictEqual(seq1, seq3);
            });

            it('seq2 & seq3 are same', function() {
                assert.deepStrictEqual(seq2, seq3);
            });
        });
    });
});