file_sorter.erl

来自「OTP是开放电信平台的简称」· ERL 代码 · 共 1,498 行 · 第 1/4 页

%% ``The contents of this file are subject to the Erlang Public License,
%% Version 1.1, (the "License"); you may not use this file except in
%% compliance with the License. You should have received a copy of the
%% Erlang Public License along with this software. If not, it can be
%% retrieved via the world wide web at http://www.erlang.org/.
%% 
%% 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 Initial Developer of the Original Code is Ericsson Utvecklings AB.
%% Portions created by Ericsson are Copyright 1999, Ericsson Utvecklings
%% AB. All Rights Reserved.''
%% 
%%     $Id $
%%
-module(file_sorter).

-export([sort/1, sort/2, sort/3, 
         keysort/2, keysort/3, keysort/4,
         merge/2, merge/3, 
         keymerge/3, keymerge/4,
         check/1, check/2, 
         keycheck/2, keycheck/3]).

-include_lib("kernel/include/file.hrl").

-define(CHUNKSIZE, 16384).
-define(RUNSIZE, 524288).
-define(NOMERGE, 16).
-define(MERGESIZE, ?CHUNKSIZE).

-define(MAXSIZE, (1 bsl 31)).

-record(w, {keypos, runs = [[]], seq = 1, in, out, fun_out, prefix, temp = [],
            format, runsize, no_files, order, chunksize, wfd, ref, z, unique,
            hdlen, inout_value}).

-record(opts, {format = binary_term_fun(), size = ?RUNSIZE, 
               no_files = ?NOMERGE, tmpdir = default, order = ascending, 
               compressed = false, unique = false, header = 4}).

-compile({inline, [{badarg, 2}, {make_key,2}, {make_stable_key,3}, {cfun,3}]}).

%%%
%%% Exported functions
%%%

sort(FileName) ->
    sort([FileName], FileName).

sort(Input, Output) ->
    sort(Input, Output, []).

sort(Input0, Output0, Options) ->
    case {is_input(Input0), maybe_output(Output0), options(Options)}  of
        {{true,Input}, {true,Output}, #opts{}=Opts} -> 
            do_sort(0, Input, Output, Opts, sort);
        T -> 
            badarg(culprit(tuple_to_list(T)), [Input0, Output0, Options])
    end.

keysort(KeyPos, FileName) ->
    keysort(KeyPos, [FileName], FileName).

keysort(KeyPos, Input, Output) ->
    keysort(KeyPos, Input, Output, []).

keysort(KeyPos, Input0, Output0, Options) ->
    R = case {is_keypos(KeyPos), is_input(Input0), 
              maybe_output(Output0), options(Options)} of
            {_, _, _, #opts{format = binary}} -> 
                {Input0,Output0,[{badarg,format}]};
            {_, _, _, #opts{order = Order}} when is_function(Order) -> 
                {Input0,Output0,[{badarg,order}]};
            {true, {true,In}, {true,Out}, #opts{}=Opts} -> 
                {In,Out,Opts};
            T -> 
                {Input0,Output0,tuple_to_list(T)}
        end,
    case R of
        {Input,Output,#opts{}=O} ->
            do_sort(KeyPos, Input, Output, O, sort);
        {_,_,O} ->
            badarg(culprit(O), [KeyPos, Input0, Output0, Options])
    end.

merge(Files, Output) ->
    merge(Files, Output, []).

merge(Files0, Output0, Options) ->
    case {is_files(Files0), maybe_output(Output0), options(Options)} of
        %% size not used
        {{true,Files}, {true,Output}, #opts{}=Opts} ->
            do_sort(0, Files, Output, Opts, merge);
        T -> 
            badarg(culprit(tuple_to_list(T)), [Files0, Output0, Options])
    end.

keymerge(KeyPos, Files, Output) ->
    keymerge(KeyPos, Files, Output, []).

keymerge(KeyPos, Files0, Output0, Options) ->
    R = case {is_keypos(KeyPos), is_files(Files0), 
              maybe_output(Output0), options(Options)} of
            {_, _, _, #opts{format = binary}} -> 
                {Files0,Output0,[{badarg,format}]};
            {_, _, _, #opts{order = Order}} when is_function(Order) -> 
                {Files0,Output0,[{badarg,order}]};
            {true, {true,Fs}, {true,Out}, #opts{}=Opts} -> 
                {Fs,Out,Opts};
            T -> 
                {Files0,Output0,tuple_to_list(T)}
        end,
    case R of
        {Files,Output,#opts{}=O} -> 
            do_sort(KeyPos, Files, Output, O, merge);
        {_,_,O} -> 
            badarg(culprit(O), [KeyPos, Files0, Output0, Options])
    end.

check(FileName) ->
    check([FileName], []).

check(Files0, Options) ->
    case {is_files(Files0), options(Options)} of
        {{true,Files}, #opts{}=Opts} ->
            do_sort(0, Files, undefined, Opts, check);
        T ->
            badarg(culprit(tuple_to_list(T)), [Files0, Options])
    end.

keycheck(KeyPos, FileName) ->
    keycheck(KeyPos, [FileName], []).

keycheck(KeyPos, Files0, Options) ->
    R = case {is_keypos(KeyPos), is_files(Files0), options(Options)} of
            {_, _, #opts{format = binary}} -> 
                {Files0,[{badarg,format}]};
            {_, _, #opts{order = Order}} when is_function(Order) -> 
                {Files0,[{badarg,order}]};
            {true, {true,Fs}, #opts{}=Opts} -> 
                {Fs,Opts};
            T -> 
                {Files0,tuple_to_list(T)}
        end,
    case R of
        {Files,#opts{}=O} -> 
            do_sort(KeyPos, Files, undefined, O, check);
        {_,O} -> 
            badarg(culprit(O), [KeyPos, Files0, Options])
    end.

%%%
%%% Local functions
%%%

%%-define(debug, true).

-ifdef(debug).
-define(DEBUG(S, A), io:format(S, A)).
-else.
-define(DEBUG(S, A), ok).
-endif.

culprit([{error, _} = E | _]) ->
    E;
culprit([{badarg, _} = B | _]) ->
    B;
culprit([_ | B]) -> 
    culprit(B).

%% Inlined.
badarg({error, _} = E, _Args) ->
    E;
badarg({badarg, _} = B, Args) ->
    erlang:fault(B, Args).

options(Options) when is_list(Options) ->
    options(Options, #opts{});
options(Option) ->
    options([Option]).

options([{format, Format} | L], Opts) when Format =:= binary; 
                                           Format =:= term;
                                           is_function(Format),
                                             is_function(Format, 1) ->
    options(L, Opts#opts{format = Format});
options([{format, binary_term} | L], Opts) ->
    options(L, Opts#opts{format = binary_term_fun()});
options([{size, Size} | L], Opts) when is_integer(Size), Size >= 0 ->
    options(L, Opts#opts{size = max(Size, 1)});
options([{no_files, NoFiles} | L], Opts) when is_integer(NoFiles), 
                                              NoFiles > 1 ->
    options(L, Opts#opts{no_files = NoFiles});
options([{tmpdir, ""} | L], Opts) ->
    options(L, Opts#opts{tmpdir = default});
options([{tmpdir, Dir} | L],  Opts) ->
    case catch filename:absname(Dir) of
        {'EXIT', _} ->
            {badarg, Dir};
        FileName -> 
            options(L, Opts#opts{tmpdir = {dir, FileName}})
    end;
options([{order, Fun} | L], Opts) when is_function(Fun), is_function(Fun, 2) ->
    options(L, Opts#opts{order = Fun});
options([{order, Order} | L], Opts) when Order =:= ascending; 
                                         Order =:= descending ->
    options(L, Opts#opts{order = Order});
options([{compressed, Bool} | L], Opts) when Bool; not Bool ->
    options(L, Opts#opts{compressed = Bool});
options([{unique, Bool} | L], Opts) when Bool; not Bool ->
    options(L, Opts#opts{unique = Bool});
options([{header, Len} | L], Opts) 
                when is_integer(Len), Len > 0, Len < ?MAXSIZE ->
    options(L, Opts#opts{header = Len});
options([], Opts) ->
    if 
        Opts#opts.format =:= term, Opts#opts.header =/= 4 ->
            {badarg, header};
        true ->
            Opts
    end;
options([Bad | _], _Opts) ->
    {badarg, Bad};
options(Bad, _Opts) ->
    {badarg, Bad}.

-define(OBJ(X, Y), {X, Y}).
-define(SK(T, I), [T | I]). % stable key

do_sort(KeyPos0, Input0, Output0, Opts, Do) ->
    #opts{format = Format0, size = Size, no_files = NoFiles,
          tmpdir = TmpDir, order = Order, compressed = Compressed,
          unique = Unique, header = HdLen} = Opts,
    Prefix = tmp_prefix(Output0, TmpDir),
    ChunkSize = ?CHUNKSIZE,
    Ref = make_ref(),
    KeyPos = case KeyPos0 of [Kp] -> Kp; _ -> KeyPos0 end,
    {Format, Input} = wrap_input(Format0, Do, Input0),
    Z = if Compressed -> [compressed]; true -> [] end,
    {Output, FunOut} = wrap_output_terms(Format0, Output0, Z),
    W = #w{keypos = KeyPos, out = Output, fun_out = FunOut,
           prefix = Prefix, format = Format, runsize = Size, 
           no_files = NoFiles, order = Order, chunksize = ChunkSize, 
           ref = Ref, z = Z, unique = Unique, hdlen = HdLen, 
           inout_value = no_value},
    try
        doit(Do, Input, W)
    catch {Ref,Error} ->
        Error
    end.
    
doit(sort, Input, W) ->
    files(1, [], 0, W, Input);
doit(merge, Input, W) ->
    last_merge(Input, W);
doit(check, Input, W) ->
    check_files(Input, W, []).

wrap_input(term, check, Files) ->
    Fun = fun(File) -> 
                  Fn = merge_terms_fun(file_rterms(no_file, [File])),
                  {fn, Fn, File}
          end,
    {binary_term_fun(), lists:map(Fun, Files)};
wrap_input(Format, check, Files) ->
    {Format, Files};
wrap_input(term, merge, Files) ->
    Fun = fun(File) -> merge_terms_fun(file_rterms(no_file, [File])) end,
    Input = lists:reverse(lists:map(Fun, Files)),
    {binary_term_fun(), Input};
wrap_input(Format, merge, Files) ->
    Input = lists:reverse(lists:map(fun merge_bins_fun/1, Files)),
    {Format, Input};
wrap_input(term, sort, InFun) when is_function(InFun), is_function(InFun, 1) ->
    {binary_term_fun(), fun_rterms(InFun)};
wrap_input(term, sort, Files) ->
    {binary_term_fun(), file_rterms(no_file, Files)};
wrap_input(Format, sort, Input) ->
    {Format, Input}.

merge_terms_fun(RFun) ->
    fun(close) ->
            RFun(close);
       ({I, [], _LSz, W}) ->
            case RFun(read) of
                end_of_input ->
                    eof;
                {Objs, NRFun} when is_function(NRFun), is_function(NRFun, 1) ->
                    {_, [], Ts, _} = fun_objs(Objs, [], 0, ?MAXSIZE, I, W),
                    {{I, Ts, ?CHUNKSIZE}, merge_terms_fun(NRFun)};
                Error ->
                    error(Error, W)
            end
    end.

merge_bins_fun(FileName) ->
    fun(close) ->
            ok;
       ({_I, _L, _LSz, W} = A) ->
            Fun = read_fun(FileName, user, W),
            Fun(A)
    end.

wrap_output_terms(term, OutFun, _Z) when is_function(OutFun),
                                         is_function(OutFun, 1) ->
    {fun_wterms(OutFun), true};
wrap_output_terms(term, File, Z) when File =/= undefined ->
    {file_wterms(name, File, Z++[write]), false};
wrap_output_terms(_Format, Output, _Z) ->
    {Output, is_function(Output) and is_function(Output, 1)}.

binary_term_fun() ->
    fun binary_to_term/1.

check_files([], _W, L) ->
    {ok, lists:reverse(L)};
check_files([FN | FNs], W, L) ->
    {IFun, FileName} = 
        case FN of
            {fn, Fun, File} ->
                {Fun, File};
            File ->
                {read_fun(File, user, W), File}
        end,
    NW = W#w{in = IFun},
    check_run(IFun, FileName, FNs, NW, L, 2, nolast).

check_run(IFun, F, FNs, W, L, I, Last) ->
    case IFun({{merge,I}, [], 0, W}) of
        {{_I, Objs, _LSz}, IFun1} ->
            NW = W#w{in = IFun1},
            check_objs0(IFun1, F, FNs, NW, L, I, Last, lists:reverse(Objs));
        eof ->
            NW = W#w{in = undefined},
            check_files(FNs, NW, L)
    end.

check_objs0(IFun, F, FNs, W, L, I, nolast, [?OBJ(T,_BT) | Os]) ->
    check_objs1(IFun, F, FNs, W, L, I, T, Os);
check_objs0(IFun, F, FNs, W, L, I, Last, []) ->
    check_run(IFun, F, FNs, W, L, I, Last);
check_objs0(IFun, F, FNs, W, L, I, {last, Last}, Os) ->
    check_objs1(IFun, F, FNs, W, L, I, Last, Os).

check_objs1(IFun, F, FNs, W, L, I, LastT, Os) ->
    case W of
        #w{order = ascending, unique = true} ->
            ucheck_objs(IFun, F, FNs, W, L, I, LastT, Os);
        #w{order = ascending, unique = false} ->
            check_objs(IFun, F, FNs, W, L, I, LastT, Os);
        #w{order = descending, unique = true} ->
            rucheck_objs(IFun, F, FNs, W, L, I, LastT, Os);
        #w{order = descending, unique = false} ->
            rcheck_objs(IFun, F, FNs, W, L, I, LastT, Os);
        #w{order = CF, unique = true} ->
            uccheck_objs(IFun, F, FNs, W, L, I, LastT, Os, CF);
        #w{order = CF, unique = false} ->
            ccheck_objs(IFun, F, FNs, W, L, I, LastT, Os, CF)
    end.

check_objs(IFun, F, FNs, W, L, I, Last, [?OBJ(T,_BT) | Os]) when T >= Last ->
    check_objs(IFun, F, FNs, W, L, I+1, T, Os);
check_objs(IFun, F, FNs, W, L, I, _Last, [?OBJ(_T,BT) | _]) ->
    culprit_found(IFun, F, FNs, W, L, I, BT);
check_objs(IFun, F, FNs, W, L, I, Last, []) ->
    check_run(IFun, F, FNs, W, L, I, {last, Last}).

rcheck_objs(IFun, F, FNs, W, L, I, Last, [?OBJ(T,_BT) | Os]) when T =< Last ->
    rcheck_objs(IFun, F, FNs, W, L, I+1, T, Os);
rcheck_objs(IFun, F, FNs, W, L, I, _Last, [?OBJ(_T,BT) | _]) ->
    culprit_found(IFun, F, FNs, W, L, I, BT);
rcheck_objs(IFun, F, FNs, W, L, I, Last, []) ->