timeit

-module(timeit).

-export([timeit/3,
         timeit/4,
         timeit/5]).

-export([simple_test/0]).

timeit(Mod, Fun, Arity) ->
    timeit(all, Mod, Fun, Arity, undefined).

timeit(Mod, Fun, Arity, CallFun) ->
    timeit(all, Mod, Fun, Arity, CallFun).

timeit(Pid, Mod, Fun, Arity, CallFun) ->
    {ok, TracePid} = dbg:tracer(process, {fun trace/2, {CallFun, []}}),
    io:format("Tracing ~p:~p/~p on ~p\n", [Mod, Fun, Arity, TracePid]),
    dbg:p(Pid, call),
    dbg:tpl(Mod, Fun, Arity, [{'_', [], [{return_trace}]}]).


trace({trace, Pid, call, {Mod, Fun, _}}, {CallFun, Acc0}) ->
    Acc = orddict:store({Pid, Mod, Fun}, os:timestamp(), Acc0),
    {CallFun, Acc};
trace({trace, Pid, return_from, {Mod, Fun, Args}, FunRes}, {CallFun, Acc}) ->
    case orddict:find({Pid, Mod, Fun}, Acc) of
        {ok, StartTime} ->
            ElapsedUs = timer:now_diff(os:timestamp(), StartTime),
            render(Pid, Mod, Fun, Args, FunRes, ElapsedUs, CallFun);
        error ->
            ok
    end,
    {CallFun, Acc}.

render(SourcePid, Mod, Fun, _Args, _FunRes, ElapsedUs, undefined) ->
    io:format(user, "~p:~p:~p:~p us\n", [SourcePid, Mod, Fun, ElapsedUs]);
render(SourcePid, Mod, Fun, Args, FunRes, ElapsedUs, CallFun) when is_function(CallFun) ->
    Res = CallFun(Mod, Fun, Args, FunRes),
    io:format(user, "~p:~p:~p: ~p us: ~p\n", [SourcePid, Mod, Fun, ElapsedUs, Res]).


simple_test() ->
    F = fun(Mod, Fun, Arg, Result) -> {Mod, Fun, Arg, Result} end,
    timeit(timer, sleep, 1, F),
    timer:sleep(100),
    dbg:stop_clear().