tree.ts

export type Tree<T> = Leaf<T> | TreeNodeArray<T> | TreeNodeObject<T>;

type GenericObject<T> = { [key: string]: T };

type ChildrenArray<T> = Tree<T>[];
type ChildrenObject<T> = GenericObject<Tree<T>>

function mapObj<T, U>(fn: ((value: T, key: string) => U), obj: GenericObject<T>): GenericObject<U> {
    let ret = {};

    for (let key in obj) {
        if (obj.hasOwnProperty(key)) {
            ret[key] = fn(obj[key], key);
        }
    }

    return ret;
}

function values<T>(obj: GenericObject<T>): T[] {
    let ret = [];

    for (let key in obj) {
        if (obj.hasOwnProperty(key)) {
            ret.push(obj[key]);
        }
    }

    return ret;
}

export class Leaf<T> {
    private value: T;

    /* tslint:disable: no-any */
    public static of(value: any): Leaf<any> {
        return new Leaf(value);
    }
    /* tslint:enable: no-any */

    constructor(value: T) {
        this.value = value;
    }

    public of(value: T): Leaf<T> {
        return new Leaf(value);
    }

    public extract(): T {
        return this.value;
    }

    public map<U>(fn: (value: T) => U): Tree<U> {
        return Leaf.of(fn(this.value));
    }

    public extend<U>(fn: (tree: Tree<T>) => U): Leaf<U> {
        return Leaf.of(fn(this));
    }

    public reduce<U>(fn: (previousValue: U, currentValue: T) => U, acc: U): U {
        return fn(acc, this.value);
    }
}

export class TreeNodeArray<T> {
    private value: T;
    private children: ChildrenArray<T>;

    /* tslint:disable: no-any */
    public static of(value: any, children: Tree<any>[]): TreeNodeArray<any> {
        return new TreeNodeArray(value, children);
    }
    /* tslint:enable: no-any */

    constructor(value: T, children: ChildrenArray<T>) {
        this.value = value;
        this.children = children;
    }

    public of(value: T, children: ChildrenArray<T>): TreeNodeArray<T> {
        return new TreeNodeArray(value, children);
    }

    /* tslint:disable: no-any */
    public extract(): any[] {
        return this.children.map(child => child.extract());
    }
    /* tslint:enable: no-any */

    public map<U>(fn: (value: T) => U): Tree<U> {
        const nextChildren = this.children.map(child => {
            return child.map(fn);
        });

        return TreeNodeArray.of(fn(this.value), nextChildren);
    }

    public extend<U>(fn: (tree: Tree<T>) => U): Tree<U> {
        const nextChildren = this.children.map(child => {
            return child.extend(fn);
        });

        return TreeNodeArray.of(fn(this), nextChildren);
    }

    public reduce<U>(fn: (previousValue: U, currentValue: T) => U, acc: U): U {
        return this.children.reduce((memo, tree) => {
            return tree.reduce(fn, memo);
        }, fn(acc, this.value));
    }
}

export class TreeNodeObject<T> {
    private value: T;
    private children: ChildrenObject<T>;

    /* tslint:disable: no-any */
    public static of(value: any, children: ChildrenObject<any>): TreeNodeObject<any> {
        return new TreeNodeObject(value, children);
    }
    /* tslint:enable: no-any */

    constructor(value: T, children: ChildrenObject<T>) {
        this.value = value;
        this.children = children;
    }

    public of(value: T, children: ChildrenObject<T>): TreeNodeObject<T> {
        return new TreeNodeObject(value, children);
    }

    /* tslint:disable: no-any */
    public extract(): GenericObject<any> {
        return mapObj((child) => child.extract(), this.children);
    }
    /* tslint:enable: no-any */

    public map<U>(fn: (value: T) => U): Tree<U> {
        const nextChildren = mapObj(child => child.map(fn), this.children);

        return TreeNodeObject.of(fn(this.value), nextChildren);
    }

    public extend<U>(fn: (tree: Tree<T>) => U): Tree<U> {
        const nextChildren = mapObj(child => child.extend(fn), this.children);

        return TreeNodeObject.of(fn(this), nextChildren);
    }

    public reduce<U>(fn: (previousValue: U, currentValue: T) => U, acc: U): U {
        return values(this.children).reduce((memo, tree) => {
            return tree.reduce(fn, memo);
        }, fn(acc, this.value));
    }
}

const myTree = TreeNodeObject.of(
    1, {
        one: Leaf.of(2),
        two: Leaf.of(3),
        three: TreeNodeArray.of(
            2, [
                Leaf.of(1),
                Leaf.of(3),
            ]
        ),
        four: TreeNodeArray.of(1, [])
    }
);

console.log(
    myTree
        .map(x => x + 1)
        .extract()
);