// AVL Tree Example

Node `alias Nullable[Table[Node left; Node right; I64 height; Any key; Any value]]

Avl_Compare_Fn `alias Fn[Result[I64 value; Error error]; Any key_a; Any key_b]

AVL `alias Table[Node root; U64 size; Avl_Compare_Fn compare_fn]

height : `{[Node node] ? node {[node] node.height } { 0 } }

max : `{[I64 a; I64 b] ? a > b { a } { b } }

get_balance : `{[Node n]
    ? n {
        - (height `sync n.left; height `sync n.right)
    } {
        0
    }
}

min_value_node : `{[Node n]
    @ n.left { n : n.left }
    n
}

new_node : `{[Any key; Any value]
    Node (.left `null; .right `null; .height 1; .key key; .value value)
}

right_rotate : `{[Node y]
    ? !y { `return `null }
    x : y.left;
    t2 : x.right;
    y.height : 1 + max `sync (height `sync y.left; height `sync y.right)
    x.height : 1 + max `sync (height `sync x.left; height `sync x.right)
    x
}

left_rotate : `{[Node x]
    ? !x { `return `null }
    y : x.right;
    t2 : y.left;
    x.height : 1 + max `sync (height `sync x.left; height `sync x.right)
    y.height : 1 + max `sync (height `sync y.left; height `sync y.right)
    y
}

`export avl_init : `{[Avl_Compare_Fn compare_fn]
    AVL (.root `null; .size 0; .compare_fn compare_fn)
}

`export avl_insert : `{[AVL avl; Any key; Any value]
    insert : `{[Node n]
        ? !n {
            avl.size +: 1
            `return new_node `sync (key; value)
        }
        cmp : avl.compare_fn `sync (n.key; key)
        ? {
            cmp < 0 { n.left : insert `sync n.left }
            cmp > 0 { n.right : insert `sync n.right }
            {
                n.value : value
                `return n
            }
        }
        n.height : 1 + max `sync (height `sync n.left; height `sync n.right)
        balance : get_balance `sync n
        ? {
            & (balance > 1; cmp < 0) { right_rotate `sync n }
            & (balance < -1; cmp > 0) { left_rotate `sync n }
            & (balance > 1; cmp > 0) {
                n.left : left_rotate `sync n.left
                right_rotate `sync n
            }
            & (balance < -1; cmp < 0) {
                n.right : right_rotate `sync n.right
                left_rotate `sync n
            }
            { n }
        }
    }
    avl.root : insert `sync avl.root
    Void
}

`export avl_delete : `{[Avl avl; Any key]
    delete : `{[Node n; Any key]
        ? !n { `return n }
        cmp : avl.compare_fn `sync (n.key; key)
        ? {
            cmp < 0 { n.left : delete `sync (n.left; key)
            cmp > 0 { n.right : delete `sync (n.right; key)
            {
                ? | (!n.left; !n.right) {
                    n : ? n.left { n.left } { n.right }
                } {
                    tmp : min_value_node `sync n.right
                    n.key : tmp.key
                    n.value : tmp.value
                    n.right : delete `sync (n.right; tmp.key)
                }
            }
        }
        ? !n { `return n }
        n.height : 1 + max `sync (height `sync n.left; height `sync n.right)
        balance : get_balance `sync n
        ? {
            & (balance > 1; 0 <= get_balance `sync n.left) { right_rotate `sync n }
            & (balance > 1; 0 > get_balance `sync n.left {
                n.left : left_rotate `sync n.left
                right_rotate `sync n
            }
            & (balance < -1; 0 >= get_balance `sync n.right) { left_rotate `sync n }
            & (balance < -1; 0 > get_balance `sync n.right) {
                n.right : right_rotate `sync n.right
                left_rotate `sync n
            }
            { n }
        }
    }
    avl.root : delete `sync (avl.root; key)
    Void
}