1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
|
%% The contents of this file are subject to the Mozilla Public License
%% Version 1.1 (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.mozilla.org/MPL/
%%
%% Software distributed under the License is distributed on an "AS IS"
%% basis, WITHOUT WARRANTY OF ANY KIND, either express or implied. See
%% the License for the specific language governing rights and
%% limitations under the License.
%%
%% The Original Code is RabbitMQ.
%%
%% The Initial Developer of the Original Code is VMware, Inc.
%% Copyright (c) 2007-2012 VMware, Inc. All rights reserved.
%%
%% A dual-index tree. Entries have a primary key and a set of
%% secondary keys. Most operations require the supply of just the
%% secondary key. Entries exists while they have a non-empty secondar
%% key set. The 'take' operations return the entries that got removed,
%% i.e. that had no remaining secondary keys. take/3 expects entries
%% to exist with the supplied primary keys and secondary key. take/2
%% can cope with the supplied secondary key having no entries.
-module(dtree).
-export([empty/0, insert/4, take/3, take/2,
is_defined/2, is_empty/1, smallest/1, size/1]).
%%----------------------------------------------------------------------------
-ifdef(use_specs).
-export_type([?MODULE/0]).
-opaque(?MODULE() :: {gb_tree(), gb_tree()}).
-type(pk() :: any()).
-type(sk() :: any()).
-type(val() :: any()).
-type(kv() :: {pk(), val()}).
-spec(empty/0 :: () -> ?MODULE()).
-spec(insert/4 :: (pk(), [sk()], val(), ?MODULE()) -> ?MODULE()).
-spec(take/3 :: ([pk()], sk(), ?MODULE()) -> {[kv()], ?MODULE()}).
-spec(take/2 :: (sk(), ?MODULE()) -> {[kv()], ?MODULE()}).
-spec(is_defined/2 :: (sk(), ?MODULE()) -> boolean()).
-spec(is_empty/1 :: (?MODULE()) -> boolean()).
-spec(smallest/1 :: (?MODULE()) -> kv()).
-spec(size/1 :: (?MODULE()) -> non_neg_integer()).
-endif.
%%----------------------------------------------------------------------------
empty() -> {gb_trees:empty(), gb_trees:empty()}.
insert(PK, SKs, V, {P, S}) ->
{gb_trees:insert(PK, {gb_sets:from_list(SKs), V}, P),
lists:foldl(fun (SK, S0) ->
case gb_trees:lookup(SK, S0) of
{value, PKS} -> PKS1 = gb_sets:insert(PK, PKS),
gb_trees:update(SK, PKS1, S0);
none -> PKS = gb_sets:singleton(PK),
gb_trees:insert(SK, PKS, S0)
end
end, S, SKs)}.
take(PKs, SK, {P, S}) ->
{KVs, P1} = take2(PKs, SK, P),
PKS = gb_sets:difference(gb_trees:get(SK, S), gb_sets:from_list(PKs)),
{KVs, {P1, case gb_sets:is_empty(PKS) of
true -> gb_trees:delete(SK, S);
false -> gb_trees:update(SK, PKS, S)
end}}.
take(SK, {P, S}) ->
case gb_trees:lookup(SK, S) of
none -> {[], {P, S}};
{value, PKS} -> {KVs, P1} = take2(gb_sets:to_list(PKS), SK, P),
{KVs, {P1, gb_trees:delete(SK, S)}}
end.
is_defined(SK, {_P, S}) -> gb_trees:is_defined(SK, S).
is_empty({P, _S}) -> gb_trees:is_empty(P).
smallest({P, _S}) -> {K, {_SKS, V}} = gb_trees:smallest(P),
{K, V}.
size({P, _S}) -> gb_trees:size(P).
%%----------------------------------------------------------------------------
take2(PKs, SK, P) ->
lists:foldl(fun (PK, {KVs, P0}) ->
{SKS, V} = gb_trees:get(PK, P0),
SKS1 = gb_sets:delete(SK, SKS),
case gb_sets:is_empty(SKS1) of
true -> {[{PK, V} | KVs], gb_trees:delete(PK, P0)};
false -> {KVs, gb_trees:update(PK, {SKS1, V}, P0)}
end
end, {[], P}, PKs).
|
