radix.erl

-module(radix).

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

-compile(export_all).

-record(tree, {value :: maybe(),
               children :: children()}).
-type maybe() :: error | {ok, term()}.
-type tree() :: #tree{}.
-type children() :: orddict:orddict().

-spec empty() -> tree().
empty() ->
    #tree{value=error,
          children=[]}.

-spec find(string(), tree()) -> maybe().
find([], #tree{value=Value}) ->
    Value;
find(String, #tree{children=Children}) ->
    case find_prefix(String, Children) of
        {ok, {Prefix, Tree}} ->
            Tail = lists:nthtail(length(Prefix), String),
            find(Tail, Tree);
        error ->
            error
    end.

find_prefix(String, Orddict) ->
    orddict:fold(find_prefix_helper(String), error, Orddict).

find_prefix_helper(String) ->
    fun(Key, Value, Acc) ->
            case lists:prefix(Key, String) of
                true ->
                    {ok, {Key, Value}};
                false ->
                    Acc
            end
    end.

-spec store(string(), term(), tree()) -> tree().
store([], Value, Tree) ->
    Tree#tree{value={ok, Value}};
store(String, Value, Tree=#tree{children=[]}) ->
    Tree#tree{children=[{String, #tree{value={ok, Value},
                                       children=[]}}]};
store(String, Value, Tree=#tree{children=Children}) ->
    {PrefixToStore, SubTreePrefix, OldVal, NewTree} = case find_splittable(String, Children) of
        {insert, {Prefix, OldTree}} ->
            Tail = lists:nthtail(length(Prefix), String),
            {Prefix, Tail, OldTree, Tree};
        {split, {TooLong, OldTree}} ->
            io:format("splitting~n"),
            Tail = lists:nthtail(length(String), TooLong),
            {ok, ValueFromTree} = OldTree#tree.value,
            SplitTree = store(Tail, ValueFromTree, empty()),
            {String, "", SplitTree, Tree#tree{children=orddict:erase(TooLong, Tree#tree.children)}};
        error ->
            {String, "", empty(), Tree}
    end,
    NewSubTree = store(SubTreePrefix, Value, OldVal),
    NewTree#tree{children=orddict:store(PrefixToStore,
                                        NewSubTree,
                                        Children)}.

find_splittable(String, Children) ->
    orddict:fold(find_split_helper(String), error, Children).

find_split_helper(String) ->
    fun(Key, Value, Acc) ->
            case lists:prefix(Key, String) of
                true ->
                    {insert, {Key, Value}};
                false ->
                    case lists:prefix(String, Key) of
                        true ->
                            {split, {Key, Value}};
                        false ->
                            Acc
                    end
            end
    end.


-spec to_list(tree()) -> [{string(), term()}].
%% @doc Return a key-sorted list. Uses a pre-order
%% traversal.
to_list(Tree) ->
    lists:flatten(to_list("", Tree)).

%% Right now we use body-recursion. Could likely rewrite to be
%% tail recursive and reverse the accumulator at the end, but this
%% is working fine for now.
to_list(Prefix, #tree{value={ok, Value},
                      children=Children}) ->
    [{Prefix, Value}] ++ list_children(Prefix, Children);
to_list(Prefix, #tree{value=error,
                      children=Children}) ->
    list_children(Prefix, Children).

list_children(Prefix, Children) ->
    [to_list(Prefix ++ C, V) || {C, V} <- Children].

%% Testing

test() ->
    test(500).

test(NumTimes) ->
    eqc:quickcheck(eqc:numtests(NumTimes, prop_to_list())).

input_list() ->
    list({list(int()), int()}).

prop_to_list() ->
    ?FORALL(Xs, input_list(),
            equiv_to_orddict(Xs)).

equiv_to_orddict(Xs) ->
    to_list(from_list(Xs)) =:= orddict:from_list(Xs).

from_list(L) ->
    lists:foldl(fun insert_fun/2, empty(), L).

insert_fun({Key, Value}, Acc) ->
    store(Key, Value, Acc).