gen_fsm_visualizer.erl

%% ---------------------------------------------------------------------
%%
%% Copyright (c) 2015 Kota UENISHI.  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 = analyze(Fsm1),
            ?CONSOLE("State Transitions:~n ~p~n", [Transitions]),
            #fsm{statenames=StateNames} = Fsm1,
            render_graph(ModName, StateNames, Transitions);
        _ ->
            ?CONSOLE("~s does not have gen_fsm behaviour.~n", [ModName])
    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 <- StateNames ],
    %% 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),
    Trans ++ lists:flatten(Trans0) ++ Trans1 ++ Trans2 ++ Trans3.

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

analyze_state_call_3(FromStateName, Form, Fsm) ->
    RetStats = return_statements(Form, Fsm),
    [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}, Fsm) ->
    {func, _Attr2, Heads} = Func,
    clauses(Heads, Fsm).

clauses(Clauses, Fsm) ->
    Ret = [begin
               {clause, Argv, _, Code} = C,
               [begin
                    {Argv,  Leaf}
                end
                || Leaf <- return_leaves(lists:last(Code), Fsm)]
           end || {tree, clause, _, C} <- 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) ->
    [begin
         {_, Leaf} = Clause,
         Leaf
     end || Clause <- clauses(Clauses0, Fsm) ];
return_leaves({tree, application, _, Application},  #fsm{fun_index=Index} =Fsm) ->
    {application, FuncName0, Argv} = Application,
    %% TODO: FuncName could be not only an atom
    {atom, _, FuncName} = FuncName0,
    Func = proplists:get_value({FuncName, length(Argv)}, Index),
    %% Currently this does not support trimming
    %% self-recursive call, which leads to infinite
    %% recursion here.
    [begin
         {_, Leaf} = Clause,
         Leaf
     end || Clause <- return_statements(Func, Fsm)];
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", []).