Red-Black Tree with Zipper in Prolog

red_black_tree.prolog

:- op(200, xfy, <--).
:- op(200, xfy, /^\).
:- op(200, xfy, /^ ).
:- op(200, xfy,  ^\).
 
tree_empty(zip(nil, [])).
 
tree_lhs(zip(L, [step(lhs, V, R)|XS]), zip(node(V, L, R), XS)).
tree_rhs(zip(R, [step(rhs, V, L)|XS]), zip(node(V, L, R), XS)).
 
tree_up(zip(node(V, L, R), T), zip(L, [step(lhs, V, R)|T])) :- !.
tree_up(zip(node(V, L, R), T), zip(R, [step(rhs, V, L)|T])) :- !.
 
tree_opp( zip(X, [step(DX, V, Y)|T])
        , zip(Y, [step(DY, V, X)|T])) :-
    ((DX, DY) = (lhs, rhs)
    ;(DY, DX) = (lhs, rhs)), !.
 
U <-- B :- tree_up(U, B).
L /^\ R :- tree_opp(L, R).
L /^  U :- tree_lhs(L, U).
U  ^\ R :- tree_rhs(R, U).
 
tree_dir(X, zip(_, [step(X, _, _)|_])).
 
tree_top(zip(X, []), zip(X, [])) :- !.
tree_top(A, B) :-
  C <-- B, tree_top(A, C).
 
tree_put(zip(node(V, nil, nil), XS), zip(_, XS), V).
 
tree_val(V, zip(node(V, _, _), _)).
tree_val(zip(node(V, L, R), XS), V, zip(node(_, L, R), XS)).
 
tree_rot(Z, O) :- (
  tree_dir(lhs, O)
  -> tree_rrot(Z, O)
   ; tree_lrot(Z, O)).
 
tree_lrot(Z, O) :-
  A <-- O,
  zip(node(OV, OL, OR), _) = O,
  zip(node(AV, AL, AR), L) = A,
  Z = zip(node(OV, node(AV, AL, OL), OR), L).
 
tree_rrot(Z, O) :-
  A <-- O,
  zip(node(OV, OL, OR), _) = O,
  zip(node(AV, AL, AR), L) = A,
  Z = zip(node(OV, OL, node(AV, OR, AR)), L).
 
tree_member(Z, Zip) :- tree_top(zip(X, _), Zip), !, tree_member(Z, X).
tree_member(E, node(_, L, _)) :- tree_member(E, L).
tree_member(E, node(E, _, _)).
tree_member(E, node(_, _, R)) :- tree_member(E, R).
 
rbt_color(   O, Col) :- tree_val(rbt(_, Col), O).
rbt_color(Z, O, Col) :-
  tree_val(   rbt(V,   _), O),
  tree_val(Z, rbt(V, Col), O).
 
rbt_put(Z, O, V) :-
  rbt_put0(A, O, V),
  rbt_balance(B, A),
  rbt_color(Z, B, black).
 
rbt_put0(Z, O, V) :-
  tree_val(rbt(V0, Col), O), !,
  ( V @< V0, !, C /^ O, rbt_put0(Z, C, V)
  ; V @> V0, !, O ^\ C, rbt_put0(Z, C, V)
  ; fail).
 
rbt_put0(Z, O, V) :-
  tree_put(Z, O, rbt(V, red)).
 
rbt_val(Z, A) :- tree_val(rbt(Z, _), A).
 
rbt_balance(Z, O1) :-
  \+(P1 <-- O1), !, tree_top(Z, O1).
 
rbt_balance(Out, O1) :-
  P1 <-- O1, (rbt_color(P1, black), !, tree_top(Out, O1);
  G1 <-- P1,
  Q1 /^\ P1,
  (rbt_color(Q1, red), !,
    rbt_color(P2, P1, black), P2 /^\ Q2,
    rbt_color(Q3, Q2, black), G2 <-- Q3,
    rbt_color(G3, G2,   red),
    rbt_balance(Out, G3)
    ;
    ((tree_dir(Dir, O1)
    , tree_dir(Dir, P1)), !
       , P2 = P1
       ; tree_rot(P2, O1)),
    tree_rot(P3, P2),
    rbt_color(P4, P3, black), P4 <-- G2,
    rbt_color(G3, G2,   red), G3 /^\ G4,
    rbt_color(G5, G4,   red),
    tree_top(Out, G5))).

rbt_member(E, A) :- tree_member(rbt(E, _), A).

usage.prolog

tree_dump(A) :- tree_top(zip(X, _), A), !, tree_dump(X).
tree_dump(A) :- tree_dump(A, ' ', '─', true).
 
tree_dump(X, Ln1, Con, End) :-
  print(Ln1),
  put_char(Con),
  put_char('─'),
  put_char(' '),
  tree_dump_node(X),
  nl,
  (( X = node(_, L, R)
    , (L \== nil; R \== nil), !
    , (End, !
        , atom_concat(Ln1, '    ', Ln2)
        ; atom_concat(Ln1,  '│  ', Ln2))
    , tree_dump(L, Ln2, '├', fail)
    , tree_dump(R, Ln2, '└', true)), !; true).
 
tree_dump_node(nil)           :- write(nil).
tree_dump_node(node(V, _, _)) :- write(V).

step(_, A, A, _) :- !, fail.
step(A, A, B, _).
step(E, A, B, U) :- A1 is A + U, step(E, A1, B, U).
 
?- tree_empty(T0)
 , findall(E, step(E, 1, 103,  3), X)
 , findall(E, step(E, 102, 0, -3), Y)
 , findall(E, step(E, 2, 104,  3), Z)
 , rbt_put_all(T1, T0, X)
 , rbt_put_all(T2, T1, Y)
 , rbt_put_all(T3, T2, Z)
 , tree_dump(T3), nl
 , findall(E, rbt_member(E, T3), Out)
 , write(Out), nl
 .