gen_fsm_visualizer.erl

%% ---------------------------------------------------------------------
%%
%% Copyright (c) 2007-2015 Basho Technologies, Inc.  All Rights Reserved.
%%
%% This file is provided to you 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.
%%
%% ---------------------------------------------------------------------
%%
%% gen_fsm Visualizer - requires graphviz to visuaize
%% usage:
%% $ escript gen_fsm_visualizer.erl riak_cs_gc_manager.erl > test.dot
%% $ dot -Tpng -o test.png test.dot
%% $ open test.png

-module(gen_fsm_visualizer).

-mode(compile).
-export([main/1]).

-type edge() :: {From::atom(), To::atom(), Msg::term()}.

-record(fsm, {
          fun_index = [] :: [erl_syntax:syntaxTree()],
          forms = [] :: [erl_syntax:syntaxTree()],
          statenames = [] :: [atom()],
          modname = undefined,
          exports = [] :: [{atom(), non_neg_integer()}],
          filename = "" :: string(),
          edges = [] :: [edge()]
         }).

-type fsm() :: #fsm{}.

-define(CONSOLE(Fmt, Term),
        begin
            io:format(standard_error, Fmt, Term)
        end).

-define(STDOUT(Fmt, Term),
        begin
            io:format(standard_io, Fmt, Term)
        end).

-define(ANYSTATE, '__anystate__').
-define(INIT, '__init__').
-define(STOP, '__stop__').

main([File]) ->
    ?CONSOLE("parsing ~s~n", [File]),
    {ok, Forms} = epp_dodger:parse_file(File),
    %% io:format("~p~n", [length(Forms)]),
    AnalyzedForms = erl_syntax_lib:analyze_forms(Forms),
    ModName = proplists:get_value(module, AnalyzedForms),
    Attributes = proplists:get_value(attributes, AnalyzedForms),
    Behaviour = proplists:get_value(behaviour, Attributes),
    ?CONSOLE("behaviour of ~s: ~s~n", [ModName, Behaviour]),
    Fsm0 = #fsm{filename = File, modname=ModName, forms=Forms},
    case Behaviour of
        gen_fsm ->
            Fsm1 = extract_statem(Fsm0),
            {Transitions, Fsm2} = repeat_analysis(Fsm1),
            ?CONSOLE("State Transitions:~n ~p~n", [Transitions]),
            #fsm{statenames=StateNames} = Fsm2,
            render_graph(ModName, StateNames, Transitions);
        _ ->
            ?CONSOLE("~s does not have gen_fsm behaviour.~n", [ModName])
    end.

repeat_analysis(#fsm{statenames=StateNames} = Fsm) ->
    Transitions = analyze(Fsm),
    NewStateNames =
        lists:usort(lists:flatten([[From, To] || {From, To, _} <- Transitions])),
    case NewStateNames -- StateNames of
        [] ->
            {Transitions, Fsm};
        _More ->
            repeat_analysis(Fsm#fsm{statenames=NewStateNames})
    end.

-spec extract_statem(fsm()) -> fsm().
extract_statem(#fsm{forms=Forms, modname=ModName} = Fsm0) ->
    AnalyzedForms = erl_syntax_lib:analyze_forms(Forms),
    Exports = proplists:get_value(exports, AnalyzedForms),
    States = extract_states(Exports, []),
    ?CONSOLE("~s is likely to have following states: ~p.~n",
             [ModName, States]),

    Index = [begin
                 {Name, Arity} = erl_syntax_lib:analyze_function(Form),
                 ?CONSOLE("Function ~s/~p found.~n", [Name, Arity]),
                 {{Name, Arity}, Form}
             end
             || Form = {tree, function, _Attr, _Func} <- Forms],
    _Fsm1 = Fsm0#fsm{exports=Exports, statenames=States, fun_index=Index}.

analyze(#fsm{fun_index=Index, statenames=StateNames} = Fsm) ->
    InitForm = proplists:get_value({init, 1}, Index),
    InitRet = return_statements(InitForm, Fsm, []),
    InitStates = init_return_state(InitRet),
    ?CONSOLE("Initial state: init -> ~p~n", [InitStates]),
    Trans = [{?INIT, InitState, 'init/1'} || InitState <- InitStates],

    Trans0 = [ begin
                   Form2 = proplists:get_value({StateName, 2}, Index),
                   Form3 = proplists:get_value({StateName, 3}, Index),
                   %% 1. analyze each state calls
                   analyze_state_call_2(StateName, Form2, Fsm) ++
                       %% 2. analyze each state sync calls
                       analyze_state_call_3(StateName, Form3, Fsm)
               end || StateName <- lists:usort(StateNames ++ InitStates)],
    %% 3. analyze handle_event
    HandleEventForm = proplists:get_value({handle_event, 3}, Index),
    Trans1 = analyze_handle_event(HandleEventForm, Fsm),
    %% 4. analyze handle_sync_event
    HandleSyncEventForm = proplists:get_value({handle_sync_event, 4}, Index),
    Trans2 = analyze_handle_sync_event(HandleSyncEventForm, Fsm),
    %% 5. analyze handle_info
    HandleInfoForm = proplists:get_value({handle_info, 3}, Index),
    Trans3 = analyze_handle_info(HandleInfoForm, Fsm),
    lists:usort(Trans ++ lists:flatten(Trans0) ++ Trans1 ++ Trans2 ++ Trans3).

-spec analyze_state_call_2(atom(), erl_syntax:syntaxTree(), fsm()) -> [edge()].
analyze_state_call_2(_, undefined, _) ->
    [];
analyze_state_call_2(FromStateName, Form, Fsm) ->
    RetStats = return_statements(Form, Fsm, [{FromStateName, 2}]),
    [begin
         [Event, _] = Argv,
         ToStateName = statename_from_return_statement(Ret),
         {FromStateName, ToStateName, {event, pp_term_tree(Event)}}
     end
     || {Argv, Ret} <- RetStats].

analyze_state_call_3(_, undefined, _) ->
    [];
analyze_state_call_3(FromStateName, Form, Fsm) ->
    ?CONSOLE("erasdf ~p, ~p~n", [FromStateName, Form]),
    RetStats = return_statements(Form, Fsm, [{FromStateName, 3}]),
    [begin
         [Event, _, _] = Argv,
         ToStateName = statename_from_return_statement(Ret),
         {FromStateName, ToStateName, {sync_event, pp_term_tree(Event)}}
     end
     || {Argv, Ret} <- RetStats].

analyze_handle_event(HandleEventForm, Fsm) ->
    RetStats = return_statements(HandleEventForm, Fsm, []),
    [begin
         [Event, StateName0, _] = Argv,
         FromStateName = statename_from_tuple(StateName0),
         ToStateName = statename_from_return_statement(Ret),
         {FromStateName, ToStateName, {all_state_event, pp_term_tree(Event)}}
     end
     || {Argv, Ret} <- RetStats].

analyze_handle_sync_event(HandleSyncEventForm, Fsm) ->
    RetStats = return_statements(HandleSyncEventForm, Fsm, []),
    [begin
         [Event, _, StateName0, _] = Argv,
         %% ?CONSOLE(">>>wawawa: ~p => ~n ~p~n", [Argv, Ret]),
         FromStateName = statename_from_tuple(StateName0),
         ToStateName = statename_from_return_statement(Ret),
         {FromStateName, ToStateName, {sync_all_state_event, pp_term_tree(Event)}}
     end
     || {Argv, Ret} <- RetStats].

analyze_handle_info(HandleInfoForm, Fsm) ->
    RetStats = return_statements(HandleInfoForm, Fsm, []),
    [begin
         [Msg0, StateName0, _] = Argv,
         FromStateName = statename_from_tuple(StateName0),
         ToStateName = statename_from_return_statement(Ret),
         {FromStateName, ToStateName, {msg, pp_term_tree(Msg0)}}
     end
     || {Argv, Ret} <- RetStats].

statename_from_tuple({atom, _, StateName0}) ->
    StateName0;
statename_from_tuple({var, _, _}) ->
    ?ANYSTATE.

statename_from_return_statement({tree, tuple, _, Items}) ->
    [A, B|Rest] = Items,
    case statename_from_tuple(A) of
        next_state ->
            statename_from_tuple(B);
        reply ->
            [C|_] = Rest,
            statename_from_tuple(C);
        ok ->
            %% used in init
            statename_from_tuple(B);
        stop ->
            ?STOP
    end.

return_statements({tree, function, _Attr, Func} = Tree, Fsm, CallStack) ->
    {func, _Attr2, Heads} = Func,
    Funcname = erl_syntax_lib:analyze_function(Tree),
    clauses(Heads, Fsm, [Funcname|CallStack]).

clauses(Clauses, Fsm, CallStack) ->
    Ret = lists:map(fun({tree, clause, _, C}) ->
                            {clause, Argv, _, Code} = C,
                            [begin
                                 {Argv,  Leaf}
                             end
                             || Leaf <- return_leaves(lists:last(Code), Fsm, CallStack)]
                    end, Clauses),
    lists:flatten(Ret).

%% Extract state State Candidates
init_return_state(RetVals) ->
    Cands = [begin
                 statename_from_return_statement(RetVal)
             end || {_, RetVal} <- RetVals],
    lists:usort(lists:flatten(Cands)).

%% Return AST return leaves
return_leaves({tree, tuple, _, _} = Tuple, _, _) ->
    [Tuple];
return_leaves({tree, case_expr, _, {case_expr, _Expr, Clauses0}}, Fsm, CallStack) ->
    [begin
         {_, Leaf} = Clause,
         Leaf
     end || Clause <- clauses(Clauses0, Fsm, CallStack) ];
return_leaves({tree, if_expr, _, Clauses0}, Fsm, CallStack) ->
    %%?CONSOLE("<><>< ~p~n", [Clauses]),
    [begin
         {_, Leaf} = Clause,
         Leaf
     end || Clause <- clauses(Clauses0, Fsm, CallStack) ];
return_leaves({tree, application, _, Application},
              #fsm{fun_index=Index} = Fsm,
              CallStack) ->
    {application, FuncName0, Argv} = Application,
    %% TODO: FuncName could be not only an atom
    case FuncName0 of
        {atom, _, FuncName} ->
            FuncArity = {FuncName, length(Argv)},
            case lists:member(FuncArity, CallStack) of
                true ->
                    %% tail recursion!
                    [];
                _ ->
                    case proplists:get_value(FuncArity, Index) of
                        undefined ->
                            ?CONSOLE("[warning] external call? ~p~n", [FuncArity]),
                            [];
                        Func ->
                            RetStmts = return_statements(Func, Fsm,
                                                         [FuncArity|CallStack]),
                            [begin
                                 {_, Leaf} = Clause,
                                 Leaf
                             end || Clause <- RetStmts]
                    end
            end;
        Func ->
            ?CONSOLE("[warning] unsupported function call style: ~p~n",
                     [Func]),
            []
    end;
return_leaves({tree, try_expr, _Attr, TryExpr}, Fsm, CallStack) ->
    {try_expr, Tryee, _, CatchClauses, _After} = TryExpr,
    [begin
         {_, Leaf} = Clause,
         Leaf
     end || Clause <- clauses(CatchClauses, Fsm, CallStack)] ++
        return_leaves(lists:last(Tryee), Fsm, CallStack);

return_leaves({tree, receive_expr, _Attr, RecvClauses0}, Fsm, CallStack) ->
    %% erlang:display(length(tuple_to_list(RecvClauses0))),
    {receive_expr, Clauses, _AfterWhat, After} = RecvClauses0,
    %% ?CONSOLE("~p: ~p~n", [?LINE, lists:last(After)]),
    [begin
         {_, Leaf} = Clause,
         Leaf
     end || Clause <- clauses(Clauses, Fsm, CallStack)] ++
        return_leaves(lists:last(After), Fsm, CallStack);

return_leaves(Tree, _, _) ->
    Tree.

pp_term_tree({atom, _, Name}) ->
    Name;
pp_term_tree({var, _, Name}) ->
    Name;
pp_term_tree({tree, tuple, _, Elems}) ->
    list_to_tuple(lists:map(fun pp_term_tree/1, Elems)).

extract_states([], States) ->
    States;
extract_states([{Fun, 2}|Rest], States0) ->
    case proplists:get_value(Fun, Rest) of
        3 ->
            extract_states(Rest, [Fun|States0]);
        _ ->
            extract_states(Rest, States0)
    end;
extract_states([{Fun, 3}|Rest], States0) ->
    case proplists:get_value(Fun, Rest) of
        2 ->
            extract_states(Rest, [Fun, States0]);
        _ ->
            extract_states(Rest, States0)
    end;
extract_states([_|Rest], States0) ->
    extract_states(Rest, States0).

-spec render_graph(atom(), atom(), [edge()]) -> ok.
render_graph(Modname, AllStateNames, Transitions) ->
    ?STDOUT("digraph ~s {~n", [Modname]),
    [
     begin
         FromTos =
             case {From0, To0} of
                 {?ANYSTATE, ?ANYSTATE} ->
                     [];
                     %% [ {State, State} || State <- AllStateNames ];
                 {?ANYSTATE, _} ->
                     [ {State, To0} || State <- AllStateNames ];
                 {_, ?ANYSTATE} ->
                     [ {From0, State} || State <- AllStateNames ];
                 _ ->
                     [{From0, To0}]
             end,
         [?STDOUT("  ~s -> ~s [ label = \"~p\" ] ~n",[From, To, Term])
          || {From, To} <- FromTos]
     end
     || {From0, To0, Term} <- Transitions],
    ?STDOUT("}~n", []).