%%%-------------------------------------------------------------------
%%% @doc
%%% A parse transform implementing partial function application.
%%% @end
%%%-------------------------------------------------------------------
-module(partial).

% -define(PARTIAL_DEBUG, true).

-ifdef(PARTIAL_DEBUG).
-define(IF_DEBUG(Expression), Expression).
-else.
-define(IF_DEBUG(Expression), (ok)).
-endif.

%% API
-export([cut/1,
         cute/1,
         parse_transform/2]).

-record(cut, {variables=[], arguments=[]}).
-record(cute, {variables=[], matches=[], arguments=[]}).

-include_lib("syntax_tools/include/merl.hrl").

%%%===================================================================
%%% API
%%%===================================================================

%% @doc
%% A dummy function used as a marker by parse_transform/2 to convert
%% calls to functions to partially applied functions. The special
%% variable '_' is used as a marker for unevaluated arguments, as it
%% is usually illegal to use on the right hand side of a match.
%%
%% All arguments are evaluated when the partially applied function is
%% called.
%%
%% The parse transform is only able to detect and rewrite simple
%% literal calls to this function. Other uses will result in an error
%% being thrown at runtime.
%% @end
%% @see parse_transform/2
cut(_Fun) ->
    missing_parse_transform().

%% @doc
%% A dummy function used as a marker by parse_transform/2 to convert
%% calls to functions to partially applied functions. The special
%% variable '_' is used as a marker for unevaluated arguments, as it
%% is usually illegal to use on the right hand side of a match.
%%
%% Given arguments are evaluated when the partially applied function
%% is constructed. This can be used as an easy way to cache expensive
%% computation in a closure.
%%
%% The parse transform is only able to detect and rewrite simple
%% literal calls to this function. Other uses will result in an error
%% being thrown at runtime.
%% @end
%% @see parse_transform/2
cute(_Fun) ->
    missing_parse_transform().

%% @doc
%% A parse transformation function which converts calls to special
%% dummy functions in this module.
%%
%% Add:
%% -compile({parse_transform, partial}).
%%
%% to the top of any module to enable.
%% @end
parse_transform(Forms, _Options) ->
    transform_forms(Forms).

%%%===================================================================
%%% Internal Functions
%%%===================================================================

missing_parse_transform() ->
    throw({missing_parse_transform,
           "This function requires that the partial:parse_trans/2 "
           "function be listed as a parse transformation for your "
           "module as \"-compile({parse_transform, partial}).\" "
           "If that is already true, this function has been called "
           "indirectly in a way the parse transform can not "
           "recognize, such as via apply/3."}).

transform_forms(Forms) ->
    % If any errors occured, we replace the entire form with them.
    % This allows the standard compiler chain to pick them up as if
    % they were erl_parse errors and display a standard error message
    % to the user instead of a traceback.
    erl_syntax:revert_forms(
      lists:append(
          [case erl_syntax_lib:mapfold(fun transform/2, [], Form) of
               {Transformed, []} ->
                   [Transformed];
               {_Transformed, Errors} ->
                   Errors
           end || Form <- Forms])).

transform(Form, Errors) ->
    Line = erl_syntax:get_pos(Form),
    Transformed = case Form of
        ?Q("partial:cut(_@@Args)") ->
            ?IF_DEBUG(ok = io:format(
                             "partial:parse_transform/2 cut Args:~n~p~n",
                             [Args])),
            Cut = cut_function(Line, Args),
            ?IF_DEBUG(ok = io:format(
                             "partial:parse_transform/2 cut Transform:~n~p~n",
                             [Cut])),
            Cut;
        ?Q("partial:cute(_@@Args)") ->
            ?IF_DEBUG(ok = io:format(
                             "partial:parse_transform/2 cute Args:~n~p~n",
                             [Args])),
            Cute = cute_function(Line, Args),
            ?IF_DEBUG(ok = io:format(
                             "partial:parse_transform/2 cute Transform:~n~p~n",
                             [Cute])),
            Cute;
        _ ->
            {ok, Form}
    end,
    case Transformed of
        {ok, NewForm} ->
            {NewForm, Errors};
        {error, Error} ->
            {Form, [Error | Errors]}
    end.

transform_error(Line, Message) ->
    % I'm probably abusing the error system by using erl_parse here,
    % but it allows my errors to show up as usual in the compiler
    % output.
    Prefix = "Error: In partial:parse_transform/2, ",
    {error, erl_syntax:error_marker({Line, erl_parse, Prefix ++ Message})}.

reverse(#cut{} = Cut) ->
    #cut{variables=lists:reverse(Cut#cut.variables),
         arguments=lists:reverse(Cut#cut.arguments)};
reverse(#cute{} = Cute) ->
    #cute{variables=lists:reverse(Cute#cute.variables),
          matches=lists:reverse(Cute#cute.matches),
          arguments=lists:reverse(Cute#cute.arguments)}.

is_cut_variable(Var) ->
    case Var of
        ?Q("_") ->
            true;
        _ ->
            false
    end.

variable_name(Type) ->
    Name = io_lib:format("PartialArgument_~s_~w", [Type, make_ref()]),
    erlang:list_to_atom(lists:flatten(Name)).

variable(Line, Type) ->
    erl_syntax:set_pos(erl_syntax:variable(variable_name(Type)), Line).

name([Name]) ->
    [Name];
name([Module, Function]) ->
    ?Q("_@Module:_@Function").

split_name(Name) ->
    case Name of
        ?Q("_@Module:_@Function") ->
            [Module, Function];
        _ ->
            [Name]
    end.

match(Line, Type, Form) ->
    Pattern = variable(Line, Type),
    Match = merl:qquote(Line,
                        "_@pattern = _@form",
                        [{pattern, Pattern}, {form, Form}]),
    {Pattern, Match}.

cuts(Line, Type, Forms) ->
    reverse(lists:foldl(
      fun (Form, #cut{} = Acc) ->
          case is_cut_variable(Form) of
              true ->
                  Variable = variable(Line, Type),
                  Acc#cut{variables=[Variable | Acc#cut.variables],
                          arguments=[Variable | Acc#cut.arguments]};
              false ->
                  Acc#cut{arguments=[Form | Acc#cut.arguments]}
          end
      end,
      #cut{},
      Forms)).

name_cuts(Line, Name) ->
    Parts = split_name(Name),
    Cut = cuts(Line, name, Parts),
    Cut#cut{arguments=name(Cut#cut.arguments)}.

cutes(Line, Type, Forms) ->
    reverse(lists:foldl(
      fun (Form, #cute{} = Acc) ->
          case {is_cut_variable(Form), erl_syntax:is_literal(Form)} of
              {true, false} ->
                  Variable = variable(Line, Type),
                  Acc#cute{variables=[Variable | Acc#cute.variables],
                           arguments=[Variable | Acc#cute.arguments]};
              {false, true} ->
                  Acc#cute{arguments=[Form | Acc#cute.arguments]};
              {false, false} ->
                  {Variable, Match} = match(Line, Type, Form),
                  Acc#cute{matches=[Match | Acc#cute.matches],
                           arguments=[Variable | Acc#cute.arguments]}
          end
      end,
      #cute{},
      Forms)).

name_cutes(Line, Name) ->
    Parts = split_name(Name),
    Cute = cutes(Line, name, Parts),
    Cute#cute{arguments=name(Cute#cute.arguments)}.

%% @doc
%% Transform the AST for:
%% 
%% F = partial:cut(f(_, x, y, z)).
%%
%% into:
%%
%% F = fun (Arg1) ->
%%     f(Arg1, x, y, z)
%% end.
%% @end
cut_function(MarkerLine, [MarkerArgument])->
    case MarkerArgument of
        ?Q("_@Name(_@@Args)") ->
            Line = erl_syntax:get_pos(MarkerArgument),
            NameCut = name_cuts(Line, Name),
            Cut = cuts(Line, cut, Args),
            CutFun = merl:qquote(Line,
                                 "fun (_@@variables) ->"
                                 " _@name(_@@arguments)"
                                 " end",
                                 [{variables, NameCut#cut.variables ++ Cut#cut.variables},
                                  {name, NameCut#cut.arguments},
                                  {arguments, Cut#cut.arguments}]),
            {ok, CutFun};
        _ ->
            transform_error(
              MarkerLine,
              "partial:cut/1 requires a function call as an argument")
    end;
cut_function(MarkerLine, MarkerArguments) ->
    transform_error(
      MarkerLine,
      io_lib:format(
        "partial:cut/1 requires a single argument, got ~b",
        [length(MarkerArguments)])).

%% @doc
%% Transform the AST for:
%% 
%% F = partial:cute(f(_, x, y, z)).
%%
%% into:
%%
%% F = (fun () ->
%%      Arg2 = x,
%%      Arg3 = y,
%%      Arg4 = z,
%%      fun (Arg1) ->
%%          f(Arg1, Arg2, Arg3, Arg4)
%%      end
%% end)().
%% @end
cute_function(MarkerLine, [MarkerArgument])->
    case MarkerArgument of
        ?Q("_@Name(_@@Args)") ->
            Line = erl_syntax:get_pos(MarkerArgument),
            NameCute = name_cutes(Line, Name),
            Cute = cutes(Line, cute, Args),
            CuteFun = merl:qquote(Line,
                                  "(fun () ->"
                                  " _@@matches,"
                                  " fun (_@@variables) -> _@name(_@@arguments) end"
                                  " end)()",
                                  [{matches, NameCute#cute.matches ++ Cute#cute.matches},
                                   {variables, NameCute#cute.variables ++ Cute#cute.variables},
                                   {name, NameCute#cute.arguments},
                                   {arguments, Cute#cute.arguments}]),
            {ok, CuteFun};
        _ ->
            transform_error(
              MarkerLine,
              "partial:cute/1 requires a function call as an argument")
    end;
cute_function(MarkerLine, MarkerArguments)->
    transform_error(
      MarkerLine,
      io_lib:format(
        "partial:cute/1 requires a single argument, got ~b",
        [length(MarkerArguments)])).