%%% Copyright (c) 2014, NORDUnet A/S.
%%% See LICENSE for licensing information.
%%%

%% tree storage, TODO: docdoc

-module(ts).
-include_lib("eunit/include/eunit.hrl").
-export_type([tree_store/0]).
-export([new/0, add/3, delete/2, retrieve/2, count/2]).

%% #tree_store{} has one member, layers, holding an array of arrays
%% with binaries, keyed on layer.
-record(tree_store, {layers :: array:array(array:array(binary()))}).
-type tree_store() :: #tree_store{}.

%%%%%%%%%%%%%%%%%%%%
%% Public.
new() ->
    #tree_store{layers = array:new()}.

-spec add(tree_store(), non_neg_integer(), binary()) -> tree_store().
add(S = #tree_store{layers = Layers}, Layer, Entry) ->
    {NewLayers, Array} = layer_rw(Layers, Layer),
    NewArray = array:set(array:size(Array), Entry, Array),
    S#tree_store{layers = array:set(Layer, NewArray, NewLayers)}.

-spec delete(tree_store(), non_neg_integer()) -> tree_store().
delete(S = #tree_store{layers = Layers}, Layer) ->
    Array = layer_ro(Layers, Layer),
    NewArray = array:resize(array:size(Array) - 1, Array),
    S#tree_store{layers = array:set(Layer, NewArray, Layers)}.

-spec retrieve(tree_store(), tuple()) -> binary() | undefined.
retrieve(#tree_store{layers = Layers}, {Layer, Index}) ->
    Array = layer_ro(Layers, Layer),
    Len = array:size(Array),
    case Index < Len of
        true -> array:get(Index, Array);
        false -> undefined
    end.

-spec count(tree_store(), non_neg_integer()) -> non_neg_integer().
count(#tree_store{layers = Layers}, Layer) ->
    array:size(layer_ro(Layers, Layer)).

%%%%%%%%%%%%%%%%%%%%
%% Private.
-spec layer_ro(array:array(array:array(binary())), non_neg_integer()) ->
                      array:array(binary).
layer_ro(Layers, Layer) ->
    case array:get(Layer, Layers) of
        undefined -> array:new();
        Array -> Array
    end.

-spec layer_rw(array:array(array:array(binary())), non_neg_integer()) ->
                      {array:array(), array:array(binary)}.
layer_rw(Layers, Layer) ->
    case array:get(Layer, Layers) of
        undefined -> {array:set(Layer, array:new(), Layers), array:new()};
        Array -> {Layers, Array}
    end.

%%%%%%%%%%%%%%%%%%%%
%% Testing ts.
add_retrieve_test_() ->
    S = new(),
    S0 = add(S, 0, <<00>>),
    S1 = add(S0, 0, <<01>>),
    S2 = add(S1, 1, <<10>>),
    [?_assertEqual(2, count(S2, 0)),
     ?_assertEqual(1, count(S2, 1)),
     ?_assertEqual(<<10>>, retrieve(S2, {1,0})),
     ?_assertEqual(<<00>>, retrieve(S2, {0,0})),
     ?_assertEqual(<<01>>, retrieve(S2, {0,1})),
     ?_assertEqual(<<00>>, retrieve(S2, {0,0})),
     ?_assertEqual(undefined, retrieve(S2, {3,0})),
     ?_assertEqual(undefined, retrieve(S1, {1,0}))].

delete_test_() ->
    S = new(),
    S0 = add(S, 0, <<00>>),
    S1 = add(S0, 0, <<01>>),
    S2 = add(S1, 1, <<10>>),
    S3 = delete(S2, 0),
    S4 = add(S3, 0, <<99>>),
    [?_assertEqual(2, count(S2, 0)),
     ?_assertEqual(1, count(S3, 0)),
     ?_assertEqual(2, count(S4, 0)),
     ?_assertEqual(retrieve(S2, {0,1}), <<01>>),
     ?_assertEqual(retrieve(S3, {0,1}), undefined),
     ?_assertEqual(retrieve(S3, {0,0}), <<00>>),
     ?_assertEqual(retrieve(S4, {0,1}), <<99>>)].