Binary tree example

tree.erl

-module(tree).

-compile({no_auto_import,[{size, 1}]}).

%% API
-export([
    new/0, new/1,
    insert/2,
    search/2,
    to_list/1,
    size/1
]).

-define(NODE(Left, Data, Right), {Left, Data, Right}).
-define(KV(Key, Value), {Key, Value}).

new() ->
    ?NODE(undefined, undefined, undefined).

new(KV = {_, _}) ->
    ?NODE(undefined, KV, undefined).

insert(undefined, KV = {_Key, _Value}) ->
    new(KV);
insert(Node = ?NODE(_, undefined, _), KV = {_Key, _Value}) ->
    setelement(2, Node, KV);
insert(Node = ?NODE(Left, {K, _}, _), KV = {Key, _Value}) when K > Key ->
    setelement(1, Node, insert(Left, KV));
insert(Node = ?NODE(_, {Key, _}, _), KV = {Key, _Value}) ->
    setelement(2, Node, KV);
insert(Node = ?NODE(_, {K, _}, Right), KV = {Key, _Value}) when K < Key ->
    setelement(3, Node, insert(Right, KV)).

search(undefined, _Key) ->
    false;
search(?NODE(_, KV = {Key, _}, _), Key) ->
    KV;
search(?NODE(Left, {K, _}, _), Key) when K > Key ->
    search(Left, Key);
search(?NODE(_, {K, _}, Right), Key) when K < Key ->
    search(Right, Key).

to_list(Tree) ->
    to_list_(Tree, []).

to_list_(undefined, Acc) ->
    Acc;
to_list_(?NODE(_, undefined, _), Acc) ->
    Acc;
to_list_(?NODE(Left, KV, Right), Acc) ->
    to_list_(Left, [KV | to_list_(Right, Acc)]).

size(undefined) ->
    0;
size(?NODE(_, undefined, _)) ->
    0;
size(?NODE(Left, _, Right)) ->
    1 + size(Left) + size(Right).

-ifdef(TEST).

-include_lib("eunit/include/eunit.hrl").

eleven_number_sequence_test() ->
    TenNums = lists:seq(-5, 5),
    InsertFun =
        fun(N, Tree) ->
            insert(Tree, {N, $a+N})
        end,
    Expected = {undefined,{-5,92},{undefined,{-4,93},{undefined,{-3,94},{undefined,{-2,95},{undefined,{-1,96},{undefined,{0,97},{undefined,{1,98},{undefined,{2,99},{undefined,{3,100},{undefined,{4,101},{undefined,{5,102},undefined}}}}}}}}}}},
    ?assertEqual(Expected, lists:foldl(InsertFun, new(), TenNums)).

search_test() ->
    Tree = {undefined,{-5,92},{undefined,{-4,93},{undefined,{-3,94},{undefined,{-2,95},{undefined,{-1,96},{undefined,{0,97},{undefined,{1,98},{undefined,{2,99},{undefined,{3,100},{undefined,{4,101},{undefined,{5,102},undefined}}}}}}}}}}},
    {Time, Value} = timer:tc(fun search/2, [Tree, 3]),
    ?debugFmt("Took ~bus to search tree for key~n", [Time]),
    ?assertEqual({3, 100}, Value).

random_generation_test_() ->
    {timeout, 60, fun() ->
        MaxElems = 1000000,
        Elements = [begin
            <<Key:32>> = rand:bytes(4),
            {Key, N}
         end || N <- lists:seq(1, MaxElems)],
        Tree = lists:foldl(fun(KV, Tree) -> insert(Tree, KV) end, new(), Elements),
        {Key, _} = lists:last(to_list(Tree)),
        {SizeTime, Size} = timer:tc(fun size/1, [Tree]),
        ?debugFmt("Key to look for is ~p, tree is ~b large (calculated in ~bus)~n", [Key, Size, SizeTime]),
        {SearchTime, Value} = timer:tc(fun search/2, [Tree, Key]),
        ?debugFmt("Took ~bus to search tree for key~n", [SearchTime]),
        ?assertMatch({Key, _}, Value)
    end}.

-endif.