1%%% vi:ts=4 sw=4 et
2%%% Copyright (c) 2008 Robert Virding. All rights reserved.
3%%%
4%%% Redistribution and use in source and binary forms, with or without
5%%% modification, are permitted provided that the following conditions
6%%% are met:
7%%%
8%%% 1. Redistributions of source code must retain the above copyright
9%%%    notice, this list of conditions and the following disclaimer.
10%%% 2. Redistributions in binary form must reproduce the above copyright
11%%%    notice, this list of conditions and the following disclaimer in the
12%%%    documentation and/or other materials provided with the distribution.
13%%%
14%%% THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
15%%% "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
16%%% LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
17%%% FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
18%%% COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
19%%% INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
20%%% BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
21%%% LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
22%%% CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
23%%% LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
24%%% ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
25%%% POSSIBILITY OF SUCH DAMAGE.
26%%%-------------------------------------------------------------------
27%%% @copyright 2008 Robert Verding
28%%%
29%%% @doc
30%%%
31%%% Rbdict implements a Key - Value dictionary. An rbdict is a
32%%% representation of a dictionary, where a red-black tree is used to
33%%% store the keys and values.
34%%%
35%%% This module implents exactly the same interface as the module
36%%% ec_dictionary but with a defined representation. One difference is
37%%% that while dict considers two keys as different if they do not
38%%% match (=:=), this module considers two keys as different if and
39%%% only if they do not compare equal (==).
40%%%
41%%% The algorithms here are taken directly from Okasaki and Rbset
42%%% in ML/Scheme. The interface is compatible with the standard dict
43%%% interface.
44%%%
45%%% The following structures are used to build the the RB-dict:
46%%%
47%%% {r,Left,Key,Val,Right}
48%%% {b,Left,Key,Val,Right}
49%%% empty
50%%%
51%%% It is interesting to note that expanding out the first argument of
52%%% l/rbalance, the colour, in store etc. is actually slower than not
53%%% doing it. Measured.
54%%%
55%%% see ec_dictionary
56%%% @end
57%%%-------------------------------------------------------------------
58-module(ec_rbdict).
59
60-behaviour(ec_dictionary).
61
62%% Standard interface.
63-export([add/3, from_list/1, get/2, get/3, has_key/2,
64         has_value/2, new/0, remove/2, size/1, to_list/1,
65         keys/1]).
66
67-export_type([dictionary/2]).
68
69%%%===================================================================
70%%% Types
71%%%===================================================================
72%% This should be opaque, but that kills dialyzer so for now we export it
73%% however you should not rely on the internal representation here
74-type dictionary(K, V) :: empty | {color(),
75                                   dictionary(K, V),
76                                   ec_dictionary:key(K),
77                                   ec_dictionary:value(V),
78                                   dictionary(K, V)}.
79
80-type color() :: r | b.
81
82%%%===================================================================
83%%% API
84%%%===================================================================
85
86-spec new() -> dictionary(_K, _V).
87new() -> empty.
88
89-spec has_key(ec_dictionary:key(K), dictionary(K, _V)) -> boolean().
90has_key(_, empty) ->
91    false;
92has_key(K, {_, Left, K1, _, _}) when K < K1 ->
93    has_key(K, Left);
94has_key(K, {_, _, K1, _, Right}) when K > K1 ->
95    has_key(K, Right);
96has_key(_, {_, _, _, _, _}) ->
97    true.
98
99-spec get(ec_dictionary:key(K), dictionary(K, V)) -> ec_dictionary:value(V).
100get(_, empty) ->
101    throw(not_found);
102get(K, {_, Left, K1, _, _}) when K < K1 ->
103    get(K, Left);
104get(K, {_, _, K1, _, Right}) when K > K1 ->
105    get(K, Right);
106get(_, {_, _, _, Val, _}) ->
107    Val.
108
109-spec get(ec_dictionary:key(K),
110          ec_dictionary:value(V),
111          dictionary(K, V)) -> ec_dictionary:value(V).
112get(_, Default, empty) ->
113    Default;
114get(K, Default, {_, Left, K1, _, _}) when K < K1 ->
115    get(K, Default, Left);
116get(K, Default, {_, _, K1, _, Right}) when K > K1 ->
117    get(K, Default, Right);
118get(_, _, {_, _, _, Val, _}) ->
119    Val.
120
121-spec add(ec_dictionary:key(K), ec_dictionary:value(V),
122          dictionary(K, V)) -> dictionary(K, V).
123add(Key, Value, Dict) ->
124    {_, L, K1, V1, R} = add1(Key, Value, Dict),
125    {b, L, K1, V1, R}.
126
127-spec remove(ec_dictionary:key(K), dictionary(K, V)) -> dictionary(K, V).
128remove(Key, Dictionary) ->
129    {Dict1, _} = erase_aux(Key, Dictionary), Dict1.
130
131-spec has_value(ec_dictionary:value(V), dictionary(_K, V)) -> boolean().
132has_value(Value, Dict) ->
133    fold(fun (_, NValue, _) when NValue == Value -> true;
134             (_, _, Acc) -> Acc
135         end,
136         false, Dict).
137
138-spec size(dictionary(_K, _V)) -> non_neg_integer().
139size(T) ->
140    size1(T).
141
142-spec to_list(dictionary(K, V)) ->
143                     [{ec_dictionary:key(K), ec_dictionary:value(V)}].
144to_list(T) ->
145    to_list(T, []).
146
147-spec from_list([{ec_dictionary:key(K), ec_dictionary:value(V)}]) ->
148                       dictionary(K, V).
149from_list(L) ->
150    lists:foldl(fun ({K, V}, D) ->
151                        add(K, V, D)
152                end, new(),
153                L).
154
155-spec keys(dictionary(K, _V)) -> [ec_dictionary:key(K)].
156keys(Dict) ->
157    keys(Dict, []).
158
159%%%===================================================================
160%%% Enternal functions
161%%%===================================================================
162-spec keys(dictionary(K, _V), [ec_dictionary:key(K)]) ->
163                  [ec_dictionary:key(K)].
164keys(empty, Tail) ->
165    Tail;
166keys({_, L, K, _, R}, Tail) ->
167    keys(L, [K | keys(R, Tail)]).
168
169
170-spec erase_aux(ec_dictionary:key(K), dictionary(K, V)) ->
171                       {dictionary(K, V), boolean()}.
172erase_aux(_, empty) ->
173    {empty, false};
174erase_aux(K, {b, A, Xk, Xv, B}) ->
175    if K < Xk ->
176            {A1, Dec} = erase_aux(K, A),
177            if Dec ->
178                    unbalright(b, A1, Xk, Xv, B);
179               true ->
180                    {{b, A1, Xk, Xv, B}, false}
181            end;
182       K > Xk ->
183            {B1, Dec} = erase_aux(K, B),
184            if Dec ->
185                    unballeft(b, A, Xk, Xv, B1);
186               true ->
187                    {{b, A, Xk, Xv, B1}, false}
188            end;
189       true ->
190            case B of
191                empty ->
192                    blackify(A);
193                _ ->
194                    {B1, {Mk, Mv}, Dec} = erase_min(B),
195                    if Dec ->
196                            unballeft(b, A, Mk, Mv, B1);
197                       true ->
198                            {{b, A, Mk, Mv, B1}, false}
199                    end
200            end
201    end;
202erase_aux(K, {r, A, Xk, Xv, B}) ->
203    if K < Xk ->
204            {A1, Dec} = erase_aux(K, A),
205            if Dec ->
206                    unbalright(r, A1, Xk, Xv, B);
207               true ->
208                    {{r, A1, Xk, Xv, B}, false}
209            end;
210       K > Xk ->
211            {B1, Dec} = erase_aux(K, B),
212            if Dec ->
213                    unballeft(r, A, Xk, Xv, B1);
214               true ->
215                    {{r, A, Xk, Xv, B1}, false}
216            end;
217       true ->
218            case B of
219                empty ->
220                    {A, false};
221                _ ->
222                    {B1, {Mk, Mv}, Dec} = erase_min(B),
223                    if Dec ->
224                            unballeft(r, A, Mk, Mv, B1);
225                       true ->
226                            {{r, A, Mk, Mv, B1}, false}
227                    end
228            end
229    end.
230
231-spec erase_min(dictionary(K, V)) ->
232                       {dictionary(K, V), {ec_dictionary:key(K), ec_dictionary:value(V)}, boolean()}.
233erase_min({b, empty, Xk, Xv, empty}) ->
234    {empty, {Xk, Xv}, true};
235erase_min({b, empty, Xk, Xv, {r, A, Yk, Yv, B}}) ->
236    {{b, A, Yk, Yv, B}, {Xk, Xv}, false};
237erase_min({b, empty, _, _, {b, _, _, _, _}}) ->
238    exit(boom);
239erase_min({r, empty, Xk, Xv, A}) ->
240    {A, {Xk, Xv}, false};
241erase_min({b, A, Xk, Xv, B}) ->
242    {A1, Min, Dec} = erase_min(A),
243    if Dec ->
244            {T, Dec1} = unbalright(b, A1, Xk, Xv, B),
245            {T, Min, Dec1};
246       true -> {{b, A1, Xk, Xv, B}, Min, false}
247    end;
248erase_min({r, A, Xk, Xv, B}) ->
249    {A1, Min, Dec} = erase_min(A),
250    if Dec ->
251            {T, Dec1} = unbalright(r, A1, Xk, Xv, B),
252            {T, Min, Dec1};
253       true -> {{r, A1, Xk, Xv, B}, Min, false}
254    end.
255
256blackify({r, A, K, V, B}) -> {{b, A, K, V, B}, false};
257blackify(Node) -> {Node, true}.
258
259unballeft(r, {b, A, Xk, Xv, B}, Yk, Yv, C) ->
260    {lbalance(b, {r, A, Xk, Xv, B}, Yk, Yv, C), false};
261unballeft(b, {b, A, Xk, Xv, B}, Yk, Yv, C) ->
262    {lbalance(b, {r, A, Xk, Xv, B}, Yk, Yv, C), true};
263unballeft(b, {r, A, Xk, Xv, {b, B, Yk, Yv, C}}, Zk, Zv,
264          D) ->
265    {{b, A, Xk, Xv,
266      lbalance(b, {r, B, Yk, Yv, C}, Zk, Zv, D)},
267     false}.
268
269unbalright(r, A, Xk, Xv, {b, B, Yk, Yv, C}) ->
270    {rbalance(b, A, Xk, Xv, {r, B, Yk, Yv, C}), false};
271unbalright(b, A, Xk, Xv, {b, B, Yk, Yv, C}) ->
272    {rbalance(b, A, Xk, Xv, {r, B, Yk, Yv, C}), true};
273unbalright(b, A, Xk, Xv,
274           {r, {b, B, Yk, Yv, C}, Zk, Zv, D}) ->
275    {{b, rbalance(b, A, Xk, Xv, {r, B, Yk, Yv, C}), Zk, Zv,
276      D},
277     false}.
278
279-spec fold(fun((ec_dictionary:key(K), ec_dictionary:value(V), any()) -> any()),
280           any(), dictionary(K, V)) -> any().
281fold(_, Acc, empty) -> Acc;
282fold(F, Acc, {_, A, Xk, Xv, B}) ->
283    fold(F, F(Xk, Xv, fold(F, Acc, B)), A).
284
285add1(K, V, empty) -> {r, empty, K, V, empty};
286add1(K, V, {C, Left, K1, V1, Right}) when K < K1 ->
287    lbalance(C, add1(K, V, Left), K1, V1, Right);
288add1(K, V, {C, Left, K1, V1, Right}) when K > K1 ->
289    rbalance(C, Left, K1, V1, add1(K, V, Right));
290add1(K, V, {C, L, _, _, R}) -> {C, L, K, V, R}.
291
292size1(empty) -> 0;
293size1({_, L, _, _, R}) -> size1(L) + size1(R) + 1.
294
295to_list(empty, List) -> List;
296to_list({_, A, Xk, Xv, B}, List) ->
297    to_list(A, [{Xk, Xv} | to_list(B, List)]).
298
299%% Balance a tree afer (possibly) adding a node to the left/right.
300-spec lbalance(color(), dictionary(K, V),
301               ec_dictionary:key(K), ec_dictionary:value(V),
302               dictionary(K, V)) ->
303                      dictionary(K, V).
304lbalance(b, {r, {r, A, Xk, Xv, B}, Yk, Yv, C}, Zk, Zv,
305         D) ->
306    {r, {b, A, Xk, Xv, B}, Yk, Yv, {b, C, Zk, Zv, D}};
307lbalance(b, {r, A, Xk, Xv, {r, B, Yk, Yv, C}}, Zk, Zv,
308         D) ->
309    {r, {b, A, Xk, Xv, B}, Yk, Yv, {b, C, Zk, Zv, D}};
310lbalance(C, A, Xk, Xv, B) -> {C, A, Xk, Xv, B}.
311
312-spec rbalance(color(), dictionary(K, V),
313               ec_dictionary:key(K), ec_dictionary:value(V),
314               dictionary(K, V)) ->
315                      dictionary(K, V).
316rbalance(b, A, Xk, Xv,
317         {r, {r, B, Yk, Yv, C}, Zk, Zv, D}) ->
318    {r, {b, A, Xk, Xv, B}, Yk, Yv, {b, C, Zk, Zv, D}};
319rbalance(b, A, Xk, Xv,
320         {r, B, Yk, Yv, {r, C, Zk, Zv, D}}) ->
321    {r, {b, A, Xk, Xv, B}, Yk, Yv, {b, C, Zk, Zv, D}};
322rbalance(C, A, Xk, Xv, B) -> {C, A, Xk, Xv, B}.
323