Mercury binary tree example
:- module main.
%=============================================================================%
% Implementation of an extremely simple binary tree data type.
% Has very modest optimizations to avoid unnecessary allocations.
:- interface.
%=============================================================================%
:- use_module io.
%-----------------------------------------------------------------------------%
:- pred main(io.io::di, io.io::uo) is det.
%-----------------------------------------------------------------------------%
:- type xbtree(K, V).
%-----------------------------------------------------------------------------%
:- func init = xbtree(K, V).
%-----------------------------------------------------------------------------%
:- func singleton(K, V) = xbtree(K, V).
%-----------------------------------------------------------------------------%
:- pred insert(K::in, V::in, xbtree(K, V)::in, xbtree(K, V)::out) is semidet.
%-----------------------------------------------------------------------------%
:- pred update(K::in, V::in, xbtree(K, V)::in, xbtree(K, V)::out) is semidet.
%-----------------------------------------------------------------------------%
:- pred det_insert(K::in, V::in, xbtree(K, V)::in, xbtree(K, V)::out) is det.
:- pred det_update(K::in, V::in, xbtree(K, V)::in, xbtree(K, V)::out) is det.
%-----------------------------------------------------------------------------%
:- pred set(K::in, V::in, xbtree(K, V)::in, xbtree(K, V)::out) is det.
%-----------------------------------------------------------------------------%
:- pred search(xbtree(K, V)::in, K::in, V::out) is semidet.
%-----------------------------------------------------------------------------%
:- pred lookup(xbtree(K, V)::in, K::in, V::out) is det.
%=============================================================================%
% Implementation of a red/black tree
:- implementation.
%=============================================================================%
:- use_module maybe.
:- use_module bool.
:- use_module exception.
%-----------------------------------------------------------------------------%
det_insert(K, V, In, Out) :-
( insert(K, V, In, Tree) -> Tree = Out
; exception.throw(exception.software_error("Could not insert!")) ).
%-----------------------------------------------------------------------------%
det_update(K, V, In, Out) :-
( update(K, V, In, Tree) -> Tree = Out
; exception.throw(exception.software_error("Could not update!")) ).
%-----------------------------------------------------------------------------%
lookup(Tree, K, V) :-
( search(Tree, K, X) -> V = X
; exception.throw(exception.software_error("Could not lookup!")) ).
%-----------------------------------------------------------------------------%
:- use_module string.
:- pred verify(xbtree(string, int)::in, string::in, int::in, int::out) is det.
verify(Tree, Key, In, Out) :-
( search(Tree, Key, X) ->
( X = In ->
Out = X,
trace [io(!IO)] (
io.write_string("Key ", !IO),
io.write_string(Key, !IO),
io.write_string(" was ", !IO),
io.write_int(X, !IO),
io.nl(!IO)
)
; exception.throw(exception.software_error(
string.append(string.append(string.append(
"Expected ", Key),
" to equal "),
string.from_int(In)))) )
;
exception.throw(exception.software_error(string.append("Missing key ", Key)))
).
main(!IO) :-
init = Tree0,
det_insert("X", 216, Tree0, Tree1),
det_insert("Y", 97, Tree1, Tree2),
det_insert("Z", 312, Tree2, Tree3),
det_update("Y", 409, Tree3, Tree4),
set("W", 1717, Tree4, Tree5),
set("X", 292, Tree5, Tree6),
verify(Tree6, "X", 292, _),
verify(Tree1, "X", 216, _),
verify(Tree6, "Y", 409, _),
verify(Tree6, "Z", 312, _),
verify(Tree6, "W", 1717, _),
( not search(Tree6, "A", _) -> true
; exception.throw(exception.software_error("Unexpected search success!")) ),
io.write_string("Success!\n", !IO).
%-----------------------------------------------------------------------------%
:- type children(T) --->
none ;
one(T) ;
two(T, T).
:- type node(K, V) --->
node(k::K, v::V, less::maybe.maybe(node(K, V)), greater::maybe.maybe(node(K, V))).
:- type xbtree(K, V) --->
empty ;
root(node(K, V)).
%-----------------------------------------------------------------------------%
:- pred search_node(node(K, V)::in, K::in, V::out) is semidet.
%-----------------------------------------------------------------------------%
:- pred insert_node(K::in, V::in, node(K, V)::in, node(K, V)::out) is semidet.
%-----------------------------------------------------------------------------%
:- pred update_node(K::in, V::in, node(K, V)::in, node(K, V)::out, bool.bool::out) is semidet.
%-----------------------------------------------------------------------------%
:- pred set_node(K::in, V::in, node(K, V)::in, node(K, V)::out, bool.bool::out) is det.
%-----------------------------------------------------------------------------%
:- func node(K, V) = node(K, V).
:- mode node(in, in) = (out) is det.
:- mode node(out, out) = (in) is det.
%-----------------------------------------------------------------------------%
:- pragma promise_pure(node/2).
node(K::in, V::in) = (node(K, V, maybe.no, maybe.no)::out).
node(K::out, V::out) = (node(K, V, _, _)::in).
%-----------------------------------------------------------------------------%
init = empty.
%-----------------------------------------------------------------------------%
singleton(K, V) = root(node(K, V)).
%-----------------------------------------------------------------------------%
search(root(Node), K, V) :- search_node(Node, K, V).
%-----------------------------------------------------------------------------%
insert(K, V, empty, root(node(K, V))).
insert(K, V, root(In), root(Out)) :-
insert_node(K, V, In, Out).
%-----------------------------------------------------------------------------%
update(K, V, root(In), root(Out)) :-
update_node(K, V, In, Out, _).
%-----------------------------------------------------------------------------%
set(K, V, empty, root(node(K, V))).
set(K, V, root(In), root(Out)) :-
set_node(K, V, In, Out, _).
%-----------------------------------------------------------------------------%
search_node(node(NodeK, NodeV, MaybeLess, MaybeGreater), K, V) :-
builtin.compare(Cmp, K, NodeK),
require_complete_switch [Cmp] (
Cmp = (=),
V = NodeV
;
Cmp = (<),
MaybeLess = maybe.yes(Less),
search_node(Less, K, V)
;
Cmp = (>),
MaybeGreater = maybe.yes(Greater),
search_node(Greater, K, V)
).
%-----------------------------------------------------------------------------%
insert_node(K, V, !Node) :-
!.Node = node(NodeK, _NodeV, MaybeLess, MaybeGreater),
builtin.compare(Cmp, K, NodeK),
(
Cmp = (<),
require_complete_switch [MaybeLess] (
MaybeLess = maybe.no,
NewLess = node(K, V)
;
MaybeLess = maybe.yes(Less),
insert_node(K, V, Less, NewLess)
% TODO: NOT TAIL RECURSIVE!
),
!Node ^ less := maybe.yes(NewLess)
;
Cmp = (>),
require_complete_switch [MaybeGreater] (
MaybeGreater = maybe.no,
NewGreater = node(K, V)
;
MaybeGreater = maybe.yes(Greater),
insert_node(K, V, Greater, NewGreater)
% TODO: NOT TAIL RECURSIVE!
),
!Node ^ greater := maybe.yes(NewGreater)
).
%-----------------------------------------------------------------------------%
update_node(K, V, !Node, Changed) :-
!.Node = node(NodeK, NodeV, MaybeLess, MaybeGreater),
builtin.compare(Cmp, K, NodeK),
require_complete_switch [Cmp] (
Cmp = (=),
( if
NodeV = V
then
Changed = bool.no % No update needed
else
Changed = bool.yes,
!Node ^ v := V
)
;
Cmp = (<),
MaybeLess = maybe.yes(Less),
update_node(K, V, Less, NewLess, Changed),
% TODO: NOT TAIL RECURSIVE!
(
Changed = bool.yes,
!Node ^ less := maybe.yes(NewLess)
;
Changed = bool.no % No update needed
)
;
Cmp = (>),
MaybeGreater = maybe.yes(Greater),
update_node(K, V, Greater, NewGreater, Changed),
% TODO: NOT TAIL RECURSIVE!
(
Changed = bool.yes,
!Node ^ greater := maybe.yes(NewGreater)
;
Changed = bool.no % No update needed
)
).
%-----------------------------------------------------------------------------%
set_node(K, V, !Node, Changed) :-
!.Node = node(NodeK, NodeV, MaybeLess, MaybeGreater),
builtin.compare(Cmp, K, NodeK),
(
Cmp = (=),
( if
NodeV = V
then
Changed = bool.no % No update needed
else
Changed = bool.yes,
!Node ^ v := V
)
;
Cmp = (<),
(
MaybeLess = maybe.yes(Less),
set_node(K, V, Less, NewLess, Changed)
% TODO: NOT TAIL RECURSIVE!
;
MaybeLess = maybe.no,
Changed = bool.yes,
NewLess = node(K, V)
),
(
Changed = bool.yes,
!Node ^ less := maybe.yes(NewLess)
;
Changed = bool.no % No update needed
)
;
Cmp = (>),
(
MaybeGreater = maybe.yes(Greater),
set_node(K, V, Greater, NewGreater, Changed)
% TODO: NOT TAIL RECURSIVE!
;
MaybeGreater = maybe.no,
Changed = bool.yes,
NewGreater = node(K, V)
),
(
Changed = bool.yes,
!Node ^ greater := maybe.yes(NewGreater)
;
Changed = bool.no % No update needed
)
).
INFO