interveraltree.ts

declare function binarySearch<T>(array: T[], value: T, comparer?: (v1: T, v2: T) => number, offset?: number): number;
declare function insertItemAt<T>(array: T[], offset: number, value: T): void;
declare function orderedRemoveItemAt<T>(array: T[], index: number): void;
declare function lastOrUndefined<T>(array: T[]): T;

interface TextRange { pos: number, end: number }

class IntervalTree<T extends TextRange> {
    public count: number = 0;

    private root: IntervalNode<T>;

    public get(pos: number, end: number) {
        return this._get(this.root, pos, end);
    }

    public add(value: T) {
        const context = { wasTreeModified: false };
        this.root = this._add(this.root, value, context);
        this.count++;
        return true;
    }

    public remove(pos: number, end: number) {
        const context = { wasTreeModified: false, wasDeleted: false };
        this.root = this._remove(this.root, pos, end, context);
        if (context.wasDeleted) {
            this.count--;
            return true;
        }
        return false;
    }

    public clear() {
        this.root = undefined;
        this.count = 0;
    }

    public findFirstStartingAt(pos: number) {
        return this._findStartingAt(this.root, pos);
    }

    public findFirstContaining(pos: number) {
        return this.findFirstIntersectingWith(pos, pos);
    }

    public findFirstIntersectingWith(pos: number, end: number) {
        return this._findIntersectingWith(this.root, pos, end);
    }

    public findAllStartingAt(pos: number): T[] {
        const results: T[] = [];
        this._findStartingAt(this.root, pos, results);
        return results;
    }

    public findAllContaining(pos: number): T[] {
        return this.findAllIntersectingWith(pos, pos);
    }

    public findAllIntersectingWith(pos: number, end: number): T[] {
        const results: T[] = [];
        this._findIntersectingWith(this.root, pos, end, results);
        return results;
    }

    private rotateLeft(node: IntervalNode<T>) {
        const right = node.right;
        node.right = right.left;
        right.left = node;
        return right;
    }

    private rotateRight(node: IntervalNode<T>) {
        const left = node.left;
        node.left = left.right;
        left.right = node;
        return left;
    }

    private findMin(node: IntervalNode<T>) {
        if (node) {
            while (node.left) {
                node = node.left;
            }
        }
        return node;
    }

    private findMax(node: IntervalNode<T>) {
        if (node) {
            while (node.right) {
                node = node.right;
            }
        }
        return node;
    }

    private _add(node: IntervalNode<T>, value: T, context: { wasTreeModified?: boolean }) {
        if (!node) {
            node = new IntervalNode(value);
            context.wasTreeModified = true;
        }
        else {
            if (value.pos < node.pos) {
                node.left = this._add(node.left, value, context);
                if (context.wasTreeModified) {
                    node.balance--;
                    if (node.balance === 0) {
                        context.wasTreeModified = false;
                    }
                    else if (node.balance === -2) {
                        if (node.left.balance === 1) {
                            const leftRightBalance = node.left.right.balance;
                            node.left = this.rotateLeft(node.left);
                            node = this.rotateRight(node);
                            node.balance = 0;
                            node.left.balance = leftRightBalance === 1 ? -1 : 0;
                            node.right.balance = leftRightBalance === -1 ? 1 : 0;
                        }
                        else if (node.left.balance === -1) {
                            node = this.rotateRight(node);
                            node.balance = 0;
                            node.right.balance = 0;
                        }
                        context.wasTreeModified = false;
                    }
                }
            }
            else if (value.pos > node.pos) {
                node.right = this._add(node.right, value, context);
                if (context.wasTreeModified) {
                    node.balance++;
                    if (node.balance === 0) {
                        context.wasTreeModified = false;
                    }
                    else if (node.balance === 2) {
                        if (node.right.balance === -1) {
                            const rightLeftBalance = node.right.left.balance;
                            node.right = this.rotateRight(node.right);
                            node = this.rotateLeft(node);
                            node.balance = 0;
                            node.left.balance = rightLeftBalance === 1 ? -1 : 0;
                            node.right.balance = rightLeftBalance === -1 ? 1 : 0;
                        }
                        else if (node.right.balance === 1) {
                            node = this.rotateLeft(node);
                            node.balance = 0;
                            node.left.balance = 0;
                        }
                        context.wasTreeModified = false;
                    }
                }
            }
            else {
                node.add(value);
                context.wasTreeModified = false;
            }
        }
        return node;
    }

    private _remove(node: IntervalNode<T>, pos: number, end: number, context: { wasTreeModified: boolean, wasDeleted: boolean }) {
        context.wasTreeModified = false;
        context.wasDeleted = false;
        if (!node) return undefined;
        if (pos < node.pos) {
            if (node.left) {
                node.left = this._remove(node.left, pos, end, context);
                if (context.wasTreeModified) {
                    node.balance++;
                }
            }
        }
        else if (pos > node.pos) {
            if (node.right) {
                node.right = this._remove(node.right, pos, end, context);
                if (context.wasTreeModified) {
                    node.balance--;
                }
            }
        }
        else {
            if (end === node.end && node.values.length === 1) {
                if (!node.left || !node.right) {
                    context.wasTreeModified = true;
                    context.wasDeleted = true;
                    return node.right || node.left;
                }
                const minNode = this.findMin(node.right);
                const nodePos = node.pos;
                const nodeEnd = node.end;
                node.swap(minNode);
                node.right = this._remove(node.right, nodePos, nodeEnd, context);
                if (context.wasTreeModified) {
                    node.balance--;
                }
            } 
            else {
                context.wasDeleted = node.remove(end);
            }
        }
        if (context.wasTreeModified) {
            if (node.balance === 1 || node.balance === -1) {
                context.wasTreeModified = false;
                return node;
            }
            else if (node.balance == -2) {
                if (node.left.balance == 1) {
                    const leftRightBalance = node.left.right.balance;
                    node.left = this.rotateLeft(node.left);
                    node = this.rotateRight(node);
                    node.balance = 0;
                    node.left.balance = leftRightBalance == 1 ? -1 : 0;
                    node.right.balance = leftRightBalance == -1 ? 1 : 0;
                }
                else if (node.left.balance == -1) {
                    node = this.rotateRight(node);
                    node.balance = 0;
                    node.right.balance = 0;
                }
                else if (node.left.balance == 0) {
                    node = this.rotateRight(node);
                    node.balance = 1;
                    node.right.balance = -1;
                    context.wasTreeModified = false;
                }
            }
            else if (node.balance == 2) {
                if (node.right.balance == -1) {
                    const rightLeftBalance = node.right.left.balance;
                    node.right = this.rotateRight(node.right);
                    node = this.rotateLeft(node);
                    node.balance = 0;
                    node.left.balance = rightLeftBalance == 1 ? -1 : 0;
                    node.right.balance = rightLeftBalance == -1 ? 1 : 0;
                }
                else if (node.right.balance == 1) {
                    node = this.rotateLeft(node);
                    node.balance = 0;
                    node.left.balance = 0;
                }
                else if (node.right.balance == 0) {
                    node = this.rotateLeft(node);
                    node.balance = -1;
                    node.left.balance = 1;
                    context.wasTreeModified = false;
                }
            }
        }
        return node;
    }

    private _get(node: IntervalNode<T>, pos: number, end: number) {
        if (node) {
            if (pos < node.pos) {
                return this._get(node.left, pos, end);
            }
            else if (pos > node.pos) {
                return this._get(node.right, pos, end);
            }
            else {
                return node.get(end);
            }
        }
        return undefined;
    }

    private _findStartingAt(node: IntervalNode<T>, pos: number, results?: T[]) {
        if (!node) return undefined;
        if (pos < node.pos) {
            return this._findStartingAt(node.left, pos, results);
        }
        else if (pos > node.pos) {
            return this._findStartingAt(node.right, pos, results);
        }
        else {
            for (let i = node.values.length - 1; i >= 0; i--) {
                if (results) {
                    results.push(node.values[i]);
                }
                else {
                    return node.values[i];
                }
            }
        }
        return undefined;
    }

    private _findIntersectingWith(node: IntervalNode<T>, pos: number, end: number, results?: T[]) {
        if (!node) return undefined;
        if (end <= node.pos) {
            // all results must be to the left
            return this._findIntersectingWith(node.left, pos, end, results);
        } 
        if (pos >= node.max) {
            // all results must be to the right
            return this._findIntersectingWith(node.right, pos, end, results);
        }

        return this._findIntersectingWith(node.left, pos, end, results)
            || node.intersect(pos, end, results)
            || this._findIntersectingWith(node.right, pos, end, results);
    }
}

class IntervalNode<T extends TextRange> {
    public ends: number[];
    public values: T[];
    public balance: number = 0;

    private _left: IntervalNode<T>;
    private _right: IntervalNode<T>;
    private _max: number;
    private _pos: number;
    private _end: number;
    private _parent: IntervalNode<T>;

    constructor(value: T) {
        this._pos = value.pos;
        this._end = value.end;
        this._max = value.end;
        this.add(value);
    }

    public get pos() { return this._pos; }
    public set pos(value: number) {
        if (this._pos !== value) {
            this._pos = value;
            this.resetMax();
        }
    }

    public get end() { return this._end; }
    public set end(value: number) {
        if (this._end !== value) {
            this._end = value;
            this.resetMax();
        }
    }

    public get max() {
        if (this._max === -1) {
            let max = this._end;
            if (this._left) max = Math.max(max, this._left.max);
            if (this._right) max = Math.max(max, this._right.max);
            this._max = max;
        }
        return this._max; 
    }

    public get left() { return this._left; }
    public set left(value: IntervalNode<T>) {
        if (this._left !== value) {
            if (this._left) this._left._parent = undefined;
            this._left = value;
            this.resetMax();
            if (this._left) this._left._parent = this;
        }
    }

    public get right() { return this._right; }
    public set right(value: IntervalNode<T>) {
        if (this._right !== value) {
            if (this._right) this._right._parent = undefined;
            this._right = value;
            this.resetMax();
            if (this._right) this._right._parent = this;
        }
    }

    public swap(node: IntervalNode<T>) {
        const pos = this.pos;
        const end = this.end;
        const ends = this.ends;
        const values = this.values;
        this.pos = node.pos;
        this.end = node.end;
        this.ends = node.ends;
        this.values = node.values;
        node.pos = pos;
        node.end = end;
        node.ends = ends;
        node.values = values;
    }

    public get(end: number) {
        const offset = binarySearch(this.ends, -end);
        if (offset >= 0) {
            return this.values[offset];
        }
        return undefined;
    }

    public add(value: T) {
        let offset = binarySearch(this.ends, -value.end);
        if (offset < 0) offset = ~offset;
        insertItemAt(this.ends, offset, -value.end);
        insertItemAt(this.values, offset, value);
        this.end = -this.ends[0];
    }

    public remove(end: number) {
        if (this.ends.length == 1) return false;
        const offset = binarySearch(this.ends, -end);
        if (offset < 0) return false;
        orderedRemoveItemAt(this.ends, offset);
        orderedRemoveItemAt(this.values, offset);
        this.end = -this.ends[0];
        return true;        
    }

    public intersect(pos: number, end: number, results?: T[]) {
        if (pos <= this.end && end >= this.pos) {
            for (let i = 0; i < this.ends.length; i++) {
                if (pos < -this.ends[i]) {
                    if (results) {
                        results.push(this.values[i]);
                    }
                    else {
                        return this.values[i];
                    }
                }
            }
        }
        return undefined;
    }

    public forEach<U>(callback: (value: T, pos: number, end: number) => U): U | undefined {
        for (let i = 0; i < this.ends.length; i++) {
            const result = callback(this.values[i], this.pos, -this.ends[i]);
            if (result) {
                return result;
            }
        }
        return undefined;
    }

    public forEachInverted<U>(callback: (value: T, pos: number, end: number) => U): U | undefined {
        for (let i = this.ends.length - 1; i >= 0; i--) {
            const result = callback(this.values[i], this.pos, -this.ends[i]);
            if (result) {
                return result;
            }
        }
        return undefined;
    }

    private resetMax() {
        this._max = -1;
        if (this._parent && this._parent._max !== -1) {
            this._parent.resetMax();
        }
    }
}