trie.erl

-module(trie).

%% Compile Options
-compile({no_auto_import, [get/2, put/2]}).

%% API
-export([get/2,
         new/0,
         put/2,
         put/3]).

%% Trie Structure
-record(trie, {prefix = <<>>, value, children}).


%% =============================================================================
%% API Functions
%% =============================================================================

get(Key, T = #trie{prefix = Prefix}) ->
    Size = binary:longest_common_prefix([Prefix, Key]),
    if
        Size == byte_size(Prefix) ->
            <<_:Size/bytes, Rest/binary>> = Key,
            get_child(Rest, T);
        true ->
            undefined
    end.

new() ->
    #trie{}.

put(Pairs, T = #trie{prefix = Prefix}) when is_list(Pairs) ->
    Size = binary:longest_common_prefix([Prefix|keys(Pairs)]),
    NewT = if
               Size < byte_size(Prefix) ->
                   <<NewPrefix:Size/bytes, C, Key/binary>> = Prefix,
                   init_trie(NewPrefix, C, T#trie{prefix = Key});
               true ->
                   T
           end,
    lists:foldl(fun put_child/2, NewT, strip_prefix(Size, Pairs)).

put(Key, Val, T) ->
    put([{Key, Val}], T).


%% =============================================================================
%% Helper Functions
%% =============================================================================

get_child(<<C, Key/binary>>, #trie{children = Children}) ->
    get_child(C + 1, Key, Children);
get_child(<<>>, #trie{value = Val}) ->
    Val.

get_child(_, _, undefined) ->
    undefined;
get_child(I, Key, Children) ->
    case element(I, Children) of
        #trie{} = T -> get(Key, T);
        undefined   -> undefined
    end.

put_child({<<C, Key/binary>>, Val}, T = #trie{children = Children}) ->
    T#trie{children = put_child(C + 1, Key, Val, Children)};
put_child({<<>>, Val}, T) ->
    T#trie{value = Val}.

put_child(I, Key, Val, undefined) ->
    init_children(I, init_trie(Key, Val));
put_child(I, Key, Val, Children) ->
    Trie = case element(I, Children) of
               #trie{} = T -> put([{Key, Val}], T);
               undefined   -> init_trie(Key, Val)
           end,
    setelement(I, Children, Trie).

keys(Pairs) ->
    lists:map(fun ({K, _}) -> K end, Pairs).

init_children(I, Node) ->
    erlang:make_tuple(256, undefined, [{I, Node}]).

init_trie(Prefix, Value) ->
    #trie{prefix = Prefix, value = Value}.

init_trie(Prefix, C, Node) ->
    #trie{prefix = Prefix, children = init_children(C + 1, Node)}.

strip_prefix(Size, Pairs) ->
    Fun = fun ({K, Val}) ->
                  <<_:Size/bytes, Key/binary>> = K,
                  {Key, Val}
          end,
    lists:map(Fun, Pairs).