denrei

%%%-------------------------------------------------------------------
%%% @private
%%% @doc
%%% Prefix tree for storing subject to socket mappings
%%% @end
%%%-------------------------------------------------------------------
-module(denrei_tree).

%% API
-export([new/0, insert/3, delete/3, fetch/2, match/2]).

-export_type([tree/0, keys/0, values/0]).

-type tree() :: {trees(), values()}.
-type trees() :: [{key(), tree()}].
-type values() :: [any()].
-type keys() :: [key()].
-type key() :: atom().

-define(EMPTY_TREE, {[], []}).

-spec new() -> tree().
new() ->
    ?EMPTY_TREE.

-spec insert(any(), keys(), tree()) -> tree().
insert(Value, [], {Trees, Values}) ->
    {Trees, lists:usort([Value | Values])};
insert(Value, [Key | Keys], {Trees, Values}) ->
    NewTrees = case lists:keyfind(Key, 1, Trees) of
                   {Key, Tree} ->
                       lists:keyreplace(Key, 1, Trees, {Key, insert(Value, Keys, Tree)});
                   false ->
                       lists:ukeysort(1, [{Key, insert(Value, Keys, new())} | Trees])
               end,
    {NewTrees, Values}.

-spec delete(any(), keys() | all, tree()) -> tree().
delete(Value, [], {Trees, Values}) ->
    {Trees, delete_value(Value, Values)};
delete(Value, all, {Trees, Values}) ->
    NewValues = delete_value(Value, Values),
    NewTrees = lists:foldr(fun({Key, Tree}, AccIn) ->
        case delete(Value, all, Tree) of
            ?EMPTY_TREE ->
                AccIn;
            NewTree ->
                [{Key, NewTree} | AccIn]
        end
    end, [], Trees),
    {NewTrees, NewValues};
delete(Value, [Key | Keys], {Trees, Values}) ->
    NewTrees = case lists:keyfind(Key, 1, Trees) of
                   {Key, Tree} ->
                       case delete(Value, Keys, Tree) of
                           ?EMPTY_TREE ->
                               lists:keydelete(Key, 1, Trees);
                           NewTree ->
                               lists:keyreplace(Key, 1, Trees, {Key, NewTree})
                       end;
                   false ->
                       Trees
               end,
    {NewTrees, Values}.

-spec delete_value(any() | all, values()) -> values().
delete_value(all, _Values) ->
    [];
delete_value(Value, Values) ->
    lists:delete(Value, Values).

-spec fetch(keys() | all, tree()) -> false | values().
fetch(all, Tree) ->
    case fetch_all(Tree) of
        [] ->
            false;
        Values ->
            Values
    end;
fetch([], {_, Values}) ->
    Values;
fetch([Key | Keys], {Trees, _}) ->
    case lists:keyfind(Key, 1, Trees) of
        false ->
            false;
        {Key, Tree} ->
            fetch(Keys, Tree)
    end.

-spec fetch_all(tree()) -> values().
fetch_all({Trees, Values}) ->
    lists:umerge([Values | [fetch_all(Tree) || {_, Tree} <- Trees]]).

-spec match(keys(), tree()) -> values().
match([], {_, Values}) ->
    Values;
match([Key | Keys], {Trees, _}) ->
    lists:foldl(fun
        ({values, Values}, AccIn) ->
            lists:umerge(AccIn, Values);
        ({match, Tree}, AccIn) ->
            lists:umerge(AccIn, match(Keys, Tree))
    end, [], keymatch(Key, Trees)).

-spec keymatch(key(), trees()) -> [{atom(), tree()}].
keymatch(_, []) ->
    [];
keymatch(Key, [{'*', {[], Values}} | Trees]) ->
    [{values, Values} | keymatch(Key, Trees)];
keymatch(Key, [{'?', Tree} | Trees]) ->
    [{match, Tree} | keymatch(Key, Trees)];
keymatch(Key, [{Key, Tree} | Trees]) ->
    [{match, Tree} | keymatch(Key, Trees)];
keymatch(Key, [_ | Trees]) ->
    keymatch(Key, Trees).

%% ===================================================================
%%  Tests
%% ===================================================================

-ifdef(TEST).

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

setup_tree() ->
    Tree0 = insert(s1, ['a', 'b', '*'], new()),
    Tree1 = insert(s0, ['a', '?', 'c'], Tree0),
    Tree2 = insert(s2, ['a', 'b', 'c'], Tree1),
    Tree3 = insert(s3, ['a', 'b', '?'], Tree2),
    Tree4 = insert(s4, ['a', 'd', 'c'], Tree3),
    Tree5 = insert(s5, ['a', '?', 'c', '?', 'e'], Tree4),
    Tree6 = insert(s6, ['a', '*'], Tree5),
    Tree7 = insert(s7, ['a', 'b', 'c'], Tree6),
    Tree8 = insert(s5, ['a', '*'], Tree7),
    Tree8.

fetch_test(Tree) ->
    [
        ?_assertEqual([s5, s6], fetch(['a', '*'], Tree)),
        ?_assertEqual([s3], fetch(['a', 'b', '?'], Tree)),
        ?_assertEqual([s2, s7], fetch(['a', 'b', 'c'], Tree)),
        ?_assertEqual(false, fetch(['a', 'b', 'z'], Tree)),
        ?_assertEqual([s0, s1, s2, s3, s4, s5, s6, s7], fetch(all, Tree))
    ].

match_test(Tree) ->
    [
        ?_assertEqual([s5, s6], match(['a', 'z', 'd'], Tree)),
        ?_assertEqual([s1, s3, s5, s6], match(['a', 'b', 'z'], Tree)),
        ?_assertEqual([s0, s1, s2, s3, s5, s6, s7], match(['a', 'b', 'c'], Tree)),
        ?_assertEqual([s1, s5, s6], match(['a', 'b', 'c', 'd', 'e'], Tree))
    ].

delete_test(Tree) ->
    Tree0 = delete(s3, ['a', 'b', '?'], Tree),
    Tree1 = delete(all, ['a', 'b', 'c'], Tree0),
    Tree2 = delete(s5, all, Tree0),
    Tree3 = delete(s2, ['a', 'z', 'z'], Tree2),
    [
        ?_assertEqual([s5, s6], match(['a', 'z'], Tree0)),
        ?_assertEqual([s1, s5, s6], match(['a', 'b', 'z'], Tree0)),
        ?_assertEqual([s0, s1, s2, s5, s6, s7], match(['a', 'b', 'c'], Tree0)),
        ?_assertEqual([s1, s5, s6], match(['a', 'b', 'c', 'd', 'e'], Tree0)),
        ?_assertEqual([s0, s1, s5, s6], match(['a', 'b', 'c'], Tree1)),
        ?_assertEqual([s6], fetch(['a', '*'], Tree2)),
        ?_assertEqual([s6], match(['a', 'z'], Tree2)),
        ?_assertEqual([s1, s6], match(['a', 'b', 'z'], Tree2)),
        ?_assertEqual([s0, s1, s2, s6, s7], match(['a', 'b', 'c'], Tree2)),
        ?_assertEqual([s1, s6], match(['a', 'b', 'c', 'd', 'e'], Tree2)),
        ?_assertEqual(Tree2, Tree3)
    ].

tree_test_() ->
    {setup,
        fun setup_tree/0,
        fun(Tree) ->
            [fetch_test(Tree), match_test(Tree), delete_test(Tree)]
        end
    }.

-endif.