% Fuad Tabba (cs.auckland.ac.nz at fuad OR altabba.org at fuad)
% June, 2008
% Posted at http://www.altabba.org/2008/06/full-red-black-tree-implementation-in.html

% In case anyone might actually want to use this code, I've decided to make it
% public domain, unless there are portions of it that infringe on other people's
% copyrights, in that case those people's permissions must be sought.

% References:-
% http://semanticvector.blogspot.com/2008/05/red-black-tree-in-2-hours.html
% http://www.eecs.usma.edu/webs/people/okasaki/jfp99.ps
% http://sage.mc.yu.edu/kbeen/teaching/algorithms/resources/red-black-tree.html
% http://erlang.org/download/erlang-book-part1.pdf


% Thanks:-
% Nick Hay, Robert Virding, Gleb Peregud and all the nice folks at erlang-questions@erlang.org

% Node Description:-
% {Key, Value, Color, LeftSubTree, RightSubTree}
% Color is: r (Red), b (Black), bb (Doubly Black - a transient state)

% Known bug: if the Key is fakenil, and its value is 0, then it might
% dissappear. This is because fakenil is a special key used sometimes.
% Need to tacke it slightly differently.


-module(rbtree).
-export([insert/3, delete/2, lookup/2, isSane/1]).
-compile(export_all). % For debugging purposes


% Whether the node is black or not (empty trees {} are black)
% Done this way to be able to use it as a guard.
% (Thanks Gleb Peregud)
-define(IS_BLACK(Tree), ((Tree =:= {}) orelse (element(3, Tree) =:= b))).


% Inserts a new node into the tree
insert(Key, Value, Tree) ->
    % Invariant: Root must be black
    makeBlack(ins(Key, Value, Tree)).


% Removes an existing node from the tree
delete(Key, Tree) ->
    % Invariant: Root must be black
    fakenilfix(makeBlack(del(Key, Tree))).


% Gets the {Key, Value} tuple corresponding the the key
% Normal BST lookup

% Not found - returns {}.
% TODO: Maybe should throw an exception instead?
lookup(_, {}) ->
    {};

lookup(Key, {Key, Value, _, _, _}) ->
    {Key, Value};

lookup(Key, {K2, _, _, Left, _}) when Key < K2  ->
    lookup(Key, Left);

lookup(Key, {K2, _, _, _, Right}) when Key > K2  ->
    lookup(Key, Right).


%
% Insertion Algorithm: do a BST insertion and balance as you go along
%
% Does nothing if there's a collision
% TODO: maybe throw an exception?
%

% Inserting a node into an empty tree: just add the node and color it red
ins(Key, Value, {}) ->
    {Key, Value, r, {}, {}};

% Collision with an existing key: leave the tree as it is
ins(Key, _, {Key, Value, Color, Left, Right}) ->
    {Key, Value, Color, Left, Right};

% Normal BST insertion (while rebalancing the tree)
ins(Key, Value, {Key2, Value2, Color, Left, Right}) when Key < Key2 ->
    balance({Key2, Value2, Color, ins(Key, Value, Left), Right});

ins(Key, Value, {Key2, Value2, Color, Left, Right}) when Key > Key2 ->
    balance({Key2, Value2, Color, Left, ins(Key, Value, Right)}).


% Rules for rebalancing the tree as described by Okasaki:-
% http://www.eecs.usma.edu/webs/people/okasaki/jfp99.ps

balance({Kz, Vz, b, {Ky, Vy, r, {Kx, Vx, r, A, B}, C}, D}) ->
    {Ky, Vy, r, {Kx, Vx, b, A, B}, {Kz, Vz, b, C, D}};

balance({Kz, Vz, b, {Kx, Vx, r, A, {Ky, Vy, r, B, C}}, D}) ->
    {Ky, Vy, r, {Kx, Vx, b, A, B}, {Kz, Vz, b, C, D}};

balance({Kx, Vx, b, A, {Kz, Vz, r, {Ky, Vy, r, B, C}, D}}) ->
    {Ky, Vy, r, {Kx, Vx, b, A, B}, {Kz, Vz, b, C, D}};

balance({Kx, Vx, b, A, {Ky, Vy, r, B, {Kz, Vz, r, C, D}}}) ->
    {Ky, Vy, r, {Kx, Vx, b, A, B}, {Kz, Vz, b, C, D}};

% No rebalancing needed
balance(Tree) ->
    Tree.


% Only used to ensure that the root is always black (an RB-Tree invariant)
makeBlack({Key, Value, _, Left, Right}) ->
    {Key, Value, b, Left, Right}.


%
% Deletion: do a BST deletion and fix the mess as you go
%
% Does nothing if the key doesn't exists.
% TODO: maybe throw an exception?
%

% Deleting from an empty tree: do nothing
del(_, {}) ->
    {};

% Deleting a red node, doesn't violate RB invariants
del(Key, {Key, _, r, Child, {}}) ->
    Child;

del(Key, {Key, _, r, {}, Child}) ->
    Child;

% Deleting a black node - violates invariants, give child a black token to
% indicate that subtree is missing a black node
del(Key, {Key, _, b, Child, {}}) ->
    blackToken(Child);

del(Key, {Key, _, b, {}, Child}) ->
    blackToken(Child);

% Normal BST deletion (while fixing any messes being made)
del(Key, {Key, _, Color, Left, Right}) ->
    {Km, Vm} = max(Left),
    delfix({Km, Vm, Color, del(Km, Left), Right});

del(Key, {K2, V2, C2, Left, Right}) when Key < K2  ->
    delfix({K2, V2, C2, del(Key, Left), Right});

del(Key, {K2, V2, C2, Left, Right}) when Key > K2  ->
    delfix({K2, V2, C2, Left, del(Key, Right)}).


% Fixing the red black tree after a deletion, based on:-
% http://sage.mc.yu.edu/kbeen/teaching/algorithms/resources/red-black-tree.html

% Case A (partial - doesn't handle making the root black):
blackToken({Key, Value, r, Left, Right}) ->
    {Key, Value, b, Left, Right};

% A black node becomes doubly black
blackToken({Key, Value, b, Left, Right}) ->
    {Key, Value, bb, Left, Right};

% An empty leaf (which is actually black, becomes doubly black).
blackToken({}) ->
    {fakenil, 0, bb, {}, {}}. % transient node, removed before the end of the operation


% Reverts a fakenil into {} once it loses its token
% Apply to all cases where a node potentially loses a token
fakenilfix({fakenil, 0, b, {}, {}}) ->
    {};

fakenilfix(Tree) ->
    Tree.


% delfix: Applies cases B, C and D

% Case B
delfix({Ky, Vy, Cy, {Kx, Vx, bb, A, B}, {Kz, Vz, b, C, D}}) when ?IS_BLACK(C) andalso ?IS_BLACK(D) ->
    blackToken({Ky, Vy, Cy, fakenilfix({Kx, Vx, b, A, B}), {Kz, Vz, r, C, D}});

delfix({Ky, Vy, Cy, {Kx, Vx, b, A, B}, {Kz, Vz, bb, C, D}}) when ?IS_BLACK(A) andalso ?IS_BLACK(B)->
    blackToken({Ky, Vy, Cy, {Kx, Vx, r, A, B}, fakenilfix({Kz, Vz, b, C, D})});

% Case C
delfix({Ky, Vy, b, {Kx, Vx, bb, A, B}, {Kz, Vz, r, C, D}}) when ?IS_BLACK(C) andalso ?IS_BLACK(D)->
    {Kz, Vz, b, delfix({Ky, Vy, r, {Kx, Vx, bb, A, B}, C}), D};

delfix({Ky, Vy, b, {Kx, Vx, r, A, B}, {Kz, Vz, bb, C, D}}) when ?IS_BLACK(A) andalso ?IS_BLACK(B)->
    {Kx, Vx, b, A, delfix({Ky, Vy, r, B, {Kz, Vz, bb, C, D}})};

% Case D

% NOTE: in terms of performance it might be better if subcase i got transformed
% immediately to the result of subcase ii, but for clarity I'm leaving it this
% way for the time being.

% Subcase i
delfix({Ky, Vy, Cy, {Kx, Vx, bb, A, B}, {Kw, Vw, b, {Kz, Vz, r, C, D}, E}}) when ?IS_BLACK(E) ->
    delfix({Ky, Vy, Cy, {Kx, Vx, bb, A, B}, {Kz, Vz, b, C, {Kw, Vw, r, D, E}}});

delfix({Kz, Vz, Cz, {Kx, Vx, b, A, {Ky, Vy, r, B, C}}, {Kw, Vw, bb, D, E}}) when ?IS_BLACK(A) ->
    delfix({Kz, Vz, Cz, {Ky, Vy, b, {Kx, Vx, r, A, B}, C}, {Kw, Vw, bb, D, E}});

% Subcase ii
delfix({Ky, Vy, Cy, {Kx, Vx, bb, A, B}, {Kz, Vz, b, C, {Kw, Vw, r, D, E}}}) when ?IS_BLACK(E) ->
    {Kz, Vz, Cy, {Ky, Vy, b, fakenilfix({Kx, Vx, b, A, B}), C}, {Kw, Vw, b, D, E}}; % z takes on y's color

delfix({Kz, Vz, Cz, {Ky, Vy, b, {Kx, Vx, r, A, B}, C}, {Kw, Vw, bb, D, E}}) when ?IS_BLACK(A) ->
    {Ky, Vy, Cz, {Kx, Vx, b, A, B}, {Kz, Vz, b, C, fakenilfix({Kw, Vw, b, D, E})}}; % y takes on z's color

% All other cases, leave it as it is
delfix(Tree) ->
    Tree.


% Get the maximum value in the tree (used when deleting nodes)
max({Key, Value, _, _, {}}) ->
    {Key, Value};

max({_, _, _, _, Right}) ->
    max(Right).


% True if the tree is rooted in a black node (applies to empty trees too)
% NOTE: Dead code, I'm using a macro instead
isBlack({}) ->
    true;

isBlack({_, _, b, _, _}) ->
    true;

isBlack(_) ->
    false.


% Performs a sanity check on the tree, check for all the BST and RB invariants
isSane({}) ->
    true;

% Root must be black
isSane({_, _, r, _, _}) ->
    false;

isSane(T) ->
    checkBSTOrder(T) andalso checkRBProperty(T) andalso (countBlack(T) =/= false).


% Check that the binary search tree invariants (ordering) are held
checkBSTOrder({_, _, _, {}, {}}) ->
    true;

checkBSTOrder({Key, _, _, {}, Right}) ->
   {Kr, _, _, _, _} = Right,
   (Key < Kr) andalso checkBSTOrder(Right);

checkBSTOrder({Key, _, _, Left, {}}) ->
   {Kl, _, _, _, _} = Left,
   (Kl < Key) andalso checkBSTOrder(Left);

checkBSTOrder({Key, _, _, Left, Right}) ->
   {Kl, _, _, _, _} = Left,
   {Kr, _, _, _, _} = Right,
   (Kl < Key) andalso (Key < Kr) andalso checkBSTOrder(Left) andalso checkBSTOrder(Right).


% Check that there aren't two red nodes in a row.
checkRBProperty({}) ->
    true;

checkRBProperty({_, _, b, Left, Right}) ->
    checkRBProperty(Left) andalso checkRBProperty(Right);

checkRBProperty({_, _, r, Left, Right}) ->
    checkRBRed(Left) andalso checkRBRed(Right).


checkRBRed({}) ->
    true;

checkRBRed({_, _, b, Left, Right}) ->
    checkRBProperty(Left) andalso checkRBProperty(Right);

checkRBRed({_, _, r, _, _}) ->
    false.


% Check that the number of black nodes is equal in all paths to the leaves

% Leaves aren't counted (Even though they are technically black, doesn't matter)
countBlack({}) ->
    0;

countBlack({_, _, b, Left, Right}) ->
    CountLeft = countBlack(Left),
    CountRight = countBlack(Right),
    if
        (CountLeft =:= CountRight) andalso (CountLeft =/= false) ->
            1 + CountLeft;
        true -> false
    end;

countBlack({_, _, r, Left, Right}) ->
    CountLeft = countBlack(Left),
    CountRight = countBlack(Right),
    if
        (CountLeft =:= CountRight) andalso (CountLeft =/= false) ->
            CountLeft;
        true -> false
    end.


% Creates the tree from:-
% http://sage.mc.yu.edu/kbeen/teaching/algorithms/resources/red-black-tree.html
createTree() ->
    S1 = {},
    S2 = insert(260,b,S1),
    S3 = insert(240,b,S2),
    S4 = insert(140,b,S3),
    S5 = insert(320,b,S4),
    S6 = insert(250,r,S5),
    S7 = insert(170,r,S6),
    S8 = insert(60,b,S7),
    S9 = insert(100,b,S8),
    S10 = insert(20,b,S9),
    S11 = insert(40,b,S10),
    S12 = insert(290,b,S11),
    S13 = insert(280,b,S12),
    S14 = insert(270,r,S13),
    S15 = insert(30,r,S14),
    S16 = insert(265,r,S15),
    S17 = insert(275,b,S16),
    S18 = insert(277,r,S17),
    insert(278,r,S18).