%%% Copyright (c) 2014-2015, NORDUnet A/S.
%%% See LICENSE for licensing information.
%%%
%%% Implementation of a history tree as described in Efficient Data
%%% Structures for Tamper-Evident Logging [0]. This implementation
%%% follows RFC 6962 and differs from [0] only in how non-full trees
%%% are handled.
%%%
%%% The idea with placeholder nodes on incomplete layers is borrowed
%%% from Emilia Käspers implementation in Google's
%%% certificate-transparency project.
%%%
%%% Hashes of inner nodes and leaves are stored in whatever
%%% datastructure the ts module uses, one per layer. Layer 0 is where
%%% the leaves are. The total number of datastructures used in ts is
%%% equal to the depth of the tree. The depth of the tree is
%%% ceil(lg2(number of leaves)).
%%%
%%% Let {r,i} denote the hash with index i on layer r. The first leaf
%%% is {0,0}, second is {0,1} and n:th is {0,n-1}.
%%% The parent of {r,i} is {r+1,floor(i/2)} (not strictly true because
%%% of "placeholder nodes", see update_parent/4).
%%% The sibling of {r,i} is {r,i+1} when i is even and {r,i-1} when i
%%% is odd.
%%%
%%% [0] https://www.usenix.org/event/sec09/tech/full_papers/crosby.pdf

-module(ht).
-behaviour(gen_server).

-export([reset_tree/1, load_tree/1, size/0, leaf_hash/1]).
-export([add/1, root/0, root/1, path/2, consistency/2, dump/1]).
-export([start_link/0, start_link/1, stop/0]).
-export([init/1, handle_call/3, terminate/2, handle_cast/2, handle_info/2,
         code_change/3]).
-export([testing_get_state/0, print_tree/0, print_tree/1]).

-import(stacktrace, [call/2, call/3]).
-include_lib("eunit/include/eunit.hrl").
-import(lists, [foreach/2, foldl/3, reverse/1]).
-include("timeouts.hrl").

-define(MAX_READ_ENTRIES, 10000).

%% Data types.
-record(tree, {version :: integer(),
               evaluated :: integer(),
               store :: ts:tree_store()}).
-type tree() :: #tree{}.

%%%%%%%%%%%%%%%%%%%%
%% Public interface.
start_link() ->
    gen_server:start_link({local, ?MODULE}, ?MODULE, [db:size() - 1], []).
start_link(NEntries) ->
    gen_server:start_link({local, ?MODULE}, ?MODULE, [NEntries], []).
reset_tree(Arg) ->
    call(?MODULE, {reset_tree, Arg}, infinity).
load_tree(Version) ->
    case call(?MODULE, {load_tree, Version}, ?HT_LOAD_TREE_TIMEOUT) of
        eagain ->
            load_tree(Version);
        Result ->
            Result
    end.
stop() ->
    call(?MODULE, stop).
size() ->
    call(?MODULE, size).
add(Hash) ->
    call(?MODULE, {add, Hash}).
root() ->
    case call(?MODULE, root) of
        eagain ->
            root();
        Result ->
            Result
    end.
root(Version) ->
    case call(?MODULE, {root, Version}) of
        eagain ->
            root(Version);
        Result ->
            Result
    end.
path(I, V) ->
    call(?MODULE, {path, I, V}).
consistency(V1, V2) ->
    call(?MODULE, {consistency, V1, V2}).
leaf_hash(Data) ->
    mkleafhash(Data).
%% Testing and debugging.
testing_get_state() ->
    call(?MODULE, testing_get_state).
print_tree() ->
    call(?MODULE, {print_tree, 4}).
print_tree(HashOutputLen) ->
    call(?MODULE, {print_tree, HashOutputLen}).
dump(Filename) ->
    call(?MODULE, {dump, Filename}).

%% gen_server callbacks
init(Args) ->
    lager:info("reading and building tree"),
    Tree = new_and_update(Args),
    lager:info("finished"),
    {ok, Tree}.
handle_cast(_Request, State) ->
    {noreply, State}.
handle_info(_Info, State) ->
    {noreply, State}.
code_change(_OldVersion, State, _Extra) ->
    {ok, State}.
terminate(_Reason, _State) ->
    ok.

% public api
handle_call({reset_tree, Arg}, _From, _State) ->
    NewTree = new_and_update(Arg),
    {reply, NewTree, NewTree};
handle_call({load_tree, Version}, _From, State) ->
    {Reply, NewTree} = read_new_entries(State, Version),
    {reply, Reply, NewTree};
handle_call(stop, _From, State) ->
    {stop, normal, stopped, State};
handle_call(size, _From, State) ->
    {reply, State#tree.version + 1, State};
handle_call({add, Hash}, _From, State) ->
    {reply, ok, add_with_update(State, Hash)};
handle_call(root, _From, State) ->
    {NewState, Hash} = head(State, State#tree.version),
    {reply, Hash, NewState};
handle_call({root, Version}, _From, State) ->
    {NewState, Hash} = head(State, Version),
    {reply, Hash, NewState};
handle_call({path, Index, Version}, _From, State) ->
    {NewState, Path} = path(State, Index, Version),
    {reply, Path, NewState};
handle_call({consistency, Version1, Version2}, _From, State) ->
    {NewState, ConsProof} = consistency(State, Version1, Version2),
    {reply, ConsProof, NewState};
% testing and debugging
handle_call(testing_get_state, _From, State) ->
    {reply, State, State};
handle_call({print_tree, HashOutputLen}, _From, State) ->
    {reply, print_tree(State, HashOutputLen), State};
handle_call({dump, Filename}, _From, State) ->
    {reply, dump(State, Filename), State}.

%%%%%%%%%%%%%%%%%%%%
%% Private.

-spec consistency(tree(), non_neg_integer(), non_neg_integer()) -> {tree(), list()}.
consistency(Tree, -1, _V2) ->
    {Tree, []};
consistency(Tree, V1, V2) when V1 >= V2 ->
    {Tree, []};
consistency(Tree, _V1, V2) when V2 > Tree#tree.version ->
    {Tree, []};
consistency(Tree, V1, V2) ->
    %% Walk up the tree from V1 to first left child.
    {StartLayer, StartIndex} = first_left_node(0, V1),
    First = case StartIndex of
                0 -> [];
                _ -> [get_hash(Tree, {StartLayer, StartIndex})]
            end,
    %% Get path from first left child to head of V2.
    {_, Path} = path(Tree, StartLayer, StartIndex, V2),
    {Tree, First ++ Path}.

%% @doc Return a list of hashes showing the path from leaf Index to
%% the tree head in the tree of version Version.
-spec path(tree(), non_neg_integer(), non_neg_integer()) -> {tree(), list()}.
path(Tree, Index, Version) ->
    path(Tree, 0, Index, Version).

-spec path(tree(), non_neg_integer(), non_neg_integer(), non_neg_integer()) ->
                  {tree(), list()}.
path(Tree, _Layer, _Index, -1) ->
    {Tree, []};
path(Tree = #tree{version = V}, _, _, Version) when Version > V ->
    {Tree, []};                              % FIXME: Return an error?
path(Tree, Layer, Index, Version) ->
    %% The magic here is to tell path/6 to stop at Version >> Layer.
    {Tree, path(Tree, Layer, Index, Version bsr Layer, Version, [])}.

%% @doc Return path from {Layer,I} to head of tree Version. I is the
%% leftmost and ILast the rightmost node to consider, at Layer.
-spec path(tree(), non_neg_integer(), non_neg_integer(), non_neg_integer(), non_neg_integer(), list()) -> list().
path(_, _, _, 0, _, Acc) ->
    reverse(Acc);
path(Tree, Layer, I, ILast, Version, Acc) ->
    path(Tree, Layer + 1, parent(I), parent(ILast), Version,
         case sibling(I) of
             Sib when Sib == ILast ->
                 %% We're at the edge of the layer and might need to
                 %% recompute an old tree.
                 [old_version_tree_head(Tree, Version, Layer) | Acc];
             Sib when Sib < ILast ->
                 %% Just use sibling.
                 [get_hash(Tree, {Layer, Sib}) | Acc];
             _ ->
                 %% Sibling is larger than ILast so doesn't exist.
                 Acc
         end).

-spec get_hash(tree(), tuple()) -> binary().
get_hash(Tree, {R, I}) ->
    true = Tree#tree.evaluated >= I,            % ASSERTION
    ts:retrieve(Tree#tree.store, {R, I}).

-spec head(tree(), integer()) -> {tree(), binary()}.
head(Tree, -1) ->
    {Tree, hash(<<"">>)};
head(Tree = #tree{version = V}, Version) when Version == V ->
    {Tree, get_hash(Tree, {depth(Tree) - 1, 0})};
head(Tree = #tree{version = V}, Version) when Version > V ->
    {Tree, enotimetravel};
head(Tree, Version) ->
    {Tree, old_version_tree_head(Tree, Version)}.

-spec old_version_tree_head(tree(), non_neg_integer()) -> binary().
old_version_tree_head(Tree, Version) ->
    old_version_tree_head(Tree, Version, -1).

-spec old_version_tree_head(tree(), non_neg_integer(), integer()) -> binary().
old_version_tree_head(Tree, Version, BreakAtLayer) ->
    true = Tree#tree.evaluated >= Version,      % ASSERTION
    %% Go up the tree from the rightmost leaf (index=Version) until a
    %% left node is found. (There is always one -- the head is a left
    %% node.)
    {FirstLeftR, FirstLeftI} = first_left_node(0, Version, BreakAtLayer),

    %% Walk up the tree from this lowest left node up to and including
    %% the last right node, rehashing as we go. Calculate the parent
    %% hash of that node and its sibling. Return that hash.
    last_right_node_rehash(Tree, Version, FirstLeftR, FirstLeftI,
                           get_hash(Tree, {FirstLeftR, FirstLeftI}),
                           BreakAtLayer).

-spec last_right_node_rehash(tree(), non_neg_integer(), non_neg_integer(),
                             non_neg_integer(), binary(), integer()) ->
                                    binary().
last_right_node_rehash(_, _, Layer, _, RightNodeHash, BAL) when Layer == BAL ->
    %% Bailing out at Layer.
    RightNodeHash;
last_right_node_rehash(_, _, _, 0, RightNodeHash, _) ->
    %% Index is 0, we're done.
    RightNodeHash;
last_right_node_rehash(Tree, Version, Layer, Index, RightNodeHash, BAL) ->
    last_right_node_rehash(
      Tree, Version, Layer + 1, parent(Index),
      case right_node_p(Index) of
          true ->
              %% Rehash parent using sibling.
              mkinnerhash(get_hash(Tree, {Layer, Index - 1}), RightNodeHash);
          false ->
              %% Just use the incoming hash.
              RightNodeHash
      end,
      BAL).

-spec first_left_node(non_neg_integer(), non_neg_integer()) ->
                             {non_neg_integer(), non_neg_integer()}.
first_left_node(Layer, Index) ->
    first_left_node(Layer, Index, -1).

-spec first_left_node(non_neg_integer(), non_neg_integer(), integer()) ->
                             {non_neg_integer(), non_neg_integer()}.
first_left_node(Layer, Index, BAL) when Layer == BAL ->
    {Layer, Index};
first_left_node(Layer, Index, BAL) ->
    case right_node_p(Index) of
        true -> first_left_node(Layer + 1, parent(Index), BAL);
        false -> {Layer, Index}
    end.

%% @doc Add a hash and update the tree.
-spec add_with_update(tree(), binary()) -> tree().
add_with_update(Tree, Hash) ->
    update(add(Tree, Hash)).

%% @doc Add a hash but don't update the tree.
-spec add(tree(), binary()) -> tree().
add(Tree = #tree{version = V, store = Store}, Hash) ->
    Tree#tree{version = V + 1, store = ts:add(Store, 0, Hash)}.

read_new_entries(Tree, Version) when is_integer(Version) ->
    EndBound = min(Version, Tree#tree.version + ?MAX_READ_ENTRIES),
    NewEntries = db:get_by_indices(Tree#tree.version + 1, EndBound, {sorted, true}),
    NewHashes = [H || {_I, H, _E} <- NewEntries],
    NewTree = foldl(fun(H, T) ->
                            add(T, H)
                    end, Tree, NewHashes),
    UpdatedTree = update(NewTree),
    case EndBound of
        Version ->
            {ok, UpdatedTree};
        _ ->
            {eagain, UpdatedTree}
    end.

%% @doc Return a new updated tree.
-spec new_and_update(list()) -> tree().
new_and_update(List) ->
    update(new(List)).

%% @doc Return a new tree, not updated.
-spec new(list()) -> tree().
new([]) ->
    #tree{version = -1,
          evaluated = -1,
          store = ts:new()};
new([-1]) ->
    new([]);
%% Initialise tree from db.
new([Version]) when is_integer(Version) ->
    foldl(fun(Hash, Tree) ->
                  %% Return value becomes Tree in next invocation.
                  add(Tree, Hash)
          end, new([]), [H || {_I, H, _E} <-
                                  db:get_by_indices(0, Version, {sorted, true})]);
%% Initialise tree from List with hashes.
new([List]) when is_list(List) ->
    foldl(fun(Hash, Tree) -> add(Tree, Hash) end,
          new([]), List).

%% @doc Calculate hashes in Tree. Update Tree.evaluated to reflect the
%% new state.
-spec update(tree()) -> tree().
update(Tree = #tree{version = Version, evaluated = E}) when E >= Version ->
    %% Evaluated enough already. Nothing to do.
    true = E == Version,                        % ASSERTION.
    Tree;
update(Tree = #tree{version = Version, evaluated = Evaluated}) ->
    NewTree = update_layer(Tree, 0, Evaluated + 1, Version),
    NewTree#tree{evaluated = Version}.

%% @doc Update the tree wrt the leaves ICur..ILast.
-spec update_layer(tree(), non_neg_integer(), non_neg_integer(),
                   non_neg_integer()) -> tree().
update_layer(Tree, _Layer, _ICur, 0) ->         % Done
    Tree;
update_layer(Tree = #tree{store = S}, Layer, ICur, ILast) ->
    %% Before updating parent layer, delete potential placeholder.
    Store = case ts:count(S, Layer + 1) == parent(ICur) + 1 of
                true -> ts:delete(S, Layer + 1);
                false -> S
            end,

    %% Update parents on next upper layer, starting with a left
    %% child <= ICur and ending with ILast. Recurse with next layer.
    NewStore = update_parent(Store, Layer,
                             strip_bits_bottom(ICur, 1), ILast),
    update_layer(Tree#tree{store = NewStore}, Layer + 1,
                 parent(ICur), parent(ILast)).

%% @doc Update parents of I..ILast, on Layer+1. I has to be a left child.
-spec update_parent(ts:tree_store(), non_neg_integer(), non_neg_integer(),
                    non_neg_integer()) -> ts:tree_store().
update_parent(S, Layer, I, ILast) when I >= ILast ->
    %% We're done updating parents. If ILast is a left child, copy it
    %% to where its parent would've been were it a right child. This
    %% is a "placeholder node" which simplifies creating incomplete
    %% ("non-frozen") trees.
    case right_node_p(ILast) of
        true -> S;
        _ -> ts:add(S, Layer + 1, ts:retrieve(S, {Layer, ILast}))
    end;
update_parent(S, Layer, I, ILast) ->
    false = right_node_p(I),                    % ASSERTION
    %% Make an inner node hash of I and its sibling. Store it as
    %% parent. Recurse with next pair of leaves.
    update_parent(ts:add(S, Layer + 1,
                  mkinnerhash(ts:retrieve(S, {Layer, I}),
                              ts:retrieve(S, {Layer, I + 1}))),
                  Layer, I + 2, ILast).

%% @doc Parent of {r, i} is at {r+1, i/2} (unless it's a placeholder).
parent(I) ->
    I bsr 1.

-spec right_node_p(integer()) -> boolean().
right_node_p(Index) ->
    case Index band 1 of
        1 -> true;
        _ -> false
    end.

-spec sibling(non_neg_integer()) -> non_neg_integer().
sibling(Index) ->
    case right_node_p(Index) of
        true -> Index - 1;
        false -> Index + 1
    end.

strip_bits_bottom(N, Nbits) ->
    (N bsr Nbits) bsl Nbits.

%% @doc Return position of highest bit set, counting from the least
%% significant bit, starting at 1.
bitpos_first_set(N) ->
    L = [Bit || <<Bit:1>> <= binary:encode_unsigned(N)],
    length(L) - ffs(L, 0).
ffs([], Acc) ->
    Acc;
ffs([H|T], Acc) ->
    case H of
        0 -> ffs(T, Acc + 1);
        _ -> Acc
    end.

depth(#tree{version = -1}) ->
    0;
depth(#tree{version = V}) ->
    bitpos_first_set(V) + 1.

-spec mkleafhash(binary()) -> binary().
mkleafhash(Data) ->
    hash([<<"\x00">>, Data]).

-spec mkinnerhash(binary(), binary()) -> binary().
mkinnerhash(Hash1, Hash2) ->
    hash([<<"\x01">>, Hash1, Hash2]).

-spec hash(binary()) -> binary() | iolist().
hash(Data) ->
    crypto:hash(sha256, Data).

%%%%%%%%%%%%%%%%%%%%
%% Debugging helpers.
print_tree(Tree, HashOutputLen) ->
    print_tree(update(Tree), HashOutputLen, 0, Tree#tree.version, depth(Tree)).

print_tree(_, _, _, _, 0) ->
    ok;
print_tree(Tree, HashOutputLen, Layer, ILast, LayersLeft) ->
    print_layer(Tree, HashOutputLen, Layer, ILast),
    print_tree(Tree, HashOutputLen, Layer + 1, ILast bsr 1, LayersLeft - 1).

print_layer(Tree, HashOutputLen, Layer, ILast) ->
    foreach(
      fun(I) -> io:format(
                  "~s ", [string:substr(
                            hex:bin_to_hexstr(get_hash(Tree, {Layer, I})),
                            1, HashOutputLen)])
      end,
      lists:seq(0, ILast)),
    io:format("~n").

dump(Tree, Filename) ->
    {ok, File} = file:open(Filename, [read, write, binary]),
    ok = file:write(File, "# evaluated " ++ integer_to_list(Tree#tree.evaluated) ++ "\n"),
    lists:foreach(fun (R) ->
                          write_layer_header(File, R),
                          lists:foreach(fun (I) ->
                                                Entry = ts:retrieve(Tree#tree.store, {R, I}),
                                                dumpone(File, Entry)
                                        end, lists:seq(0, ts:count(Tree#tree.store, R) - 1))
                  end, lists:seq(0, depth(Tree) - 1)),
    file:close(File).

write_layer_header(File, Layer) ->
    EntryText = "# layer " ++ integer_to_list(Layer) ++ "\n",
    ok = file:write(File, EntryText).

dumpone(File, Entry) ->
    EntryText = hex:bin_to_hexstr(Entry) ++ "\n",
    ok = file:write(File, EntryText).

%%%%%%%%%%%%%%%%%%%%
%% Testing ht.

-define(TEST_VECTOR_LEAVES,
        ["", "\x00", "\x10", " !", "01", "@ABC", "PQRSTUVW", "`abcdefghijklmno"]).

%% FIXME: Don't start and stop the server manually all the time. EUnit
%% can help!
test_init(L) ->
    stop(),
    {ok, _Pid} = start_link(L).

%% Using internal interfaces.
storage_consistency_test() ->
    LEAVES = [<<X:32>> || X <- lists:seq(0, 64)],
    test_init(lists:map(fun mkleafhash/1, LEAVES)),
    Tree = testing_get_state(),
    ?assert(check_consistency(Tree)),
    Tree2 = update(Tree),
    ?assert(check_consistency(Tree2)),
    Tree3 = add_with_update(Tree2, <<>>),
    ?assert(check_consistency(Tree3)).

%%%% Test helpers, internal interfaces.
levelsize(0, 1) ->
    0;
levelsize(LastIndex, Level) ->
    trunc((LastIndex)/math:pow(2, Level))+1.

deptheval(#tree{evaluated = -1}) ->
    0;
deptheval(#tree{evaluated = V}) ->
    bitpos_first_set(V) + 1.

%% check_consistency(#tree{evaluated = Evaluated}) when Evaluated == -1 ->
%%     true;
check_consistency(Tree) ->
    check_consistency_tree(Tree) and check_consistency_storage(Tree).

check_consistency_tree(Tree) ->
    Tree#tree.evaluated == Tree#tree.version.

check_consistency_storage(Tree) ->
    lists:all(fun (R) ->
                      ActualSize = ts:count(Tree#tree.store, R),
                      ShouldBeSize = levelsize(Tree#tree.evaluated, R),
                      case ActualSize of
                          ShouldBeSize ->
                              true;
                          _ ->
                              lager:error("lastindex: ~p level ~p actual size: ~p should be: ~p", [Tree#tree.evaluated, R, ActualSize, ShouldBeSize]),
                              false
                      end
              end, lists:seq(1, deptheval(Tree) - 1)).

%% Using only exported functions.
%% TODO: Move all these tests to a separate file in ../test. They're
%% only using external functions.

%% @doc Verify trees using add/2.
add_test() ->
    lists:foreach(
      fun(X) -> L = lists:sublist(?TEST_VECTOR_LEAVES, X),
                test_init(lists:map(fun mkleafhash/1, L)),
                ?assertEqual(mth(L), root()) end,
      random_entries(length(?TEST_VECTOR_LEAVES))).

old_versions_test() ->
    test_init(lists:map(fun mkleafhash/1, (?TEST_VECTOR_LEAVES))),
    ?assertEqual(mth(?TEST_VECTOR_LEAVES), root()),
    lists:foreach(
      fun(X) -> ?assertEqual(mth(lists:sublist(?TEST_VECTOR_LEAVES, X)),
                             root(X - 1)) end,
      random_entries(length(?TEST_VECTOR_LEAVES))).

old_versions_bigger_test() ->
    LEAVES = [<<X:32>> || X <- lists:seq(0, 64)], % 1024 is not unreasonable
    test_init(lists:map(fun mkleafhash/1, LEAVES)),
    ?assertEqual(mth(LEAVES), root()),
    lists:foreach(
      fun(X) -> ?assertEqual(mth(lists:sublist(LEAVES, X)),
                             root(X - 1)) end,
      random_entries(length(LEAVES))).

%% Test vector from Googles C++ implementation, "Generated from
%% ReferenceMerklePath."
-define(TEST_VECTOR_PATHS,
        %% {leaf_index+1, version+1, path}
        [{0, 0, []},
         {1, 1, []},
         {1, 8,
          ["96a296d224f285c67bee93c30f8a309157f0daa35dc5b87e410b78630a09cfc7",
           "5f083f0a1a33ca076a95279832580db3e0ef4584bdff1f54c8a360f50de3031e",
           "6b47aaf29ee3c2af9af889bc1fb9254dabd31177f16232dd6aab035ca39bf6e4"]},
         {6, 8,
          ["bc1a0643b12e4d2d7c77918f44e0f4f79a838b6cf9ec5b5c283e1f4d88599e6b",
           "ca854ea128ed050b41b35ffc1b87b8eb2bde461e9e3b5596ece6b9d5975a0ae0",
           "d37ee418976dd95753c1c73862b9398fa2a2cf9b4ff0fdfe8b30cd95209614b7"]},
         {3, 3,
          ["fac54203e7cc696cf0dfcb42c92a1d9dbaf70ad9e621f4bd8d98662f00e3c125"]},
         {2, 5,
          ["6e340b9cffb37a989ca544e6bb780a2c78901d3fb33738768511a30617afa01d",
           "5f083f0a1a33ca076a95279832580db3e0ef4584bdff1f54c8a360f50de3031e",
           "bc1a0643b12e4d2d7c77918f44e0f4f79a838b6cf9ec5b5c283e1f4d88599e6b"]}]).

%% @doc Test paths on a single version 7 tree.
path_test() ->
    test_init(lists:map(fun mkleafhash/1, ?TEST_VECTOR_LEAVES)),
    foreach(
      fun(N) ->
              Test = lists:nth(N, ?TEST_VECTOR_PATHS),
              ?assertEqual(
                 path_ref(element(1, Test) - 1,
                          lists:sublist(?TEST_VECTOR_LEAVES, element(2, Test))),
                 path(element(1, Test) - 1, element(2, Test) - 1))
      end,
      lists:seq(1, length(?TEST_VECTOR_PATHS))).

%% @doc Test path on minimal sized trees.
path_inc_test() ->
    foreach(
      fun(N) ->
              Test = lists:nth(N, ?TEST_VECTOR_PATHS),
              Leaves = lists:sublist(?TEST_VECTOR_LEAVES, element(2, Test)),
              test_init(lists:map(fun mkleafhash/1, Leaves)),
              ?assertEqual(
                 path_ref(element(1, Test) - 1, Leaves),
                 path(element(1, Test) - 1, element(2, Test) - 1))
      end,
      lists:seq(1, length(?TEST_VECTOR_PATHS))).

-define(TEST_VECTOR_PROOFS,
        [
         %% Test vectors from Googles C++ implementation, "Generated
         %% from ReferenceSnapshotConsistency."
         {1, 1, []},
         {1, 8,
          ["96a296d224f285c67bee93c30f8a309157f0daa35dc5b87e410b78630a09cfc7",
           "5f083f0a1a33ca076a95279832580db3e0ef4584bdff1f54c8a360f50de3031e",
           "6b47aaf29ee3c2af9af889bc1fb9254dabd31177f16232dd6aab035ca39bf6e4"]},
         {6, 8,
          ["0ebc5d3437fbe2db158b9f126a1d118e308181031d0a949f8dededebc558ef6a",
           "ca854ea128ed050b41b35ffc1b87b8eb2bde461e9e3b5596ece6b9d5975a0ae0",
           "d37ee418976dd95753c1c73862b9398fa2a2cf9b4ff0fdfe8b30cd95209614b7"]},
         {2, 5,
          ["5f083f0a1a33ca076a95279832580db3e0ef4584bdff1f54c8a360f50de3031e",
           "bc1a0643b12e4d2d7c77918f44e0f4f79a838b6cf9ec5b5c283e1f4d88599e6b"]},
         %% RFC6962 section 2.1.3.
         {3, 7,
          ["0298D122906DCFC10892CB53A73992FC5B9F493EA4C9BADB27B791B4127A7FE7",
           "07506A85FD9DD2F120EB694F86011E5BB4662E5C415A62917033D4A9624487E7",
           "FAC54203E7CC696CF0DFCB42C92A1D9DBAF70AD9E621F4BD8D98662F00E3C125",
           "837DBB152E9B079010717E84E865DA4EBC0FA198A806D59D31BF15ACCEF22D0E"]},
         {4, 7,
          ["837DBB152E9B079010717E84E865DA4EBC0FA198A806D59D31BF15ACCEF22D0E"]},
         {6, 7,
          ["0EBC5D3437FBE2DB158B9F126A1D118E308181031D0A949F8DEDEDEBC558EF6A",
           "B08693EC2E721597130641E8211E7EEDCCB4C26413963EEE6C1E2ED16FFB1A5F",
           "D37EE418976DD95753C1C73862B9398FA2A2CF9B4FF0FDFE8B30CD95209614B7"]}
        ]).

%% @doc Test proofs on a single version 7 tree.
consistency_test() ->
    test_init(lists:map(fun mkleafhash/1, ?TEST_VECTOR_LEAVES)),
    foreach(
      fun(N) ->
              Test = lists:nth(N, ?TEST_VECTOR_PROOFS),
              ?assertEqual(
                 consistency_proof_ref(
                   %% element(1, Test) - 1,
                   %% lists:sublist(?TEST_VECTOR_LEAVES, element(2, Test))),
                   Test),                       % FIXME
                 consistency(element(1, Test) - 1, element(2, Test) - 1))
      end,
      lists:seq(1, length(?TEST_VECTOR_PROOFS))).

%% FIXME: implement and move
consistency_proof_ref(Test) ->
    [hex:hexstr_to_bin(X2) || X2 <- element(3, Test)].

%%%%%%%%%%%%%%%%%%%%
%% Test helpers.
random_entries(N) ->
    [V || {_, V} <- lists:sort(
                      [{random:uniform(N), E} || E <- lists:seq(1, N)])].

%% @doc Return the Merkle Tree Head for the leaves in L. Reference
%% implementation for testing. Implements the algorithm in section 2.1
%% of RFC 6962.
-spec mth(list()) -> binary().
mth([]) ->
    hash(<<"">>);
mth([E]) ->
    hash([<<"\x00">>, E]);
mth(L) ->
    Split = 1 bsl (bitpos_first_set(length(L) - 1) - 1),
    {L1, L2} = lists:split(Split, L),
    hash([<<"\x01">>, mth(L1), mth(L2)]).

%% @doc Return the Merkle Audit Path from I to the root of the tree
%% with leaves L. Reference implementation for testing. Implements the
%% algorithm in section 2.1.1 of RFC 6962.
-spec path_ref(non_neg_integer(), list()) -> list().
path_ref(I, _) when I < 0 ->
    [];
path_ref(I, L) when I >= length(L) ->
    [];
path_ref(0, [_]) ->
    [];
path_ref(I, L) ->
    Split = 1 bsl (bitpos_first_set(length(L) - 1) - 1),
    {L1, L2} = lists:split(Split, L),
    case I of
        I when I < Split ->
            path_ref(I, L1) ++ [mth(L2)];
        _ ->
            path_ref(I - Split, L2) ++ [mth(L1)]
    end.

%%%%%%%%%%%%%%%%%%%%
%% Testing the test helpers. It's turtles all the way down.
-define(TEST_VECTOR_HASHES,
        ["6e340b9cffb37a989ca544e6bb780a2c78901d3fb33738768511a30617afa01d",
         "fac54203e7cc696cf0dfcb42c92a1d9dbaf70ad9e621f4bd8d98662f00e3c125",
         "aeb6bcfe274b70a14fb067a5e5578264db0fa9b51af5e0ba159158f329e06e77",
         "d37ee418976dd95753c1c73862b9398fa2a2cf9b4ff0fdfe8b30cd95209614b7",
         "4e3bbb1f7b478dcfe71fb631631519a3bca12c9aefca1612bfce4c13a86264d4",
         "76e67dadbcdf1e10e1b74ddc608abd2f98dfb16fbce75277b5232a127f2087ef",
         "ddb89be403809e325750d3d263cd78929c2942b7942a34b77e122c9594a74c8c",
         "5dc9da79a70659a9ad559cb701ded9a2ab9d823aad2f4960cfe370eff4604328"]).
mth_test() ->
    lists:foreach(
      fun(X) -> ?assertEqual(
		   hex:hexstr_to_bin(lists:nth(X, ?TEST_VECTOR_HASHES)),
                   mth(lists:sublist(?TEST_VECTOR_LEAVES, X)))
      end,
      lists:seq(1, length(?TEST_VECTOR_LEAVES))).

path_ref_test() ->
    foreach(
      fun(N) ->
              Test = lists:nth(N, ?TEST_VECTOR_PATHS),
              ?assertEqual(
                 [hex:hexstr_to_bin(X) || X <- element(3, Test)],
                 path_ref(element(1, Test) - 1,
                          lists:sublist(?TEST_VECTOR_LEAVES, element(2, Test))))
      end,
      lists:seq(1, length(?TEST_VECTOR_PATHS))).