partisan

%% -------------------------------------------------------------------
%%
%% Copyright (c) 2014 Basho Technologies, Inc.  All Rights Reserved.
%%
%% This file is provided to you under the Apache License,
%% Version 2.0 (the "License"); you may not use this file
%% except in compliance with the License.  You may obtain
%% a copy of the License at
%%
%%   http://www.apache.org/licenses/LICENSE-2.0
%%
%% Unless required by applicable law or agreed to in writing,
%% software distributed under the License is distributed on an
%% "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
%% KIND, either express or implied.  See the License for the
%% specific language governing permissions and limitations
%% under the License.
%%
%% -------------------------------------------------------------------

-module(partisan_plumtree_util).


-include("partisan_logger.hrl").

-export([build_tree/3]).
-export([log/2]).
-export([log/3]).


%% -----------------------------------------------------------------------------
%% @doc Convert a list of elements into an N-ary tree. This conversion
%%      works by treating the list as an array-based tree where, for
%%      example in a binary 2-ary tree, a node at index i has children
%%      2i and 2i+1. The conversion also supports a "cycles" mode where
%%      the array is logically wrapped around to ensure leaf nodes also
%%      have children by giving them backedges to other elements.
%% @end
%% -----------------------------------------------------------------------------
-spec build_tree(N :: integer(), Nodes :: [term()], Opts :: [term()]) ->
    orddict:orddict().

build_tree(N, Nodes, Opts) ->
    Expand =
        case lists:member(cycles, Opts) of
            true ->
                lists:flatten(lists:duplicate(N + 1, Nodes));
            false ->
                Nodes
        end,
    {Tree, _} =
        lists:foldl(fun(Elm, {Result, Worklist}) ->
                            Len = erlang:min(N, length(Worklist)),
                            {Children, Rest} = lists:split(Len, Worklist),
                            NewResult = [{Elm, Children} | Result],
                            {NewResult, Rest}
                    end, {[], tl(Expand)}, Nodes),
    orddict:from_list(Tree).


%% -----------------------------------------------------------------------------
%% @doc
%% @end
%% -----------------------------------------------------------------------------
-spec log(debug | info | error, String :: string(), Args :: list(term())) -> ok.

-ifdef(TEST).

log(Level, String) ->
    log(Level, String, []).


log(debug, String, Args) ->
    ?LOG_DEBUG(String, Args);

log(info, String, Args) ->
    ?LOG_INFO(String, Args);

log(error, String, Args) ->
    ?LOG_ERROR(String, Args).

-else.

log(_Level, _String) -> ok.
log(_Level, _String, _Args) -> ok.

-endif.




%% =============================================================================
%% TESTS
%% =============================================================================



-ifdef(TEST).

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

arity_test() ->
    %% 1-ary tree
    ?assertEqual([{node1, []}], orddict:to_list(build_tree(1, [node1], []))),
    ?assertEqual([{node1, [node2]},
                  {node2, []}], orddict:to_list(build_tree(1, [node1, node2], []))),
    ?assertEqual([{node1, [node2]},
                  {node2, [node3]},
                  {node3, []}], orddict:to_list(build_tree(1, [node1, node2, node3], []))),
    ?assertEqual([{node1, [node2]},
                  {node2, [node3]},
                  {node3, [node4]},
                  {node4, []}], orddict:to_list(build_tree(1, [node1, node2,
                                                               node3, node4], []))),

    %% 2-ary tree
    ?assertEqual([{node1, []}], orddict:to_list(build_tree(2, [node1], []))),
    ?assertEqual([{node1, [node2]},
                  {node2, []}], orddict:to_list(build_tree(2, [node1, node2], []))),
    ?assertEqual([{node1, [node2, node3]},
                  {node2, []},
                  {node3, []}], orddict:to_list(build_tree(2, [node1, node2, node3], []))),
    ?assertEqual([{node1, [node2, node3]},
                  {node2, [node4]},
                  {node3, []},
                  {node4, []}], orddict:to_list(build_tree(2, [node1, node2,
                                                               node3, node4], []))),
    ?assertEqual([{node1, [node2, node3]},
                  {node2, [node4, node5]},
                  {node3, []},
                  {node4, []},
                  {node5, []}], orddict:to_list(build_tree(2, [node1, node2,
                                                               node3, node4,
                                                               node5], []))),
    ?assertEqual([{node1, [node2, node3]},
                  {node2, [node4, node5]},
                  {node3, [node6]},
                  {node4, []},
                  {node5, []},
                  {node6, []}], orddict:to_list(build_tree(2, [node1, node2,
                                                               node3, node4,
                                                               node5, node6], []))),

    %% 3-ary tree
    ?assertEqual([{node1, []}], orddict:to_list(build_tree(3, [node1], []))),
    ?assertEqual([{node1, [node2]},
                  {node2, []}], orddict:to_list(build_tree(3, [node1, node2], []))),
    ?assertEqual([{node1, [node2, node3]},
                  {node2, []},
                  {node3, []}], orddict:to_list(build_tree(3, [node1, node2, node3], []))),
    ?assertEqual([{node1, [node2, node3, node4]},
                  {node2, []},
                  {node3, []},
                  {node4, []}], orddict:to_list(build_tree(3, [node1, node2,
                                                               node3, node4], []))),
    ?assertEqual([{node1, [node2, node3, node4]},
                  {node2, [node5]},
                  {node3, []},
                  {node4, []},
                  {node5, []}], orddict:to_list(build_tree(3, [node1, node2,
                                                               node3, node4,
                                                               node5], []))),
    ?assertEqual([{node1, [node2, node3, node4]},
                  {node2, [node5, node6]},
                  {node3, []},
                  {node4, []},
                  {node5, []},
                  {node6, []}], orddict:to_list(build_tree(3, [node1, node2,
                                                               node3, node4,
                                                               node5, node6], []))),
    ?assertEqual([{node1, [node2, node3, node4]},
                  {node2, [node5, node6, node7]},
                  {node3, []},
                  {node4, []},
                  {node5, []},
                  {node6, []},
                  {node7, []}], orddict:to_list(build_tree(3, [node1, node2,
                                                               node3, node4,
                                                               node5, node6,
                                                               node7], []))),
    ?assertEqual([{node1, [node2, node3, node4]},
                  {node2, [node5, node6, node7]},
                  {node3, [node8]},
                  {node4, []},
                  {node5, []},
                  {node6, []},
                  {node7, []},
                  {node8, []}], orddict:to_list(build_tree(3, [node1, node2,
                                                               node3, node4,
                                                               node5, node6,
                                                               node7, node8], []))).

cycles_test() ->
    %% 1-ary tree
    ?assertEqual([{node1, [node1]}], orddict:to_list(build_tree(1, [node1], [cycles]))),
    ?assertEqual([{node1, [node2]},
                  {node2, [node1]}], orddict:to_list(build_tree(1, [node1, node2], [cycles]))),
    ?assertEqual([{node1, [node2]},
                  {node2, [node3]},
                  {node3, [node1]}], orddict:to_list(build_tree(1, [node1, node2, node3], [cycles]))),
    ?assertEqual([{node1, [node2]},
                  {node2, [node3]},
                  {node3, [node4]},
                  {node4, [node1]}], orddict:to_list(build_tree(1, [node1, node2,
                                                                   node3, node4], [cycles]))),

    %% 2-ary tree
    ?assertEqual([{node1, [node1, node1]}], orddict:to_list(build_tree(2, [node1], [cycles]))),
    ?assertEqual([{node1, [node2, node1]},
                  {node2, [node2, node1]}], orddict:to_list(build_tree(2, [node1, node2], [cycles]))),
    ?assertEqual([{node1, [node2, node3]},
                  {node2, [node1, node2]},
                  {node3, [node3, node1]}], orddict:to_list(build_tree(2, [node1, node2, node3], [cycles]))),
    ?assertEqual([{node1, [node2, node3]},
                  {node2, [node4, node1]},
                  {node3, [node2, node3]},
                  {node4, [node4, node1]}], orddict:to_list(build_tree(2, [node1, node2,
                                                                           node3, node4], [cycles]))),
    ?assertEqual([{node1, [node2, node3]},
                  {node2, [node4, node5]},
                  {node3, [node1, node2]},
                  {node4, [node3, node4]},
                  {node5, [node5, node1]}], orddict:to_list(build_tree(2, [node1, node2,
                                                                           node3, node4,
                                                                           node5], [cycles]))),
    ?assertEqual([{node1, [node2, node3]},
                  {node2, [node4, node5]},
                  {node3, [node6, node1]},
                  {node4, [node2, node3]},
                  {node5, [node4, node5]},
                  {node6, [node6, node1]}], orddict:to_list(build_tree(2, [node1, node2,
                                                               node3, node4,
                                                               node5, node6], [cycles]))),

    %% 3-ary tree
    ?assertEqual([{node1, [node1, node1, node1]}], orddict:to_list(build_tree(3, [node1], [cycles]))),
    ?assertEqual([{node1, [node2, node1, node2]},
                  {node2, [node1, node2, node1]}], orddict:to_list(build_tree(3, [node1, node2], [cycles]))),
    ?assertEqual([{node1, [node2, node3, node1]},
                  {node2, [node2, node3, node1]},
                  {node3, [node2, node3, node1]}], orddict:to_list(build_tree(3, [node1, node2, node3], [cycles]))),
    ?assertEqual([{node1, [node2, node3, node4]},
                  {node2, [node1, node2, node3]},
                  {node3, [node4, node1, node2]},
                  {node4, [node3, node4, node1]}], orddict:to_list(build_tree(3, [node1, node2,
                                                                                  node3, node4], [cycles]))),
    ?assertEqual([{node1, [node2, node3, node4]},
                  {node2, [node5, node1, node2]},
                  {node3, [node3, node4, node5]},
                  {node4, [node1, node2, node3]},
                  {node5, [node4, node5, node1]}], orddict:to_list(build_tree(3, [node1, node2,
                                                                                  node3, node4,
                                                                                  node5], [cycles]))),
    ?assertEqual([{node1, [node2, node3, node4]},
                  {node2, [node5, node6, node1]},
                  {node3, [node2, node3, node4]},
                  {node4, [node5, node6, node1]},
                  {node5, [node2, node3, node4]},
                  {node6, [node5, node6, node1]}], orddict:to_list(build_tree(3, [node1, node2,
                                                                                  node3, node4,
                                                                                  node5, node6], [cycles]))).
-endif.