Add functionality to get roots of tracebacks

There are two options:

1) TBs = tr:tracebacks(..., #{format => tree}),
   tr:roots(TBs)
   % alternatively: tb:root(TB) on a single traceback

2) tr:tracebacks(..., #{format => root})

Author: chrzaszcz

src/tr.erl

@@ -22,6 +22,7 @@
          filter/1, filter/2,
          traceback/1, traceback/2,
          tracebacks/1, tracebacks/2,
+         roots/1, root/1,
          range/1, range/2,
          ranges/1, ranges/2,
          call_tree_stat/0, call_tree_stat/1,
@@ -162,14 +163,18 @@
 -type tb_output() :: shortest | longest | all.
 %% Which tracebacks to return if they overlap. Default: `shortest'.
 
--type tb_format() :: tree | list.
-%% Merge tracebacks into a `tree' or return a `list' (default) of them.
+-type tb_format() :: list | tree | root.
+%% Format in which tracebacks are returned.
+%%
+%% `list' (default) returns a list of tracebacks.
+%% `tree' merges them into a list of trees.
+%% `root' returns only the root of each tree.
 
 -type tb_order() :: top_down | bottom_up.
 %% Order of calls in each returned traceback. Default: `top_down'.
 
--type tb_tree() :: [tr() | {tr(), tb_tree()}].
-%% Multiple tracebacks merged into a tree structure.
+-type tb_tree() :: tr() | {tr(), [tb_tree()]}.
+%% Multiple tracebacks with a common root merged into a tree structure.
 
 -type tb_acc_tree() :: [{tr(), tb_acc_tree()}].
 -type tb_acc_list() :: [[tr()]].
@@ -379,12 +384,12 @@ traceback(PredF, Options) when is_function(PredF, 1) ->
 %% @doc Returns tracebacks of all matching traces from `tab()'.
 %%
 %% @see tracebacks/2
--spec tracebacks(pred()) -> [[tr()]] | tb_tree().
+-spec tracebacks(pred()) -> [[tr()]] | [tb_tree()].
 tracebacks(PredF) ->
     tracebacks(PredF, #{}).
 
 %% @doc Returns tracebacks of all matching traces from `t:tr_source()'.
--spec tracebacks(pred(), tb_options()) -> [[tr()]] | tb_tree().
+-spec tracebacks(pred(), tb_options()) -> [[tr()]] | [tb_tree()].
 tracebacks(PredF, Options) when is_map(Options) ->
     Tab = maps:get(tab, Options, tab()),
     Output = maps:get(output, Options, shortest),
@@ -397,6 +402,16 @@ tracebacks(PredF, Options) when is_map(Options) ->
         foldl(fun(T, State) -> tb_step(PredF, T, State) end, InitialState, Tab),
     finalize_tracebacks(TBs, Output, Format, Options).
 
+%% @doc Returns the root call of each `t:tb_tree()` from the provided list.
+-spec roots([tb_tree()]) -> [tr()].
+roots(Trees) ->
+    lists:map(fun root/1, Trees).
+
+%% @doc Returns the root call of the provided `t:tb_tree()`.
+-spec root(tb_tree()) -> tr().
+root(#tr{} = T) -> T;
+root({#tr{} = T, _}) -> T.
+
 %% @doc Returns a list of traces from `tab()' between the first matched call and the corresponding return.
 %%
 %% Matching can be done with a predicate function, an index value or a `t:tr()' record.
@@ -827,6 +842,8 @@ finalize_tracebacks(TBs, all, list, Options) ->
     reorder_tb(lists:reverse(TBs), maps:get(order, Options, top_down));
 finalize_tracebacks(TBs, _, list, Options) ->
     reorder_tb(tree_to_list(TBs), maps:get(order, Options, top_down));
+finalize_tracebacks(TBs, _, root, _Options) ->
+    lists:map(fun({T, _}) -> T end, lists:reverse(TBs));
 finalize_tracebacks(TBs, _, tree, _Options) ->
     finalize_tree(TBs).
 
@@ -837,7 +854,8 @@ reorder_tb(TBs, bottom_up) -> TBs.
 -spec tree_to_list(tb_acc_tree()) -> [[tr()]].
 tree_to_list(Tree) ->
     lists:foldl(fun({K, []}, Res) -> [[K] | Res];
-                   ({K, V}, Res) -> [[K | Rest] || Rest <- tree_to_list(V)] ++ Res end, [], Tree).
+                   ({K, V}, Res) -> [[K | Rest] || Rest <- tree_to_list(V)] ++ Res
+                end, [], Tree).
 
 %% Reverse order and simplify leaf nodes
 -spec finalize_tree(tb_acc_tree()) -> tb_tree().

test/tr_SUITE.erl

@@ -56,7 +56,8 @@ groups() ->
                   tb_all,
                   tb_all_limit,
                   tb_tree,
-                  tb_tree_longest]},
+                  tb_tree_longest,
+                  tb_roots]},
      {util, [do]},
      {call_stat, [simple_total,
                   tree_total,
@@ -429,6 +430,24 @@ tb_tree_longest(_Config) ->
                                      {#tr{data = [2]}, [#tr{data = [1]},
                                                         #tr{data = [0]}]}]}], TBs).
 
+tb_roots(_Config) ->
+    tr:trace([{?MODULE, fib, 1}]),
+    ?MODULE:fib(4),
+    wait_for_traces(18),
+    tr:stop_tracing(),
+
+    %% Option 1: call root/1 or root/2 on a tree
+    TBs = tr:tracebacks(fun(#tr{event = return, data = N}) when N < 2 -> true end,
+                               #{format => tree}),
+    Roots = tr:roots(TBs),
+    ct:pal("Roots: ~p~n", [Roots]),
+    ?assertMatch([#tr{data = [4]}], Roots),
+    ?assertEqual(hd(Roots), tr:root(hd(TBs))),
+
+    %% Option 2: directly use the root format
+    ?assertEqual(Roots, tr:tracebacks(fun(#tr{event = return, data = N}) when N < 2 -> true end,
+                                      #{format => root})).
+
 simple_total(_Config) ->
     tr:trace([{?MODULE, factorial, 1}]),
     ?MODULE:factorial(2),