%% This file is a copy of erl_id_trans.erl from the R14B02 Erlang/OTP
%% distribution, with modifications to make it implement Haskell-style
%% 'do' syntax in Erlang.

%% All modifications are (C) 2011-2013 VMware, Inc; Eduard Sergeev.

%%
%% ``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.''
%%

-module(do).

-export([parse_transform/2, format_error/1]).

parse_transform(Forms, _Options) ->
    Forms1 = forms(Forms),
    %%io:format("~s~n", [erl_prettypr:format(erl_syntax:form_list(Forms1))]),
    Forms1.

%% forms(Fs) -> lists:map(fun (F) -> form(F) end, Fs).

forms([F0 | Fs0]) ->
    F1 =
        try
            form(F0)
        catch
            throw:{Error, Line} ->
                {error, {Line, ?MODULE, Error}}
        end,
    Fs1 = forms(Fs0),
    [F1 | Fs1];
forms([]) ->
    [].

%% -type form(Form) -> Form.

%The general attribute.
form({attribute, Line, Attr, Val}) ->
    {attribute, Line, Attr, Val};
form({function, Line, Name0, Arity0, Clauses0}) ->
    {Name, Arity, Clauses} = function(Name0, Arity0, Clauses0),
    {function, Line, Name, Arity, Clauses};
% Mnemosyne, ignore...
form({rule, Line, Name, Arity, Body}) ->
    % Dont dig into this
    {rule, Line, Name, Arity, Body};
%% Extra forms from the parser.
form({error, E}) ->
    {error, E};
form({warning, W}) ->
    {warning, W};
form({eof, Line}) ->
    {eof, Line}.

%% -type function(atom(), integer(), [Clause]) -> {atom(),integer(),[Clause]}.

function(Name, Arity, Clauses0) ->
    Clauses1 = clauses(Clauses0),
    {Name, Arity, Clauses1}.

%% -type clauses([Clause]) -> [Clause].

clauses([C0 | Cs]) ->
    C1 = clause(C0, []),
    [C1 | clauses(Cs)];
clauses([]) ->
    [].

%% -type clause(Clause) -> Clause.

clause({clause, Line, Head, Guard, Body}, MonadStack) ->
    {clause, Line, Head, Guard, exprs(Body, MonadStack)}.

%% -type pattern(Pattern) -> Pattern.
%%  N.B. Only valid patterns are included here.

pattern({var, Line, V}) ->
    {var, Line, V};
pattern({match, Line, L0, R0}) ->
    L1 = pattern(L0),
    R1 = pattern(R0),
    {match, Line, L1, R1};
pattern({integer, Line, I}) ->
    {integer, Line, I};
pattern({char, Line, C}) ->
    {char, Line, C};
pattern({float, Line, F}) ->
    {float, Line, F};
pattern({atom, Line, A}) ->
    {atom, Line, A};
pattern({string, Line, S}) ->
    {string, Line, S};
pattern({nil, Line}) ->
    {nil, Line};
pattern({cons, Line, H0, T0}) ->
    H1 = pattern(H0),
    T1 = pattern(T0),
    {cons, Line, H1, T1};
pattern({tuple, Line, Ps0}) ->
    Ps1 = pattern_list(Ps0),
    {tuple, Line, Ps1};
%% OTP 17.0: EEP 443: Map pattern
pattern({map, Line, Fields0}) ->
    Fields1 = map_fields(Fields0, []),
    {map, Line, Fields1};
%%pattern({struct,Line,Tag,Ps0}) ->
%%    Ps1 = pattern_list(Ps0),
%%    {struct,Line,Tag,Ps1};
pattern({record, Line, Name, Pfs0}) ->
    Pfs1 = pattern_fields(Pfs0),
    {record, Line, Name, Pfs1};
pattern({record_index, Line, Name, Field0}) ->
    Field1 = pattern(Field0),
    {record_index, Line, Name, Field1};
%% record_field occurs in query expressions
pattern({record_field, Line, Rec0, Name, Field0}) ->
    Rec1 = expr(Rec0, []),
    Field1 = expr(Field0, []),
    {record_field, Line, Rec1, Name, Field1};
pattern({record_field, Line, Rec0, Field0}) ->
    Rec1 = expr(Rec0, []),
    Field1 = expr(Field0, []),
    {record_field, Line, Rec1, Field1};
pattern({bin, Line, Fs}) ->
    Fs2 = pattern_grp(Fs),
    {bin, Line, Fs2};
pattern({op, Line, Op, A}) ->
    {op, Line, Op, A};
pattern({op, Line, Op, L, R}) ->
    {op, Line, Op, L, R}.

pattern_grp([{bin_element, L1, E1, S1, T1} | Fs]) ->
    S2 =
        case S1 of
            default ->
                default;
            _ ->
                expr(S1, [])
        end,
    T2 =
        case T1 of
            default ->
                default;
            _ ->
                bit_types(T1)
        end,
    [{bin_element, L1, expr(E1, []), S2, T2} | pattern_grp(Fs)];
pattern_grp([]) ->
    [].

bit_types([]) ->
    [];
bit_types([Atom | Rest]) when is_atom(Atom) ->
    [Atom | bit_types(Rest)];
bit_types([{Atom, Integer} | Rest]) when
    is_atom(Atom) andalso is_integer(Integer)
->
    [{Atom, Integer} | bit_types(Rest)].

%% -type pattern_list([Pattern]) -> [Pattern].
%%  These patterns are processed "in parallel" for purposes of variable
%%  definition etc.

pattern_list([P0 | Ps]) ->
    P1 = pattern(P0),
    [P1 | pattern_list(Ps)];
pattern_list([]) ->
    [].

%% -type pattern_fields([Field]) -> [Field].
%%  N.B. Field names are full expressions here but only atoms are allowed
%%  by the *linter*!.

pattern_fields([{record_field, Lf, {atom, La, F}, P0} | Pfs]) ->
    P1 = pattern(P0),
    [{record_field, Lf, {atom, La, F}, P1} | pattern_fields(Pfs)];
pattern_fields([{record_field, Lf, {var, La, '_'}, P0} | Pfs]) ->
    P1 = pattern(P0),
    [{record_field, Lf, {var, La, '_'}, P1} | pattern_fields(Pfs)];
pattern_fields([]) ->
    [].

%% -type exprs([Expression]) -> [Expression].
%%  These expressions are processed "sequentially" for purposes of variable
%%  definition etc.

exprs([E0 | Es], MonadStack) ->
    E1 = expr(E0, MonadStack),
    [E1 | exprs(Es, MonadStack)];
exprs([], _MonadStack) ->
    [].

%% -type expr(Expression) -> Expression.

expr({var, Line, V}, _MonadStack) ->
    {var, Line, V};
expr({integer, Line, I}, _MonadStack) ->
    {integer, Line, I};
expr({float, Line, F}, _MonadStack) ->
    {float, Line, F};
expr({atom, Line, A}, _MonadStack) ->
    {atom, Line, A};
expr({string, Line, S}, _MonadStack) ->
    {string, Line, S};
expr({char, Line, C}, _MonadStack) ->
    {char, Line, C};
expr({nil, Line}, _MonadStack) ->
    {nil, Line};
expr({cons, Line, H0, T0}, MonadStack) ->
    H1 = expr(H0, MonadStack),
    %% They see the same variables
    T1 = expr(T0, MonadStack),
    {cons, Line, H1, T1};
expr({lc, Line, E0, Qs0}, MonadStack) ->
    Qs1 = lc_bc_quals(Qs0, MonadStack),
    E1 = expr(E0, MonadStack),
    {lc, Line, E1, Qs1};
expr({bc, Line, E0, Qs0}, MonadStack) ->
    Qs1 = lc_bc_quals(Qs0, MonadStack),
    E1 = expr(E0, MonadStack),
    {bc, Line, E1, Qs1};
expr({tuple, Line, Es0}, MonadStack) ->
    Es1 = expr_list(Es0, MonadStack),
    {tuple, Line, Es1};
%% OTP 17.0: EEP 443: Map construction
expr({map, Line, Fields0}, MonadStack) ->
    Fields1 = map_fields(Fields0, MonadStack),
    {map, Line, Fields1};
%% OTP 17.0: EEP 443: Map update
expr({map, Line, Expr0, Fields0}, MonadStack) ->
    Expr1 = expr(Expr0, MonadStack),
    Fields1 = map_fields(Fields0, MonadStack),
    {map, Line, Expr1, Fields1};
expr({record_index, Line, Name, Field0}, MonadStack) ->
    Field1 = expr(Field0, MonadStack),
    {record_index, Line, Name, Field1};
expr({record, Line, Name, Inits0}, MonadStack) ->
    Inits1 = record_inits(Inits0, MonadStack),
    {record, Line, Name, Inits1};
expr({record_field, Line, Rec0, Name, Field0}, MonadStack) ->
    Rec1 = expr(Rec0, MonadStack),
    Field1 = expr(Field0, MonadStack),
    {record_field, Line, Rec1, Name, Field1};
expr({record, Line, Rec0, Name, Upds0}, MonadStack) ->
    Rec1 = expr(Rec0, MonadStack),
    Upds1 = record_updates(Upds0, MonadStack),
    {record, Line, Rec1, Name, Upds1};
expr({record_field, Line, Rec0, Field0}, MonadStack) ->
    Rec1 = expr(Rec0, MonadStack),
    Field1 = expr(Field0, MonadStack),
    {record_field, Line, Rec1, Field1};
expr({block, Line, Es0}, MonadStack) ->
    %% Unfold block into a sequence.
    Es1 = exprs(Es0, MonadStack),
    {block, Line, Es1};
expr({'if', Line, Cs0}, MonadStack) ->
    Cs1 = icr_clauses(Cs0, MonadStack),
    {'if', Line, Cs1};
expr({'case', Line, E0, Cs0}, MonadStack) ->
    E1 = expr(E0, MonadStack),
    Cs1 = icr_clauses(Cs0, MonadStack),
    {'case', Line, E1, Cs1};
expr({'receive', Line, Cs0}, MonadStack) ->
    Cs1 = icr_clauses(Cs0, MonadStack),
    {'receive', Line, Cs1};
expr({'receive', Line, Cs0, To0, ToEs0}, MonadStack) ->
    To1 = expr(To0, MonadStack),
    ToEs1 = exprs(ToEs0, MonadStack),
    Cs1 = icr_clauses(Cs0, MonadStack),
    {'receive', Line, Cs1, To1, ToEs1};
expr({'try', Line, Es0, Scs0, Ccs0, As0}, MonadStack) ->
    Es1 = exprs(Es0, MonadStack),
    Scs1 = icr_clauses(Scs0, MonadStack),
    Ccs1 = icr_clauses(Ccs0, MonadStack),
    As1 = exprs(As0, MonadStack),
    {'try', Line, Es1, Scs1, Ccs1, As1};
expr({'fun', Line, Body}, MonadStack) ->
    case Body of
        {clauses, Cs0} ->
            Cs1 = fun_clauses(Cs0, MonadStack),
            {'fun', Line, {clauses, Cs1}};
        {function, F, A} ->
            {'fun', Line, {function, F, A}};
        %% R10B-6: fun M:F/A.
        {function, M, F, A} ->
            {'fun', Line, {function, M, F, A}}
    end;
%% OTP 17.0: EEP 37: Funs with names
expr({named_fun, Line, Name, Cs0}, MonadStack) ->
    Cs1 = fun_clauses(Cs0, MonadStack),
    {named_fun, Line, Name, Cs1};
%%  do syntax detection:
expr(
    {call, Line, {atom, _Line1, do}, [{lc, _Line2, {AtomOrVar, _Line3, _MonadModule} = Monad, Qs}]},
    MonadStack
) when AtomOrVar =:= atom orelse AtomOrVar =:= var ->
    %% 'do' calls of a particular form:
    %%  do([ MonadMod || Qualifiers ])
    {call, Line,
        {'fun', Line, {clauses, [{clause, Line, [], [], do_syntax(Qs, [Monad | MonadStack])}]}},
        []};
%%  'return' and 'fail' syntax detection and transformation:
expr(
    {call, Line, {atom, Line1, ReturnOrFail}, As0},
    [Monad | _Monads] = MonadStack
) when
    ReturnOrFail =:= return orelse
        ReturnOrFail =:= fail
->
    %% 'return' calls of a particular form:
    %%  return(Arguments), and
    %% 'fail' calls of a particular form:
    %%  fail(Arguments)
    %% Transformed to:
    %% "Monad:return(Args)" or "Monad:fail(Args)" in monadic context
    {call, Line, {remote, Line1, Monad, {atom, Line1, ReturnOrFail}}, expr_list(As0, MonadStack)};
expr({call, Line, F0, As0}, MonadStack) ->
    %% N.B. If F an atom then call to local function or BIF,  if F a
    %% remote structure (see below) then call to other module,
    %% otherwise apply to "function".
    F1 = expr(F0, MonadStack),
    As1 = expr_list(As0, MonadStack),
    {call, Line, F1, As1};
expr({'catch', Line, E0}, MonadStack) ->
    %% No new variables added.
    E1 = expr(E0, MonadStack),
    {'catch', Line, E1};
expr({'query', Line, E0}, MonadStack) ->
    %% lc expression
    E = expr(E0, MonadStack),
    {'query', Line, E};
expr({match, Line, P0, E0}, MonadStack) ->
    E1 = expr(E0, MonadStack),
    P1 = pattern(P0),
    {match, Line, P1, E1};
expr({bin, Line, Fs}, _MonadStack) ->
    Fs2 = pattern_grp(Fs),
    {bin, Line, Fs2};
expr({op, Line, Op, A0}, MonadStack) ->
    A1 = expr(A0, MonadStack),
    {op, Line, Op, A1};
expr({op, Line, Op, L0, R0}, MonadStack) ->
    L1 = expr(L0, MonadStack),
    %% They see the same variables
    R1 = expr(R0, MonadStack),
    {op, Line, Op, L1, R1};
%% The following are not allowed to occur anywhere!
expr({remote, Line, M0, F0}, MonadStack) ->
    M1 = expr(M0, MonadStack),
    F1 = expr(F0, MonadStack),
    {remote, Line, M1, F1}.

%% -type expr_list([Expression]) -> [Expression].
%%  These expressions are processed "in parallel" for purposes of variable
%%  definition etc.

expr_list([E0 | Es], MonadStack) ->
    E1 = expr(E0, MonadStack),
    [E1 | expr_list(Es, MonadStack)];
expr_list([], _MonadStack) ->
    [].

%% -type map_fields([MapField]) -> [MapField].
map_fields([{map_field_assoc, Line, ExpK0, ExpV0} | Fs], MonadStack) ->
    ExpK1 = expr(ExpK0, MonadStack),
    ExpV1 = expr(ExpV0, MonadStack),
    [{map_field_assoc, Line, ExpK1, ExpV1} | map_fields(Fs, MonadStack)];
map_fields([{map_field_exact, Line, ExpK0, ExpV0} | Fs], MonadStack) ->
    ExpK1 = expr(ExpK0, MonadStack),
    ExpV1 = expr(ExpV0, MonadStack),
    [{map_field_exact, Line, ExpK1, ExpV1} | map_fields(Fs, MonadStack)];
map_fields([], _MoandStack) ->
    [].

%% -type record_inits([RecordInit]) -> [RecordInit].
%%  N.B. Field names are full expressions here but only atoms are allowed
%%  by the *linter*!.

record_inits([{record_field, Lf, {atom, La, F}, Val0} | Is], MonadStack) ->
    Val1 = expr(Val0, MonadStack),
    [{record_field, Lf, {atom, La, F}, Val1} | record_inits(Is, MonadStack)];
record_inits([{record_field, Lf, {var, La, '_'}, Val0} | Is], MonadStack) ->
    Val1 = expr(Val0, MonadStack),
    [{record_field, Lf, {var, La, '_'}, Val1} | record_inits(Is, MonadStack)];
record_inits([], _MonadStack) ->
    [].

%% -type record_updates([RecordUpd]) -> [RecordUpd].
%%  N.B. Field names are full expressions here but only atoms are allowed
%%  by the *linter*!.

record_updates([{record_field, Lf, {atom, La, F}, Val0} | Us], MonadStack) ->
    Val1 = expr(Val0, MonadStack),
    [{record_field, Lf, {atom, La, F}, Val1} | record_updates(Us, MonadStack)];
record_updates([], _MonadStack) ->
    [].

%% -type icr_clauses([Clause]) -> [Clause].

icr_clauses([C0 | Cs], MonadStack) ->
    C1 = clause(C0, MonadStack),
    [C1 | icr_clauses(Cs, MonadStack)];
icr_clauses([], _MonadStack) ->
    [].

%% -type lc_bc_quals([Qualifier]) -> [Qualifier].
%%  Allow filters to be both guard tests and general expressions.

lc_bc_quals([{generate, Line, P0, E0} | Qs], MonadStack) ->
    E1 = expr(E0, MonadStack),
    P1 = pattern(P0),
    [{generate, Line, P1, E1} | lc_bc_quals(Qs, MonadStack)];
lc_bc_quals([{b_generate, Line, P0, E0} | Qs], MonadStack) ->
    E1 = expr(E0, MonadStack),
    P1 = pattern(P0),
    [{b_generate, Line, P1, E1} | lc_bc_quals(Qs, MonadStack)];
lc_bc_quals([E0 | Qs], MonadStack) ->
    E1 = expr(E0, MonadStack),
    [E1 | lc_bc_quals(Qs, MonadStack)];
lc_bc_quals([], _MonadStack) ->
    [].

%% -type fun_clauses([Clause]) -> [Clause].

fun_clauses([C0 | Cs], MonadStack) ->
    C1 = clause(C0, MonadStack),
    [C1 | fun_clauses(Cs, MonadStack)];
fun_clauses([], _MonadStack) ->
    [].

%%  'do' syntax transformation:
do_syntax([], [{_AtomOrVar, MLine, _MonadModule} | _MonadStack]) ->
    transform_error("A 'do' construct cannot be empty", MLine);
do_syntax([{GenerateOrMatch, Line, _Pattern, _Expr}], _MonadStack) when
    GenerateOrMatch =:= generate orelse GenerateOrMatch =:= match
->
    transform_error("The last statement in a 'do' construct must be an expression", Line);
do_syntax(
    [{generate, Line, {var, _Line, _Var} = Pattern, Expr} | Exprs],
    [Monad | _Monads] = MonadStack
) ->
    %% "Pattern <- Expr, Tail" where Pattern is a simple variable
    %% is transformed to
    %% "Monad:'>>='(Expr, fun (Pattern) -> Tail')"
    %% without a fail to match clause
    [
        {call, Line, {remote, Line, Monad, {atom, Line, '>>='}}, [
            expr(Expr, MonadStack),
            {'fun', Line, {clauses, [{clause, Line, [Pattern], [], do_syntax(Exprs, MonadStack)}]}}
        ]}
    ];
do_syntax(
    [{generate, Line, Pattern, Expr} | Exprs],
    [Monad | _Monads] = MonadStack
) ->
    %% "Pattern <- Expr, Tail" where Pattern is not a simple variable
    %% is transformed to
    %% "Monad:'>>='(Expr, fun (Pattern) -> Tail')"
    %% with a fail clause if the function does not match
    [
        {call, Line, {remote, Line, Monad, {atom, Line, '>>='}}, [
            expr(Expr, MonadStack),
            {'fun', Line,
                {clauses, [
                    {clause, Line, [Pattern], [], do_syntax(Exprs, MonadStack)},
                    {clause, Line, [{var, Line, '_'}], [], [
                        {call, Line, {remote, Line, Monad, {atom, Line, 'fail'}}, [
                            {atom, Line, 'monad_badmatch'}
                        ]}
                    ]}
                ]}}
        ]}
    ];
do_syntax([Expr], MonadStack) ->
    %% Don't do '>>' chaining on the last elem
    [expr(Expr, MonadStack)];
do_syntax(
    [{match, _Line, _Pattern, _Expr} = Expr | Exprs],
    MonadStack
) ->
    %% Handles 'let binding' in do expression a-la Haskell
    [expr(Expr, MonadStack) | do_syntax(Exprs, MonadStack)];
do_syntax([Expr | Exprs], [Monad | _Monads] = MonadStack) ->
    %% "Expr, Tail" is transformed to "Monad:'>>='(Expr, fun (_) -> Tail')"
    %% Line is always the 2nd element of Expr
    Line = element(2, Expr),
    [
        {call, Line, {remote, Line, Monad, {atom, Line, '>>='}}, [
            expr(Expr, MonadStack),
            {'fun', Line,
                {clauses, [{clause, Line, [{var, Line, '_'}], [], do_syntax(Exprs, MonadStack)}]}}
        ]}
    ].

%% Use this function to report any parse_transform error. The
%% resulting error message will be displayed as an ordinary
%% compilation error in a standard format.
transform_error(Message, Line) ->
    throw({Message, Line}).

%% This function is called by the Erlang compiler to obtain an error
%% message which will be shown to the user.
format_error(Message) ->
    case io_lib:deep_char_list(Message) of
        true -> Message;
        _ -> io_lib:write(Message)
    end.