(*
    BSD 2-Clause License

    Copyright (c) 2018, 2019, Anton Krotov
    All rights reserved.
*)

MODULE AVLTREES;

IMPORT C := COLLECTIONS;


TYPE

    DATA* = POINTER TO RECORD (C.ITEM) END;

    NODE* = POINTER TO RECORD (C.ITEM)

        data*: DATA;

        height: INTEGER;

        left*, right*: NODE

    END;

    CMP* = PROCEDURE (a, b: DATA): INTEGER;

    DESTRUCTOR* = PROCEDURE (VAR data: DATA);


VAR

    nodes: C.COLLECTION;


PROCEDURE NewNode (data: DATA): NODE;
VAR
    node:  NODE;
    citem: C.ITEM;

BEGIN
    citem := C.pop(nodes);
    IF citem = NIL THEN
        NEW(node)
    ELSE
        node := citem(NODE)
    END;

    node.data := data;
    node.left := NIL;
    node.right := NIL;
    node.height := 1

    RETURN node
END NewNode;


PROCEDURE height (p: NODE): INTEGER;
VAR
    res: INTEGER;

BEGIN
    IF p = NIL THEN
        res := 0
    ELSE
        res := p.height
    END

    RETURN res
END height;


PROCEDURE bfactor (p: NODE): INTEGER;
    RETURN height(p.right) - height(p.left)
END bfactor;


PROCEDURE fixheight (p: NODE);
BEGIN
    p.height := MAX(height(p.left), height(p.right)) + 1
END fixheight;


PROCEDURE rotateright (p: NODE): NODE;
VAR
    q: NODE;

BEGIN
    q := p.left;
    p.left := q.right;
    q.right := p;
    fixheight(p);
    fixheight(q)

    RETURN q
END rotateright;


PROCEDURE rotateleft (q: NODE): NODE;
VAR
    p: NODE;

BEGIN
    p := q.right;
    q.right := p.left;
    p.left := q;
    fixheight(q);
    fixheight(p)

    RETURN p
END rotateleft;


PROCEDURE balance (p: NODE): NODE;
VAR
    res: NODE;

BEGIN
    fixheight(p);

    IF bfactor(p) = 2 THEN
        IF bfactor(p.right) < 0 THEN
            p.right := rotateright(p.right)
        END;
        res := rotateleft(p)

    ELSIF bfactor(p) = -2 THEN
        IF bfactor(p.left) > 0 THEN
            p.left := rotateleft(p.left)
        END;
        res := rotateright(p)

    ELSE
        res := p
    END

    RETURN res
END balance;


PROCEDURE insert* (p: NODE; data: DATA; cmp: CMP; VAR newnode: BOOLEAN; VAR node: NODE): NODE;
VAR
    res: NODE;
    rescmp: INTEGER;

BEGIN
    IF p = NIL THEN
        res := NewNode(data);
        node := res;
        newnode := TRUE
    ELSE

        rescmp := cmp(data, p.data);
        IF rescmp < 0 THEN
            p.left := insert(p.left, data, cmp, newnode, node);
            res := balance(p)
        ELSIF rescmp > 0 THEN
            p.right := insert(p.right, data, cmp, newnode, node);
            res := balance(p)
        ELSE
            res := p;
            node := res;
            newnode := FALSE
        END

    END

    RETURN res
END insert;


PROCEDURE destroy* (VAR node: NODE; destructor: DESTRUCTOR);
VAR
    left, right: NODE;

BEGIN
    IF node # NIL THEN
        left  := node.left;
        right := node.right;

        IF destructor # NIL THEN
            destructor(node.data)
        END;

        C.push(nodes, node);
        node := NIL;

        destroy(left, destructor);
        destroy(right, destructor)
    END
END destroy;


BEGIN
    nodes := C.create()
END AVLTREES.