doctest

%%%---------------------------------------------------------------------
%%% Copyright 2024 William Fank Thomé
%%%
%%% Licensed under the Apache License, Version 2.0 (the "License");
%%% you may not use this file except in compliance with the License.
%%% You may obtain a copy of the License at
%%%
%%%     http://www.apache.org/licenses/LICENSE-2.0
%%%
%%% Unless required by applicable law or agreed to in writing, software
%%% distributed under the License is distributed on an "AS IS" BASIS,
%%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
%%% See the License for the specific language governing permissions and
%%% limitations under the License.
%%%---------------------------------------------------------------------
-module(doctest_parse_transform).
-moduledoc """
Generate tests for -doc attributes via parse transform.

Just plug the header on your module:
```
-ifdef(TEST).
-include_lib("doctest/include/doctest.hrl").
-endif.
```
""".
-moduledoc #{ author => "William Fank Thomé [https://github.com/williamthome]" }.

% Check OTP version >= 27.
-include("doctest_otp_check.hrl").

% API functions
-export([parse_transform/2]).

% Settings that the user can override
-record(doctest, {enabled, funs}).

%%%=====================================================================
%%% API functions
%%%=====================================================================

parse_transform(Forms, _Opt) ->
    File = file(Forms),
    Docs = docs(doctest(doctest_attrs(Forms), File), doc_attrs(Forms)),
    doctest:run(tests(File, Forms, Docs)),
	Forms.

%%%=====================================================================
%%% Internal functions
%%%=====================================================================

docs(Doctest, AllDocs) ->
    case Doctest#doctest.funs of
        all ->
            AllDocs;
        Funs when is_list(Funs) ->
            maps:with(Funs, AllDocs)
    end.

doctest(Attrs, SrcFile) ->
    parse_to_doctest(Attrs, SrcFile, #doctest{
        enabled = true,
        funs = all
    }).

parse_to_doctest([Enabled|T], SrcFile, DT) when is_boolean(Enabled) ->
    parse_to_doctest(T, SrcFile, DT#doctest{enabled = Enabled});
parse_to_doctest([Funs|T], SrcFile, DT) when is_list(Funs) ->
    parse_to_doctest(T, SrcFile, DT#doctest{funs = Funs});
parse_to_doctest([all|T], SrcFile, DT) ->
    parse_to_doctest(T, SrcFile, DT#doctest{funs = all});
parse_to_doctest([Map|T], SrcFile, DT) when is_map(Map) ->
    parse_to_doctest(T, SrcFile, DT#doctest{
        enabled = maps:get(enabled, Map, DT#doctest.enabled),
        funs = maps:get(funs, Map, DT#doctest.funs)
    });
parse_to_doctest([], _SrcFile, #doctest{} = DT) ->
    DT.

tests(File, Forms, Docs) ->
    maps:fold(fun({F, A, Ln}, Md, Acc) ->
        case doctest:code_block(Md) of
            {ok, CodeBlock} ->
                {ok, M, Bin} = compile:forms(Forms, [
                    {i, "eunit/include/eunit.hrl"}
                ]),
                {module, M} = code:load_binary(M, File, Bin),
                [doctest:parse({M, F, A}, Ln, doctest:chunk(CodeBlock)) | Acc];
            none ->
                Acc
        end
    end, [], Docs).

file(Forms) ->
    [File, _Loc] = hd(attributes(file, Forms)),
    File.

doctest_attrs(Forms) ->
    attributes(doctest, Forms).

doc_attrs(Forms) ->
    do_doc_attrs(filtermap_forms(
        fun(Type) -> lists:member(Type, [attribute, function]) end,
        fun
            ({attribute, Attr}) ->
                case attribute_name(Attr) =:= doc of
                    true ->
                        {true, {doc, normalize_attribute(Attr)}};
                    false ->
                        false
                end;
            ({function, Fun}) ->
                {true, {function, normalize_function(Fun)}}
        end,
        Forms
    ), []).

do_doc_attrs([{doc, Doc},{function, Fun}|T], Acc) ->
    do_doc_attrs(T, [{Fun, Doc} | Acc]);
do_doc_attrs([{function, _}|T], Acc) ->
    do_doc_attrs(T, Acc);
do_doc_attrs([], Acc) ->
    maps:from_list(Acc).

attributes(Name, Forms) ->
    filtermap_forms(
        fun(Type) -> Type =:= attribute end,
        fun({attribute, Attr}) ->
            case attribute_name(Attr) =:= Name of
                true ->
                    {true, normalize_attribute(Attr)};
                false ->
                    false
            end
        end,
        Forms
    ).

filtermap_forms(Validate, Filtermap, Forms) when
    is_function(Validate, 1), is_function(Filtermap, 1), is_list(Forms) ->
    lists:filtermap(fun(Form) ->
        Type = erl_syntax:type(Form),
        case Validate(Type) of
            true ->
                Filtermap({Type, Form});
            false ->
                false
        end
    end, Forms).

attribute_name(Attr) ->
    erl_syntax:atom_value(erl_syntax:attribute_name(Attr)).

normalize_attribute(Attr) ->
    Exprs = erl_syntax:revert_forms(erl_syntax:attribute_arguments(Attr)),
    case lists:map(fun(Expr) ->
        {value, Value, []} = erl_eval:expr(Expr, []),
        Value
    end, Exprs) of
        [H] ->
            H;
        List ->
            List
    end.

normalize_function(Fun) ->
    {
        erl_syntax:atom_value(erl_syntax:function_name(Fun)),
        erl_syntax:function_arity(Fun),
        erl_anno:line(erl_syntax:get_pos(Fun))
    }.