1 files changed, 111 insertions, 0 deletions
diff --git a/src/dtree.erl b/src/dtree.erl
new file mode 100644
index 0000000000..265bb34009
--- /dev/null
+++ b/src/dtree.erl
@@ -0,0 +1,111 @@
+%% 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.
+%%
+%% Conceptually, what we want is a map that has two distinct sets of
+%% keys (referred to here as primary and secondary, although that
+%% shouldn't imply a hierarchy) pointing to one set of
+%% values. However, in practice what we'll always want to do is insert
+%% a value that's pointed at by (one primary, many secondaries) and
+%% remove values that are pointed at by (one secondary, many
+%% primaries) or (one secondary, all primaries). Thus the API.
+%%
+%% Entries exists while they have a non-empty secondary 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).