An N-Way Ordered Dictionary

lib_nkeys.erl

%%
%% 一个支持多键值记录的库
%%
%% 底层其实是1个gb_tree() + N个dict()，N是键值的数量。
%% 由于最初做世界聊天记录的时候要求按时间排序，所以用了gb_tree()，
%% 所以插入和删除会慢一点（大概几十微秒），但是查找会比较快（几微秒）
%%
%% 如果不在意顺序，完全可以把gb_tree()和dict()换成ets，插入/删除应该能快很~多倍
%%
%% 用法：
%%
%%     > DB = lib_nkeys:new(2).                                     %% 2是键值数量
%%     > Rec1 = #some_record{key = {1, a}, value = some_value}.     %% key必须是第一个字段
%%     > DB2 = lib_nkeys:enter(Rec1, DB).
%%     > Rec2 = #some_record{key = {2, b}, value = some_value}.
%%     > DB3 = lib_nkeys:enter(Rec2, DB2).
%%     > lib_nkeys:smallest(DB3).                                   %% 返回Rec1，因为是按key排序
%%     > lib_nkeys:key_lookup(b, 2, DB3).                           %% 返回{value, [Rec2]}，“2”是单个键值的位置
%%

-module(lib_nkeys).

-export([
        new/1,
        enter/2,
        lookup/2,
        key_lookup/3,
        delete_any/2,
        key_delete_any/3,
        largest/1,
        take_largest/1,
        smallest/1,
        take_smallest/1,
        size/1,
        to_list/1
    ]).

-record(nkeys, {
        data_tree,      %% gb_tree()
        key_maps        %% [dict()]
    }).


-spec new(non_neg_integer()) -> #nkeys{}.
new(NumKeys) ->
    #nkeys{
        data_tree = gb_trees:empty(),
        key_maps = lists:duplicate(NumKeys, dict:new())
    }.


-spec enter(tuple(), #nkeys{}) -> #nkeys{}.
enter(Entry, DB) ->
    Keys = get_keys_from_entry(Entry),
    NewDataTree = gb_trees:enter(Keys, Entry, DB#nkeys.data_tree),
    NewKeyMaps = enter_key_maps(tuple_to_list(Keys), DB#nkeys.key_maps, Keys, []),
    DB#nkeys{data_tree = NewDataTree, key_maps = NewKeyMaps}.


-spec lookup(tuple(), #nkeys{}) -> none | {value, tuple()}.
lookup(Keys, DB) ->
    gb_trees:lookup(Keys, DB#nkeys.data_tree).


-spec key_lookup(any(), non_neg_integer(), #nkeys{}) -> none | {value, [tuple()]}.
key_lookup(Key, KeyIdx, DB) ->
    Map = lists:nth(KeyIdx, DB#nkeys.key_maps),
    case dict:find(Key, Map) of
        {ok, KeysList} ->
            EntryList = lookup_helper(KeysList, DB#nkeys.data_tree, []),
            {value, EntryList};
        _ ->        %% error
            none
    end.


-spec delete_any(tuple(), #nkeys{}) -> #nkeys{}.
delete_any(Keys, DB) ->
    case gb_trees:is_defined(Keys, DB#nkeys.data_tree) of
        true ->
            delete(Keys, DB);
        _ ->
            DB
    end.


-spec key_delete_any(any(), non_neg_integer(), #nkeys{}) -> #nkeys{}.
key_delete_any(Key, KeyIdx, DB) ->
    Map = lists:nth(KeyIdx, DB#nkeys.key_maps),
    case dict:find(Key, Map) of
        {ok, KeysList} ->
            key_delete_any_helper(KeysList, DB);

        _ ->        %% error
            DB
    end.


-spec largest(#nkeys{}) -> none | tuple().
largest(DB) ->
    DataTree = DB#nkeys.data_tree,
    case gb_trees:is_empty(DataTree) of
        true ->
            none;
        _ ->
            {_, Val} = gb_trees:largest(DataTree),
            Val
    end.


-spec take_largest(#nkeys{}) -> none | {tuple(), #nkeys{}}.
take_largest(DB) ->
    take_helper(DB, fun gb_trees:take_largest/1).


-spec smallest(#nkeys{}) -> none | tuple().
smallest(DB) ->
    DataTree = DB#nkeys.data_tree,
    case gb_trees:is_empty(DataTree) of
        true ->
            none;
        _ ->
            {_, Val} = gb_trees:smallest(DataTree),
            Val
    end.


-spec take_smallest(#nkeys{}) -> none | {tuple(), #nkeys{}}.
take_smallest(DB) ->
    take_helper(DB, fun gb_trees:take_smallest/1).


-spec size(#nkeys{}) -> non_neg_integer().
size(DB) ->
    gb_trees:size(DB#nkeys.data_tree).


-spec to_list(#nkeys{}) -> [tuple()].
to_list(DB) ->
    gb_trees:values(DB#nkeys.data_tree).


%% private functions

%% Keys MUST exist in DB
delete(Keys, DB) ->
    NewDataTree = gb_trees:delete(Keys, DB#nkeys.data_tree),
    NewKeyMaps = delete_key_maps(tuple_to_list(Keys), DB#nkeys.key_maps, Keys),
    DB#nkeys{data_tree = NewDataTree, key_maps = NewKeyMaps}.


get_keys_from_entry(Entry) ->
    case element(1, Entry) of
        RecName when is_atom(RecName) ->
            element(2, Entry);
        KTuple when is_tuple(KTuple) ->
            KTuple;
        _ ->
            error(badarg)
    end.


take_helper(DB, TakeFun) ->
    DataTree = DB#nkeys.data_tree,
    case gb_trees:is_empty(DataTree) of
        true ->
            none;
        _ ->
            {Keys, Entry, NewDataTree} = TakeFun(DataTree),
            NewKeyMaps = delete_key_maps(tuple_to_list(Keys), DB#nkeys.key_maps, Keys),
            {Entry, DB#nkeys{data_tree = NewDataTree, key_maps = NewKeyMaps}}
    end.


delete_one_key_map(Key, Map, Keys) ->
    KeysList = dict:fetch(Key, Map),
    case length(KeysList) > 1 of
        false ->
            dict:erase(Key, Map);
        _ ->
            NewList = lists:delete(Keys, KeysList),
            dict:store(Key, NewList, Map)
    end.


delete_key_maps(KeysList, MapList, Keys) ->
    delete_key_maps(KeysList, MapList, Keys, []).

delete_key_maps(_, [], _, AccList) ->
    lists:reverse(AccList);
delete_key_maps([Key | RestKeys], [Map | RestMaps], Keys, AccList) ->
    NewMap = delete_one_key_map(Key, Map, Keys),
    delete_key_maps(RestKeys, RestMaps, Keys, [NewMap | AccList]).


enter_one_key_map(Key, Map, Keys) ->
    F = fun(OldKeysList) ->
        case lists:member(Keys, OldKeysList) of
            true ->
                OldKeysList;
            _ ->
                [Keys | OldKeysList]
        end
    end,
    dict:update(Key, F, [Keys], Map).


enter_key_maps(_, [], _, AccList) ->
    lists:reverse(AccList);
enter_key_maps([Key | RestKeys], [Map | RestMaps], Keys, AccList) ->
    NewMap = enter_one_key_map(Key, Map, Keys),
    enter_key_maps(RestKeys, RestMaps, Keys, [NewMap | AccList]).


lookup_helper([], _, AccList) ->
    AccList;
lookup_helper([Keys | RestKeys], DataTree, AccList) ->
    {value, Entry} = gb_trees:lookup(Keys, DataTree),
    lookup_helper(RestKeys, DataTree, [Entry | AccList]).


key_delete_any_helper([], DB) ->
    DB;
key_delete_any_helper([Keys | RestKeys], DB) ->
    %% XXX: 这样很慢吧……
    NewDB = delete(Keys, DB),
    key_delete_any_helper(RestKeys, NewDB).