file_sorter.erl

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

fun_loop(Objs, L, LSz, RunSize, I, Keypos, Fun, HdLen) when is_integer(I) ->
    fun_keyloop(Objs, L, LSz, RunSize, I, Keypos, Fun, HdLen).

fun_binloop([B | Bs], L, LSz, RunSize, HL) when LSz < RunSize ->
    Size = size(B),
    Obj = ?OBJ(B, [<<Size:HL/unit:8>> | B]),
    fun_binloop(Bs, [Obj | L], LSz+Size, RunSize, HL);
fun_binloop(Bs, L, LSz, _RunSize, _HL) ->
    {0, Bs, L, LSz}.

fun_loop([B | Bs], L, LSz, RunSize, Fun, HL) when LSz < RunSize ->
    Size = size(B),
    Obj = ?OBJ(Fun(B), [<<Size:HL/unit:8>> | B]),
    fun_loop(Bs, [Obj | L], LSz+Size, RunSize, Fun, HL);
fun_loop(Bs, L, LSz, _RunSize, _Fun, _HL) ->
    {0, Bs, L, LSz}.

fun_keyloop([B | Bs], L, LSz, RunSize, I, Kp, Fun, HL) when LSz < RunSize ->
    Size = size(B),
    UniqueKey = make_key(Kp, Fun(B)),
    E = ?OBJ(UniqueKey, [<<Size:HL/unit:8>> | B]),
    fun_keyloop(Bs, [E | L], LSz+Size, RunSize, I+1, Kp, Fun, HL);
fun_keyloop(Bs, L, LSz, _RunSize, I, _Kp, _Fun, _HL) ->
    {I, Bs, L, LSz}.

fun_mergeloop([B | Bs], L, LSz, RunSize, I, Kp, Fun, HL) when LSz < RunSize ->
    Size = size(B),
    UniqueKey = make_stable_key(Kp, I, Fun(B)),
    E = ?OBJ(UniqueKey, [<<Size:HL/unit:8>> | B]),
    fun_mergeloop(Bs, [E | L], LSz+Size, RunSize, I, Kp, Fun, HL);
fun_mergeloop(Bs, L, LSz, _RunSize, I, _Kp, _Fun, _HL) ->
    {{merge, I}, Bs, L, LSz}. % any I would do

%% Inlined.
make_key(Kp, T) when is_integer(Kp) ->
    element(Kp, T);
make_key([Kp1, Kp2], T) ->
    [element(Kp1, T), element(Kp2, T)];
make_key([Kp1, Kp2 | Kps], T) ->
    [element(Kp1, T), element(Kp2, T) | make_key2(Kps, T)].

%% Inlined.
%% A sequence number (I) is used for making the internal sort stable.
%% I is ordering number of the file from which T was read.
make_stable_key(Kp, I, T) when is_integer(Kp) ->
    ?SK(element(Kp, T), I);
make_stable_key([Kp1, Kp2], I, T) ->
    ?SK([element(Kp1, T) | element(Kp2, T)], I);
make_stable_key([Kp1, Kp2 | Kps], I, T) ->
    ?SK([element(Kp1, T), element(Kp2, T) | make_key2(Kps, T)], I).

make_key2([Kp], T) ->
    [element(Kp, T)];
make_key2([Kp | Kps], T) ->
    [element(Kp, T) | make_key2(Kps, T)].

max(A, B) when A < B -> B;
max(A, _) -> A.

infun(W) ->
    W1 = W#w{in = undefined},
    try (W#w.in)(read) of
        end_of_input ->
            {end_of_input, W1};
        {end_of_input, Value} ->
            {end_of_input, W1#w{inout_value = {value, Value}}};
        {Objs, NFun} when is_function(NFun), 
                          is_function(NFun, 1), 
                          is_list(Objs) ->
            {cont, W#w{in = NFun}, Objs};
        Error ->
            error(Error, W1)
    catch Class:Reason ->
        cleanup(W1),
        erlang:raise(Class, Reason, erlang:get_stacktrace())
    end.

outfun(A, W) when W#w.inout_value =/= no_value ->
    W1 = W#w{inout_value = no_value},
    W2 = if 
             W1#w.fun_out ->
                 outfun(W#w.inout_value, W1);
             true -> W1
         end,
    outfun(A, W2);
outfun(A, W) ->
    W1 = W#w{out = undefined},
    try (W#w.out)(A) of
        Reply when A =:= close ->
            Reply;
        NF when is_function(NF), is_function(NF, 1) ->
            W#w{out = NF};
        Error ->
            error(Error, W1)
    catch Class:Reason ->
        cleanup(W1),
        erlang:raise(Class, Reason, erlang:get_stacktrace())
    end.

is_keypos(Keypos) when is_integer(Keypos), Keypos > 0 ->
    true;
is_keypos([]) ->
    {badarg, []};
is_keypos(L) ->
    is_keyposs(L).

is_keyposs([Kp | Kps]) when is_integer(Kp), Kp > 0 ->
    is_keyposs(Kps);
is_keyposs([]) ->
    true;
is_keyposs([Bad | _]) ->
    {badarg, Bad};
is_keyposs(Bad) ->
    {badarg, Bad}.

is_input(Fun) when is_function(Fun), is_function(Fun, 1) ->
    {true, Fun};
is_input(Files) ->
    is_files(Files).

is_files(Fs) ->
    is_files(Fs, []).

is_files([F | Fs], L) ->
    case read_file_info(F) of
        {ok, File, _FI} ->
            is_files(Fs, [File | L]);
        Error ->
            Error
    end;
is_files([], L) ->
    {true, lists:reverse(L)};
is_files(Bad, _L) ->
    {badarg, Bad}.

maybe_output(Fun) when is_function(Fun), is_function(Fun, 1) ->
    {true, Fun};
maybe_output(File) ->
    case read_file_info(File) of
        {badarg, _File} = Badarg ->
            Badarg;
        {ok, FileName, _FileInfo} ->
            {true, FileName};
        {error, {file_error, FileName, _Reason}} ->
            {true, FileName}
    end.

read_file_info(File) ->
    %% Absolute names in case some process should call file:set_cwd/1.
    case catch filename:absname(File) of
        {'EXIT', _} ->
            {badarg, File};
        FileName ->
            case file:read_file_info(FileName) of
                {ok, FileInfo} ->
                    {ok, FileName, FileInfo};
                {error, einval} ->
                    {badarg, File};
                {error, Reason} ->
                    {error, {file_error, FileName, Reason}}
            end
    end.

%% No attempt is made to avoid overwriting existing files.
next_temp(W) ->
    Seq = W#w.seq,
    NW = W#w{seq = Seq + 1},
    Temp = lists:concat([W#w.prefix, Seq]),
    {NW, Temp}.

%% Would use the temporary directory (TMP|TEMP|TMPDIR), were it
%% readily accessible.
tmp_prefix(F, TmpDirOpt) when is_function(F); F =:= undefined ->
    {ok, CurDir} = file:get_cwd(),
    tmp_prefix1(CurDir, TmpDirOpt);
tmp_prefix(OutFile, TmpDirOpt) ->
    Dir = filename:dirname(OutFile),
    tmp_prefix1(Dir, TmpDirOpt).

tmp_prefix1(Dir, TmpDirOpt) ->
    U = "_",
    Node = node(),
    Pid = os:getpid(),
    {MSecs,Secs,MySecs} = now(),
    F = lists:concat(["fs_",Node,U,Pid,U,MSecs,U,Secs,U,MySecs,"."]),
    TmpDir = case TmpDirOpt of
                 default ->
                     Dir;
                 {dir, TDir} ->
                     TDir
             end,
    filename:join(filename:absname(TmpDir), F).

%% -> {Fd, NewW} | throw(Error)
open_file(FileName, W) ->
    case file:open(FileName, W#w.z ++ [raw, binary, write]) of
        {ok, Fd}  ->
            {Fd, W#w{temp = [{Fd,FileName} | W#w.temp]}};
        Error ->
            file_error(FileName, Error, W)
    end.

read(Fd, FileName, N, W) ->
    case file:read(Fd, N) of
        {ok, Bin} ->
            {ok, Bin};
        eof ->
            eof;
        {error, enomem} ->
            %% Bad N
            error({error, {bad_object, FileName}}, W);
        {error, einval} ->
            %% Bad N
            error({error, {bad_object, FileName}}, W);
        Error ->
            file_error(FileName, Error, W)
    end.

write(Fd, FileName, B, W) ->
    case file:write(Fd, B) of
        ok ->
            ok;
        Error ->
            file_error(FileName, Error, W)
    end.

file_error(File, {error, Reason}, W) ->
    error({error, {file_error, File, Reason}}, W).

error(Error, W) ->
    cleanup(W),
    throw({W#w.ref, Error}).

cleanup(W) ->
    close_out(W),
    W1 = close_input(W),
    F = fun(IFun) when is_function(IFun) -> 
                IFun(close);
           ({Fd,FileName}) ->
                file:close(Fd),
                file:delete(FileName);
           (FileName) -> 
                file:delete(FileName)
        end,
    lists:foreach(F, W1#w.temp).

close_input(W) when is_function(W#w.in) ->
    catch (W#w.in)(close),
    W#w{in = undefined};
close_input(#w{in = undefined}=W) ->
    W.

close_out(W) when is_function(W#w.out) ->
    catch (W#w.out)(close);
close_out(_) -> 
    ok.

close_file(Fd, W) ->
    {value, {Fd, FileName}} = lists:keysearch(Fd, 1, W#w.temp),
    ?DEBUG("closing ~p~n", [FileName]),
    file:close(Fd),
    W#w{temp = [FileName | lists:keydelete(Fd, 1, W#w.temp)]}.

%%%
%%% Format 'term'.
%%%

file_rterms(no_file, Files) ->
    fun(close) ->
            ok;
       (read) when Files =:= [] ->
            end_of_input;
       (read) ->
            [F | Fs] = Files,
            case file:open(F, [read, compressed]) of
                {ok, Fd} ->
                    file_rterms2(Fd, [], 0, F, Fs);
                {error, Reason} ->
                    {error, {file_error, F, Reason}}
            end
    end;
file_rterms({Fd, FileName}, Files) ->
    fun(close) ->
            file:close(Fd);
       (read) ->
            file_rterms2(Fd, [], 0, FileName, Files)
    end.

file_rterms2(Fd, L, LSz, FileName, Files) when LSz < ?CHUNKSIZE ->
    case io:read(Fd, '') of
        {ok, Term} ->
            B = term_to_binary(Term),
            file_rterms2(Fd, [B | L], LSz + size(B), FileName, Files);
        eof ->
            file:close(Fd),
            {lists:reverse(L), file_rterms(no_file, Files)};
        _Error ->
            file:close(Fd),
            {error, {bad_term, FileName}}
    end;
file_rterms2(Fd, L, _LSz, FileName, Files) ->
    {lists:reverse(L), file_rterms({Fd, FileName}, Files)}.

file_wterms(W, F, Args) ->
    fun(close) when W =:= name ->
            ok;
       (close) ->
            {fd, Fd} = W,
            file:close(Fd);
       (L) when W =:= name ->
            case file:open(F, Args) of
                {ok, Fd} ->
                    write_terms(Fd, F, L, Args);
                {error, Reason} ->
                    {error, {file_error, F, Reason}}
            end;
       (L) ->
            {fd, Fd} = W,
            write_terms(Fd, F, L, Args)
    end.

write_terms(Fd, F, [B | Bs], Args) ->
    case io:request(Fd, {format, "~p.~n", [binary_to_term(B)]}) of
        ok -> 
            write_terms(Fd, F, Bs, Args);
        {error, Reason} ->
            file:close(Fd),
            {error, {file_error, F, Reason}}
    end;
write_terms(Fd, F, [], Args) ->
    file_wterms({fd, Fd}, F, Args).

fun_rterms(InFun) ->
    fun(close) ->
            InFun(close);
       (read) ->
            case InFun(read) of
                {Ts, NInFun} when is_list(Ts), 
                                  is_function(NInFun),
                                  is_function(NInFun, 1) ->
                    {to_bin(Ts, []), fun_rterms(NInFun)};
                Else ->
                    Else
            end
    end.

fun_wterms(OutFun) ->
    fun(close) ->
            OutFun(close);
       (L) ->
            case OutFun(wterms_arg(L)) of
                NOutFun when is_function(NOutFun), is_function(NOutFun, 1) ->
                    fun_wterms(NOutFun);
                Else ->
                    Else
            end
    end.

to_bin([E | Es], L) ->
    to_bin(Es, [term_to_binary(E) | L]);
to_bin([], L) ->
    lists:reverse(L).

wterms_arg(L) when is_list(L) ->
    to_term(L, []);
wterms_arg(Value) ->
    Value.

to_term([B | Bs], L) ->
    to_term(Bs, [binary_to_term(B) | L]);
to_term([], L) ->
    lists:reverse(L).