1-module(vtree_test_util).
2
3-include_lib("../include/vtree.hrl").
4
5-export([generate_kvnodes/1, generate_kpnodes/1, create_file/1,
6         get_kvnodes/2]).
7
8
9-spec generate_kvnodes(Num :: pos_integer()) -> [#kv_node{}].
10generate_kvnodes(Num) ->
11    [random_kvnode(I) || I <- lists:seq(1, Num)].
12
13-spec random_kvnode(I :: pos_integer()) -> #kv_node{}.
14random_kvnode(I) ->
15    Max = 1000,
16    {A, B, C, D, E, F, G, H} = {
17      rand:uniform(Max), rand:uniform(Max), rand:uniform(Max),
18      rand:uniform(Max), rand:uniform(Max), rand:uniform(Max),
19      rand:uniform(Max), rand:uniform(Max)},
20    Mbb = [
21           {erlang:min(A, B), erlang:max(A, B)},
22           {erlang:min(C, D), erlang:max(C, D)},
23           {erlang:min(E, F), erlang:max(E, F)},
24           {erlang:min(G, H), erlang:max(G, H)}],
25    LineString = {[{<<"type">>, <<"LineString">>},
26                   {<<"coordinates">>,
27                    [[A, C, E, G],
28                     [(A+B)/2, (C+D)/2, (E+F)/2, (G+H)/2],
29                     [B, D, F, H]]
30                    }]},
31    {ok, BinLineString} = wkb_writer:geojson_to_wkb(LineString),
32    Id = list_to_binary("Node" ++ integer_to_list(I)),
33    Value = list_to_binary("Value" ++ integer_to_list(I)),
34    #kv_node {
35               key = Mbb,
36               docid = Id,
37               geometry = BinLineString,
38               body = Value
39             }.
40
41
42-spec generate_kpnodes(Num :: pos_integer()) -> [#kp_node{}].
43generate_kpnodes(Num) ->
44    [random_kpnode(I) || I <- lists:seq(1, Num)].
45
46-spec random_kpnode(I :: pos_integer()) -> #kp_node{}.
47random_kpnode(I) ->
48    Max = 1000,
49    {A, B, C, D, E, F, G, H} = {
50      rand:uniform(Max), rand:uniform(Max), rand:uniform(Max),
51      rand:uniform(Max), rand:uniform(Max), rand:uniform(Max),
52      rand:uniform(Max), rand:uniform(Max)},
53    Mbb = [
54           {erlang:min(A, B), erlang:max(A, B)},
55           {erlang:min(C, D), erlang:max(C, D)},
56           {erlang:min(E, F), erlang:max(E, F)},
57           {erlang:min(G, H), erlang:max(G, H)}],
58    #kp_node {
59               key = Mbb,
60               childpointer = I,
61               treesize = rand:uniform(Max),
62               mbb_orig = Mbb
63             }.
64
65
66create_file(Filename) ->
67    {ok, Fd} = case couch_file:open(Filename, [create, overwrite]) of
68    {ok, Fd2} ->
69        {ok, Fd2};
70    {error, Reason} ->
71        io:format("ERROR (~s): Couldn't open file (~s) for tree storage~n",
72                  [Reason, Filename])
73    end,
74    Fd.
75
76
77% Return a 2-tuple with a list of the depths of the KV-nodes and the
78% KV-nodes themselves
79get_kvnodes(Fd, RootPos) ->
80    Children = vtree_io:read_node(Fd, RootPos),
81    get_kvnodes(Fd, Children, 0, {[], []}).
82get_kvnodes(_Fd, [], _Depth, Acc) ->
83    Acc;
84get_kvnodes(_Fd, [#kv_node{}|_]=Children, Depth, {Depths, Nodes}) ->
85    {[Depth|Depths], Children ++ Nodes};
86get_kvnodes(Fd, [#kp_node{}=Node|Rest], Depth, Acc) ->
87    Children = vtree_io:read_node(Fd, Node#kp_node.childpointer),
88    % Move down
89    Acc2 = get_kvnodes(Fd, Children, Depth+1, Acc),
90    % Move sideways
91    get_kvnodes(Fd, Rest, Depth, Acc2).
92