%%    Copyright 2009 Igor Ribeiro Sucupira.
%%
%%    This module is free software; you can redistribute it and/or
%%    modify it under the terms of the GNU Lesser General Public
%%    License as published by the Free Software Foundation; either
%%    version 2.1 of the License, or (at your option) any later version.
%%
%%    This module is distributed in the hope that it will be useful,
%%    but WITHOUT ANY WARRANTY; without even the implied warranty of
%%    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
%%    Lesser General Public License for more details.


%%%-------------------------------------------------------------------------------------------
%%% Implements consistent hashing functionality for fragmented tables.
%%%
%%% Assumes the caller has created a disc_copies, ordered_set, non-fragmented table named
%%% like the table being created here, but with the _chash suffix added. This table should
%%% be replicated through all the pool and must have exactly 2 attributes (including the key).
%%%-------------------------------------------------------------------------------------------


-module(mnesia_frag_chash).

%%-behaviour(mnesia_frag_hash).

-export([
	 init_state/2,
	 add_frag/1,
	 del_frag/1,
	 key_to_frag_number/2,
	 match_spec_to_frag_numbers/2
	]).

-record(hash_state,
	{n_fragments,
         n_frag_pieces,
         chash_table}).


-define(FRAG_CHASH_LIMIT, 1000000007).
-define(N_FRAG_PIECES, 100).
%% This limits the number of fragments to something around 1 million (no problem).
-define(FRAG_NUM_BITS, 20).


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%     Exported functions     %%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


init_state(Tab, State) when is_atom(Tab), State == undefined ->
    ChashTable = list_to_atom(atom_to_list(Tab) ++ "_chash"),
    NFragPieces = ?N_FRAG_PIECES,
    insert_frag_entries(1, NFragPieces, ChashTable),
    #hash_state{n_fragments   = 1,
                n_frag_pieces = NFragPieces,
                chash_table   = ChashTable}.


add_frag(#hash_state{n_fragments = NFragments, n_frag_pieces = NFragPieces, chash_table = ChashTable} = State) ->
    AffectedFragments = all_upper_bounds(NFragments + 1, NFragPieces, ChashTable),
    insert_frag_entries(NFragments + 1, NFragPieces, ChashTable),
    NewState = State#hash_state{n_fragments = NFragments + 1},
    {NewState, AffectedFragments, [NFragments + 1]}.


del_frag(#hash_state{n_fragments = NFragments, n_frag_pieces = NFragPieces, chash_table = ChashTable} = State) ->
    AffectedFragments = all_upper_bounds(NFragments, NFragPieces, ChashTable),
    del_frag_entries(NFragments, NFragPieces, ChashTable),
    NewState = State#hash_state{n_fragments = NFragments - 1},
    {NewState, [NFragments], AffectedFragments}.

key_to_frag_number(#hash_state{chash_table = ChashTable}, Key) ->
    Entry = ets:next(ChashTable, make_entry(hash_val(Key), 0, 0)),
    case Entry of
        '$end_of_table' ->
            get_frag_num(ets:first(ChashTable));
        _ ->
            get_frag_num(Entry)
    end.

match_spec_to_frag_numbers(#hash_state{n_fragments = N} = State, MatchSpec) ->
    case MatchSpec of
        [{HeadPat, _, _}] when is_tuple(HeadPat), tuple_size(HeadPat) > 2 ->
            KeyPat = element(2, HeadPat),
            case mnesia:has_var(KeyPat) of
                false ->
                    %% If the key is known, so is the fragment.
                    [key_to_frag_number(State, KeyPat)];
                true ->
                    lists:seq(1, N)
            end;
        _ ->
            lists:seq(1, N)
    end.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%     Internal functions     %%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

insert_frag_entries(_, NumPieces, ChashTable) when is_integer(NumPieces), NumPieces =< 0 ->
    ChashTable;
insert_frag_entries(FragNum, NumPieces, ChashTable) when is_integer(NumPieces) ->
    HashVal = hash_val({FragNum, NumPieces}),
    mnesia:dirty_write({ChashTable, make_entry(HashVal, FragNum, NumPieces), []}),
    insert_frag_entries(FragNum, NumPieces - 1, ChashTable).

del_frag_entries(_, NumPieces, ChashTable) when is_integer(NumPieces), NumPieces =< 0 ->
    ChashTable;
del_frag_entries(FragNum, NumPieces, ChashTable) when is_integer(NumPieces) ->
    HashVal = hash_val({FragNum, NumPieces}),
    mnesia:dirty_delete({ChashTable, make_entry(HashVal, FragNum, NumPieces)}),
    del_frag_entries(FragNum, NumPieces - 1, ChashTable).

get_frag_num({_HashVal, FragEntry}) ->
    FragEntry rem (1 bsl ?FRAG_NUM_BITS).

make_entry(HashVal, FragNumber, FragPiece) ->
    {HashVal, (FragPiece bsl ?FRAG_NUM_BITS) bor FragNumber}.

hash_val(Term) ->
    erlang:phash2(Term, ?FRAG_CHASH_LIMIT).

all_upper_bounds(FragNum, NumPieces, ChashTable) when is_integer(FragNum), is_integer(NumPieces), is_atom(ChashTable) ->
    UpperBoundsSet = lists:foldl(fun(FragPiece, AccSet) ->
                                         sets:add_element(circ_upper_bound(FragNum, FragPiece, ChashTable), AccSet)
                                 end, sets:new(), lists:seq(1, NumPieces)),
    sets:to_list(sets:del_element(FragNum, UpperBoundsSet)).

%% @doc Find the first fragment after the entry for FragNum/FragPiece it in the table, circularly.
circ_upper_bound(FragNum, FragPiece, ChashTable) when is_integer(FragNum), is_integer(FragPiece), is_atom(ChashTable) ->
    HashVal = hash_val({FragNum, FragPiece}),
    Entry = make_entry(HashVal, FragNum, FragPiece),
    case ets:next(ChashTable, Entry) of
        '$end_of_table' ->
            get_frag_num(ets:first(ChashTable));
        UpperBound ->
            get_frag_num(UpperBound)
    end.