Add support for retrieving historical tree heads.

1 files changed, 69 insertions, 33 deletions

diff --git a/src/ht.erl b/src/ht.erl
index 042b969..a51924d 100644
--- a/src/ht.erl
+++ b/src/ht.erl
@@ -24,7 +24,6 @@
 
 -export([size/0, add/1, tree_hash/0, tree_hash/1]).
 -export([get_incl/2, get_cons/2]).
--export([mktree/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]).
@@ -95,14 +94,56 @@ handle_call({consistency, _Version1, _Version2}, _From, State) ->
 -spec get_hash(tree(), non_neg_integer(), tuple()) -> {tree(), binary()}.
 get_hash(Tree, Version, IR) ->
     NewTree = update(Tree, Version),
-    Hash = ts:retrieve_hash(NewTree#tree.store, IR), % TODO: Honour version!
+    Hash = ts:retrieve_hash(NewTree#tree.store, IR),
     {NewTree, Hash}.
 
+%% FIXME: rename to tree_head or maybe just head?
 -spec tree_hash(tree(), integer()) -> {tree(), binary()}.
 tree_hash(Tree, -1) ->
     {Tree, hash(<<"">>)};
+tree_hash(Tree = #tree{version = V}, Version) when Version == V ->
+    get_hash(Tree, Version, {0, depth(Tree) - 1});
+tree_hash(Tree = #tree{version = V}, Version) when Version > V ->
+    {Tree, enotimetravel};
 tree_hash(Tree, Version) ->
-    get_hash(Tree, Version, {0, depth(Tree) - 1}).
+    old_version_tree_head(update(Tree, Version), Version).
+
+-spec old_version_tree_head(tree(), non_neg_integer()) -> {tree(), binary()}.
+old_version_tree_head(Tree, Version) ->
+    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),
+
+    %% 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.
+    {NewTree, LeftHash} = get_hash(Tree, Version, {FirstLeftI, FirstLeftR}),
+    last_right_node_rehash(NewTree, Version, FirstLeftR, FirstLeftI, LeftHash).
+
+-spec last_right_node_rehash(tree(), non_neg_integer(), non_neg_integer(),
+                             non_neg_integer(), binary()) -> {tree(), binary()}.
+last_right_node_rehash(Tree, _Version, _Layer, 0, RightNodeHash) ->
+    {Tree, RightNodeHash};
+last_right_node_rehash(Tree, Version, Layer, Index, RightNodeHash) ->
+    {NewTree, Hash} =
+        case right_node_p(Index) of
+            true ->
+                {T2, LHash} = get_hash(Tree, Version, {Index - 1, Layer}),
+                {T2, mkinnerhash(LHash, RightNodeHash)};
+            false ->
+                {Tree, RightNodeHash}
+        end,
+    last_right_node_rehash(NewTree, Version, Layer + 1, parent(Index), Hash).
+
+-spec first_left_node(non_neg_integer(), non_neg_integer()) ->
+                             {non_neg_integer(), non_neg_integer()}.
+first_left_node(Layer, Index) ->
+    case right_node_p(Index) of
+        true -> first_left_node(Layer + 1, parent(Index));
+        false -> {Layer, Index}
+    end.
 
 %% @doc Add an entry but don't update the tree.
 -spec add(tree(), binary()) -> tree().
@@ -119,7 +160,7 @@ new() ->
           evaluated = -1,
           store = ts:new()}.
 
--spec new(non_neg_integer()) -> tree().
+-spec new([non_neg_integer()]) -> tree().
 new([Version]) when is_integer(Version) ->
     foldl(fun(#mtl{entry = E}, Tree) ->
                   D = (E#timestamped_entry.entry)#plop_entry.data,
@@ -130,7 +171,7 @@ new([List]) when is_list(List) ->
                   add(Tree, D) % Return value -> Tree in next invocation.
           end, new(), List).
 
-%% @doc Calculate hashes in Tree up to and including leaf with index
+%% @doc Calculate hashes in Tree up to and including node with index
 %% equal to Version. Update Tree.evaluated to reflect the new state.
 -spec update(tree(), non_neg_integer()) -> tree().
 update(Tree, 0) ->
@@ -167,12 +208,12 @@ update_parent(S, Layer, I, ILast) when I >= ILast ->
     %% to where its parent would've been were it a right child. This
     %% is a "placeholder node" which simplifies when creating
     %% incomplete ("non-frozen") trees.
-    case right_leaf_p(ILast) of
+    case right_node_p(ILast) of
         true -> S;
         _ -> ts:append(S, Layer + 1, ts:retrieve_hash(S, {ILast, Layer}))
     end;
 update_parent(S, Layer, I, ILast) ->
-    false = right_leaf_p(I),                    % ASSERTION
+    false = right_node_p(I),                    % ASSERTION
     %% Make an inner node hash of I and sibling. Store it as
     %% parent. Recurse with next pair of leaves.
     update_parent(ts:store(S, {parent(I), Layer + 1},
@@ -180,11 +221,12 @@ update_parent(S, Layer, I, ILast) ->
                               ts:retrieve_hash(S, {I + 1, Layer}))),
                   Layer, I + 2, ILast).
 
+%% @doc Parent of {i, r} is at {i/2, r+1} (unless it's a "placeholder").
 parent(I) ->
     I bsr 1.
 
--spec right_leaf_p(integer()) -> boolean().
-right_leaf_p(Index) ->
+-spec right_node_p(integer()) -> boolean().
+right_node_p(Index) ->
     case Index band 1 of
         1 -> true;
         _ -> false
@@ -237,34 +279,31 @@ hash(Data) ->
          "ddb89be403809e325750d3d263cd78929c2942b7942a34b77e122c9594a74c8c",
          "5dc9da79a70659a9ad559cb701ded9a2ab9d823aad2f4960cfe370eff4604328"]).
 
+%% FIXME: Don't start and stop the server manually all the time. EUnit
+%% can help.
+test_init(L) ->
+    stop(),
+    {ok, _Pid} = start_link(L).
+
 %% @doc Build tree using add/2 and mth/2 and compare the resulting
 %% tree hashes.
-%% FIXME: Don't start and stop the server here!
+%% FIXME: Move outside.
 add_test() ->
     lists:foreach(
       fun(X) -> L = lists:sublist(?TEST_VECTOR_LEAVES, X),
-                {ok, _Pid} = start_link(L),
-                mktree(L),
-                ?assertEqual(mth(L), tree_hash()),
-                stop() end,
+                test_init(L),
+                ?assertEqual(mth(L), tree_hash()) end,
       random_entries(length(?TEST_VECTOR_LEAVES))).
 
-random_entries(N) ->
-    [V || {_, V} <- lists:sort(
-                      [{random:uniform(N), E} || E <- lists:seq(1, N)])].
-
-update_test_OFF() ->
-    mktree(?TEST_VECTOR_LEAVES),
+%% FIXME: Move outside.
+old_versions_test() ->
+    test_init(?TEST_VECTOR_LEAVES),
+    ?assertEqual(mth(?TEST_VECTOR_LEAVES), tree_hash()),
     lists:foreach(
-      fun(Test) ->
-              L = lists:sublist(?TEST_VECTOR_LEAVES, Test),
-              ?assertEqual(mth(L),
-                           tree_hash(lists:nth(Test, ?TEST_VECTOR_HASHES))) end,
+      fun(X) -> ?assertEqual(mth(lists:sublist(?TEST_VECTOR_LEAVES, X)),
+                             tree_hash(X - 1)) end,
       random_entries(length(?TEST_VECTOR_LEAVES))).
 
-%% verify_old_versions_test() ->
-%%     {ok, _Pid} = start_link(db:size() - 1).
-
 %%%%%%%%%%%%%%%%%%%%
 %% Testing the test helpers.
 mth_test() ->
@@ -278,12 +317,9 @@ mth_test() ->
 %%%%%%%%%%%%%%%%%%%%
 %% Test helpers.
 
-mktree(L) ->
-    mktree(new(), L).
-mktree(Tree, []) ->
-    Tree;
-mktree(Tree, [H|T]) ->
-    mktree(add(Tree, H), T).
+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. Implements
 %% the algorithm in section 2.1 of RFC 6962. Used for testing.