]> rtime.felk.cvut.cz Git - boost-statechart-viewer.git/blobdiff - src/visualizer.cpp
Convert to Clang plugin API
[boost-statechart-viewer.git] / src / visualizer.cpp
index 7cd375b9ecdd2b227f0e4aa44c9ae9942bec6cbd..59008b378716ce227b081d0527d26efa65a4b221 100644 (file)
@@ -1,8 +1,25 @@
+/** @file */ 
+////////////////////////////////////////////////////////////////////////////////////////  
+//    
+//    This file is part of Boost Statechart Viewer.
+//
+//    Boost Statechart Viewer is free software: you can redistribute it and/or modify
+//    it under the terms of the GNU General Public License as published by
+//    the Free Software Foundation, either version 3 of the License, or
+//    (at your option) any later version.
+//
+//    Boost Statechart Viewer is distributed in the hope that it will be useful,
+//    but WITHOUT ANY WARRANTY; without even the implied warranty of
+//    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+//    GNU General Public License for more details.
+//
+//    You should have received a copy of the GNU General Public License
+//    along with Boost Statechart Viewer.  If not, see <http://www.gnu.org/licenses/>.
+//
+////////////////////////////////////////////////////////////////////////////////////////
+
 //standard header files
 #include <iostream>
-#include <string>
-#include <fstream>
-#include <list>
 
 //LLVM Header files
 #include "llvm/Support/raw_ostream.h"
 #include "clang/Driver/Driver.h"
 #include "clang/Driver/Compilation.h"
 
+#include "clang/Frontend/FrontendPluginRegistry.h"
+#include "clang/AST/ASTConsumer.h"
+#include "clang/AST/AST.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "llvm/Support/raw_ostream.h"
+
 //my own header files
-#include "stringoper.h"
+#include "iooper.h"
 
 using namespace clang;
 using namespace clang::driver;
 using namespace std;
 
-class MyDiagnosticClient : public TextDiagnosticPrinter // My diagnostic Client
+/**
+ * This class provides Simple diagnostic Client. It uses implementation in library for printing diagnostci information.
+ * Also it counts number of warnings, errors, ... When an error occurs the program is stopped.
+ */
+class MyDiagnosticClient : public TextDiagnosticPrinter
 {
+       int nwarnings; /** Save number of Warnings occured during diagnostic */
+       int nnotes;
+       int nignored;
+       int nerrors;
        public:
-       MyDiagnosticClient(llvm::raw_ostream &os, const DiagnosticOptions &diags, bool OwnsOutputStream = false):TextDiagnosticPrinter(os, diags, OwnsOutputStream = false){}
-       virtual void HandleDiagnostic(Diagnostic::Level DiagLevel, const DiagnosticInfo &Info)
+   /**
+    * Initialize number of warnings, errors, ...
+    */
+   MyDiagnosticClient(llvm::raw_ostream &os, const DiagnosticOptions &diags, bool OwnsOutputStream = false):TextDiagnosticPrinter(os, diags, OwnsOutputStream = false)
+       {
+               nwarnings=0;
+               nnotes=0;
+               nignored=0;
+               nerrors = 0;
+       }
+   /**
+     * This method prints diagnostic and counts diagnostic types.
+     */
+       virtual void HandleDiagnostic(DiagnosticsEngine::Level DiagLevel, const Diagnostic &Info)
+       {
+               TextDiagnosticPrinter::HandleDiagnostic(DiagLevel, Info); // print diagnostic information using library implementation
+               switch (DiagLevel) // count number of all diagnostic information
+               {
+                       case 0 : nignored+=1; break;
+                       case 1 : nnotes+=1; break;
+                       case 2 : nwarnings+=1; break;
+                       default : nerrors+=1; 
+                                                print_stats(); 
+                                                exit(1);
+               }
+       }
+   /**
+     * Print statistics about diagnostic.
+     */
+       void print_stats()
+       {
+               cout<<"\n--Diagnostic Info--\n";
+               cout<<"Number of ignored: "<<nignored<<"\n";
+               cout<<"Number of notes: "<<nnotes<<"\n";
+               cout<<"Number of warnings: "<<nwarnings<<"\n";
+               cout<<"Number of errors and fatal errors: "<<nerrors<<"\n";
+       }
+       
+       int getNbrOfWarnings() /** Return number of warnings */
+       {
+               return nwarnings;               
+       }
+       
+       int getNbrOfNotes() /** Return number of notes */
+       {
+               return nnotes;          
+       }
+
+       int getNbrOfIgnored() /** Return number of ignored */
        {
-               TextDiagnosticPrinter::HandleDiagnostic(DiagLevel, Info); // print diagnostic information
-               if(DiagLevel > 2) // if error/fatal error stop the program
-               {               
-                       exit(1);
-               }       
+               return nignored;                
        }
 };
 
+/**
+ * My ASTConsumer provides interface for traversing AST. It uses recursive traversing in namespaces.
+ */
 class FindStates : public ASTConsumer
 {
-       list<string> transitions;
-
-       list<string> cReactions;
+       list<string> transitions; 
+       list<string> cReactions; /** list of custom reactions. After all files are traversed this list should be empty. */
        list<string> events;
+       list<string> states;
        string name_of_machine;
        string name_of_start;
-       StringDecl sd;  
-       int nbrStates;
+       string dest;
+       FullSourceLoc *fsloc; /** Full Source Location instance for holding Source Manager. */
        public:
-       list<string> states;
+
+       FindStates(string dest) : dest(dest) {}
+
+       list<string> getStates() /** Return list of states of the state machine. */
+       {
+               return states;
+       }
+       
+       list<string> getTransitions() /** Return list of transitions. */
+       {
+               return transitions;
+       }
+               
+       list<string> getEvents() /** Return list of events. */
+       {
+               return events;
+       }
+
+       string getStateMachine() /** Return name of the state machine. */
+       {
+               return name_of_machine;
+       }
+
+       string getNameOfFirstState() /** Return name of start state. */
+       {
+               return name_of_start;
+       }
        
-       virtual void Initialize(ASTContext &ctx)//run after the AST is constructed before the consumer starts to work
+       virtual void Initialize(ASTContext &ctx)/** Run after the AST is constructed before the consumer starts to work. So this function works like constructor. */
        {       
+               fsloc = new FullSourceLoc(* new SourceLocation(), ctx.getSourceManager());
                name_of_start = "";
                name_of_machine = "";
-               nbrStates = 0;
        }
 
-       virtual void HandleTopLevelDecl(DeclGroupRef DGR)// traverse all top level declarations
+/**
+*   Traverse global decls using DeclGroupRef for handling all global decls. But only interesting decls are processed. Interesting decls are Struct, Class, C++ methods and Namespace.
+*   When Namespace is found it recursively traverse all decls inside this Namespace using method recursive_visit.
+*/
+       virtual bool HandleTopLevelDecl(DeclGroupRef DGR)
        {
                SourceLocation loc;
-      std::string line, output;
+               string line, output, event;
                llvm::raw_string_ostream x(output);
                for (DeclGroupRef::iterator i = DGR.begin(), e = DGR.end(); i != e; ++i) 
                {
@@ -76,10 +183,9 @@ class FindStates : public ASTConsumer
                        loc = decl->getLocation();
                        if(loc.isValid())
                        {
-                               //cout<<decl->getKind()<<"ss\n";
-                               if(const CXXMethodDecl *cMDecl = dyn_cast<CXXMethodDecl>(decl))
-                               {
-                                       //decl->dump();
+                               if(decl->getKind()==35)
+                               {                                       
+                                       method_decl(decl);
                                }
                                if (const TagDecl *tagDecl = dyn_cast<TagDecl>(decl))
                                {
@@ -92,30 +198,32 @@ class FindStates : public ASTConsumer
                                {
                                        
                                        DeclContext *declCont = namespaceDecl->castToDeclContext(namespaceDecl);
-                                       //cout<<namedDecl->getNameAsString()<<"   sss\n";
                                        recursive_visit(declCont);
                                
                                }
                        }
+                       output = "";
                }
+               return true;
        }
-       void recursive_visit(const DeclContext *declCont) //recursively visit all decls hidden inside namespaces
+
+/**
+*   It is used to recursive traverse decls in Namespaces. This method do the same as HandleTopLevelDecl.
+*/
+       void recursive_visit(const DeclContext *declCont)
        {
-               std::string line, output;
+               string line, output, event;
+               llvm::raw_string_ostream x(output);
                SourceLocation loc;
                for (DeclContext::decl_iterator i = declCont->decls_begin(), e = declCont->decls_end(); i != e; ++i)
                {
                        const Decl *decl = *i;
-                       //std::cout<<"a "<<decl->getDeclKindName()<<"\n";
                        loc = decl->getLocation();
                        if(loc.isValid())
                        {       
-                               if(const CXXMethodDecl *cMDecl = dyn_cast<CXXMethodDecl>(decl))
+                               if(decl->getKind()==35)
                                {
-                                       /*TODO
-                                               get Base class, test name and chcek if it is a state
-                                               
-                                       */
+                                       method_decl(decl);
                                }
                                else if (const TagDecl *tagDecl = dyn_cast<TagDecl>(decl))
                                {
@@ -126,258 +234,238 @@ class FindStates : public ASTConsumer
                                }
                                else if(const NamespaceDecl *namespaceDecl = dyn_cast<NamespaceDecl>(decl))
                                {
-                                       DeclContext *declCont = namespaceDecl->castToDeclContext(namespaceDecl);
-                                       //cout<<namedDecl->getNameAsString()<<"  sss\n";                        
+                                       DeclContext *declCont = namespaceDecl->castToDeclContext(namespaceDecl);                        
                                        recursive_visit(declCont);
                                }
                        }
+                       output = "";
                } 
        }
                
-       void struct_class(const Decl *decl) // works with struct or class decl
+/**
+*      This function works with class or struct. It splits the decl into 3 interesting parts.
+*      The state machine decl, state decl and event decl.
+*/
+       void struct_class(const Decl *decl)
        {
-               string output, line, ret, trans;        
+               string output, line, ret, trans, event; 
                llvm::raw_string_ostream x(output);
                decl->print(x);
-               line = sd.get_line_of_code(x.str());
+               line = get_line_of_code(x.str());
+               
                output = "";
-               int pos, num;
-               const TagDecl *tagDecl = dyn_cast<TagDecl>(decl);
-               const NamedDecl *namedDecl = dyn_cast<NamedDecl>(decl);         
-               if(sd.is_derived(line))
+               int pos;
+               const NamedDecl *namedDecl = dyn_cast<NamedDecl>(decl);
+               if(is_derived(line))
                {
                        const CXXRecordDecl *cRecDecl = dyn_cast<CXXRecordDecl>(decl);
                                        
-                       if(sd.find_events(cRecDecl, line))
+                       if(find_events(cRecDecl, line))
                        {
                                events.push_back(namedDecl->getNameAsString());
-                               cout<<"New event: "<<namedDecl->getNameAsString()<<"\n";
                        }
                        else if(name_of_machine == "")
                        {
-                               ret = sd.find_name_of_machine(cRecDecl, line);
+                               ret = find_name_of_machine(cRecDecl, line);
                                if(!ret.empty())
                                {
                                        pos = ret.find(",");
                                        name_of_machine = ret.substr(0,pos);
                                        name_of_start = ret.substr(pos+1);
-                                       cout<<"Name of the state machine: "<<name_of_machine<<"\n";
-                                       cout<<"Name of the first state: "<<name_of_start<<"\n";
                                }
                        }
                        else
                        {
-                               ret = sd.find_states(cRecDecl, line);   
+                               ret = find_states(cRecDecl, line); 
                                if(!ret.empty())
-                               {                               
-                                       const DeclContext *declCont = tagDecl->castToDeclContext(tagDecl);              
-                                       //states.push_back(namedDecl->getNameAsString());
-                                       std::cout << "New state: " << namedDecl->getNameAsString() << "\n";
-                                       states.push_back(ret);
-                                       ret = sd.find_transitions(namedDecl->getNameAsString(), declCont);
-                                       if(!ret.empty()) 
-                                       {
-                                               num = sd.count(ret,';')+1;
-                                               for(int i = 0;i<num;i++)
-                                               {
-                                                       pos = ret.find(";");
-                                                       if(pos == 0)
-                                                       {
-                                                               ret = ret.substr(1);
-                                                               pos = ret.find(";");
-                                                               if(pos==-1) cReactions.push_back(ret);
-                                                               else cReactions.push_back(ret.substr(0,pos));   
-                                                               num-=1;
-                                                       }
-                                                       else 
-                                                       {
-                                                               if(pos==-1) transitions.push_back(ret);
-                                                               else transitions.push_back(ret.substr(0,pos));
-                                                       }
-                                                       //cout<<ret<<"\n";
-                                                       if(i!=num-1) ret = ret.substr(pos+1);
-                                               }
-                                       }
+                               {
+                                       states.push_back(ret);                  
+                                       methods_in_class(decl,namedDecl->getNameAsString());
                                }
                        }
                }
        }
-       
-       void save_to_file(std::string output) // save all to the output file
+
+/**
+*      This function provides traversing all methods and other context indide class. If
+*      typedef or classic method decl is found. If typedef is found then it is being testted for transitions and custom reactions.
+*/
+       void methods_in_class(const Decl *decl, const string state)
        {
-               nbrStates = states.size();
-               std::string state, str, context, ctx;
-               int pos1, pos2, cnt, subs;
-               std::ofstream filestr(output.c_str());
-               //std::cout<<output<<"\n";
-               filestr<<"digraph "<< name_of_machine<< " {\n";
-               context = name_of_machine;
-               for(list<std::string>::iterator i = states.begin();i!=states.end();i++) // write all states in the context of the automaton
+               string output, line, ret, trans, event; 
+               llvm::raw_string_ostream x(output);
+               int pos, num;
+               const TagDecl *tagDecl = dyn_cast<TagDecl>(decl);
+               const DeclContext *declCont = tagDecl->castToDeclContext(tagDecl);                      
+               output="";
+               std::cout<<"Found state: "<<state<<std::endl;
+               for (DeclContext::decl_iterator i = declCont->decls_begin(), e = declCont->decls_end(); i != e; ++i) 
                {
-                       state = *i;
-                       cnt = sd.count(state,',');
-                       if(cnt==1)
+                       if (i->getKind()==26) // typedefs
                        {
-                               pos1 = state.find(",");
-                               ctx = sd.cut_namespaces(state.substr(pos1+1));
-                               //std::cout<<name_of_machine.length();                          
-                               if(ctx.compare(0,context.length(),context)==0)
+                               i->print(x);
+                               output = x.str();
+                               line = clean_spaces(cut_type(output));          
+                               ret = find_transitions(state,line);
+                               if(!ret.empty()) 
                                {
-                                       filestr<<sd.cut_namespaces(state.substr(0,pos1))<<";\n";
-                                       states.erase(i);
-                                       i--;
-                               }
-                       }
-                       if(cnt==2)
-                       {
-                               pos1 = state.find(",");
-                               pos2 = state.rfind(",");
-                               ctx = sd.cut_namespaces(state.substr(pos1+1,pos2-pos1-1));
-                               //std::cout<<ctx<<" "<<context<<"\n";
-                               if(ctx.compare(0,context.length(),context)==0)
-                               {                               
-                                       filestr<<sd.cut_namespaces(state.substr(0,pos1))<<";\n";
+                                       num = count(ret,';')+1;
+                                       for(int i = 0;i<num;i++)
+                                       {
+                                               pos = ret.find(";");
+                                               if(pos == 0)
+                                               {
+                                                       ret = ret.substr(1);
+                                                       pos = ret.find(";");
+                                                       if(pos==-1) cReactions.push_back(ret);
+                                                       else cReactions.push_back(ret.substr(0,pos));   
+                                                       num-=1;
+                                               }
+                                               else 
+                                               {
+                                                       if(pos==-1) transitions.push_back(ret);
+                                                       else transitions.push_back(ret.substr(0,pos));
+                                               }
+                                               if(i!=num-1) ret = ret.substr(pos+1);
+                                       }
+                                       output="";
                                }
                        }
+                       if(i->getKind()==35) method_decl(*i);// C++ method
                }
-               filestr<<name_of_start<<" [peripheries=2] ;\n";
-               subs = 0;
-               while(!states.empty()) // substates ?
+               
+       }
+
+ /**
+   * Traverse method declaration using classes with main class Stmt.
+   */
+       void method_decl(const Decl *decl)
+       {
+               string output, line, event;     
+               llvm::raw_string_ostream x(output);
+
+               if(decl->hasBody())
                {
-                       state = states.front();
-                       filestr<<"subgraph cluster"<<subs<<" {\n";                      
-                       pos1 = state.find(",");
-                       pos2 = state.rfind(",");
-                       context = sd.cut_namespaces(state.substr(0,pos1));
-                       filestr<<"label=\""<<context<<"\";\n";
-                       filestr<<sd.cut_namespaces(state.substr(pos2+1))<<" [peripheries=2] ;\n";       
-                       states.pop_front();     
-                       //std::cout<<states.size();     
-                       for(list<string>::iterator i = states.begin();i!=states.end();i++)
+                       decl->print(x);
+                       line = get_return(x.str());                     
+                       if(get_model(line)==5)
                        {
-                               state = *i;
-                               cnt = sd.count(state,',');
-                               //std::cout<<state<<" \n";
-                               if(cnt==1)
+                               //std::cout<<"metodass"<<std::endl;
+                               const FunctionDecl *fDecl = dyn_cast<FunctionDecl>(decl);
+                               const ParmVarDecl *pvd = fDecl->getParamDecl(0);
+                               QualType qt = pvd->getOriginalType();
+                               event = qt.getAsString();
+                               if(event[event.length()-1]=='&') event = event.substr(0,event.length()-2);
+                               event = event.substr(event.rfind(" ")+1);
+                               line = dyn_cast<NamedDecl>(decl)->getQualifiedNameAsString();
+                               line = cut_namespaces(line.substr(0,line.rfind("::")));
+                               line.append(",");
+                               line.append(event);
+                               find_return_stmt(decl->getBody(),line); 
+                               for(list<string>::iterator i = cReactions.begin();i!=cReactions.end();i++) // erase info about it from list of custom reactions
                                {
-                                       pos1 = state.find(",");
-                                       ctx = sd.cut_namespaces(state.substr(pos1+1));
-                                       
-                                       //std::cout<<ctx<<" "<<context<<"\n";
-                                       if(ctx.compare(0,context.length(),context)==0)
+                                       event = *i;
+                                       if(line.compare(event)==0) 
                                        {
-                                               filestr<<sd.cut_namespaces(state.substr(0,pos1))<<";\n";
-                                               states.erase(i);
-                                               i--;
-                                       }
-                               }
-                               if(cnt==2)
-                               {
-                                       pos1 = state.find(",");
-                                       pos2 = state.rfind(",");
-                                       ctx = sd.cut_namespaces(state.substr(pos1+1,pos2-pos1-1));
-                                       if(ctx.compare(0,context.length(),context)==0)
-                                       {                               
-                                               filestr<<sd.cut_namespaces(state.substr(0,pos1))<<";\n";
-                                               //std::cout<<ctx<<"\n";
+                                               cReactions.erase(i);
+                                               break;
                                        }
                                }
                        }
-                       filestr<<"}\n";
-                       subs+=1;        
-               }               
-               for(list<string>::iterator i = transitions.begin();i!=transitions.end();i++) // write all transitions
-               {
-                       state = *i;
-                       pos1 = state.find(",");
-                       filestr<<sd.cut_namespaces(state.substr(0,pos1))<<"->";
-                       pos2 = state.rfind(",");
-                       filestr<<sd.cut_namespaces(state.substr(pos2+1));
-                       filestr<<"[label=\""<<sd.cut_namespaces(state.substr(pos1+1,pos2-pos1-1))<<"\"];\n";
-               }               
-               filestr<<"}";
-               filestr.close();
-       }
-       void print_stats() // print statistics
-       {
-               cout<<"\n"<<"Statistics: \n";
-               cout<<"Number of states: "<<nbrStates<<"\n";
-               cout<<"Number of events: "<<events.size()<<"\n";
-               cout<<"Number of transitions: "<<transitions.size()<<"\n";
-               return;
+               }
        }
 
-};
-
-int main(int argc, char **argv)
-{ 
-       string inputFilename = "";
-       string outputFilename = "graph.dot"; // initialize output Filename
-       MyDiagnosticClient *mdc = new MyDiagnosticClient(llvm::errs(), * new DiagnosticOptions());
-       llvm::IntrusiveRefCntPtr<DiagnosticIDs> dis(new DiagnosticIDs());       
-       Diagnostic diag(dis,mdc);
-       FileManager fm( * new FileSystemOptions());
-       SourceManager sm (diag, fm);
-       HeaderSearch *headers = new HeaderSearch(fm);
-       
-       Driver TheDriver(LLVM_PREFIX "/bin", llvm::sys::getHostTriple(), "", false, false, diag);
-       TheDriver.setCheckInputsExist(true);
-       TheDriver.CCCIsCXX = 1; 
-       CompilerInvocation compInv;
-       llvm::SmallVector<const char *, 16> Args(argv, argv + argc);
-       llvm::OwningPtr<Compilation> C(TheDriver.BuildCompilation(Args.size(),
-                                                            Args.data()));
-       const driver::JobList &Jobs = C->getJobs();
-       const driver::Command *Cmd = cast<driver::Command>(*Jobs.begin());
-       const driver::ArgStringList &CCArgs = Cmd->getArguments();
-       for(unsigned i = 0; i<Args.size();i++) // find -o in ArgStringList
-       {       
-               if(strncmp(Args[i],"-o",2)==0) 
+       void find_return_stmt(Stmt *statemt,string event) /** Traverse all statements in function for finding all return Statements.*/
+       {
+               if(statemt->getStmtClass() == 99) test_stmt(dyn_cast<CaseStmt>(statemt)->getSubStmt(), event);
+               else
                {
-                       if(strlen(Args[i])>2)
+                       for (Stmt::child_range range = statemt->children(); range; ++range)    
                        {
-                               string str = Args[i];
-                               outputFilename = str.substr(2);
+                               test_stmt(*range, event);
                        }
-                       else outputFilename = Args[i+1];
-                       break;
                }
        }
-               
-       CompilerInvocation::CreateFromArgs(compInv,
-                                         const_cast<const char **>(CCArgs.data()),
-                                         const_cast<const char **>(CCArgs.data())+CCArgs.size(),
-                                         diag);
+       
+       void test_stmt(Stmt *stmt, string event) /** test statement for its kind Using number as identifier for all Statement Classes.*/
+       {
+               const SourceManager &sman = fsloc->getManager();
+               int type;
+               string line, param;
+               type = stmt->getStmtClass();
+               switch(type)
+               {       
+                       case 8 :        find_return_stmt(dyn_cast<DoStmt>(stmt)->getBody(), event); // do
+                                               break;
+                       case 86 :       find_return_stmt(dyn_cast<ForStmt>(stmt)->getBody(), event); // for
+                                               break;
+                       case 88 :       find_return_stmt(dyn_cast<IfStmt>(stmt)->getThen(), event); //if then
+                                               find_return_stmt(dyn_cast<IfStmt>(stmt)->getElse(), event); //if else
+                                               break;
+                       case 90 :       find_return_stmt(dyn_cast<LabelStmt>(stmt)->getSubStmt(), event); //label
+                                               break;
+                       case 98 :       line = sman.getCharacterData(dyn_cast<ReturnStmt>(stmt)->getReturnLoc()); 
+                                               line = get_line_of_code(line).substr(6);
+                                               line = line.substr(0,line.find("("));
+                                               if(get_model(line)==6)
+                                               {
+                                                       param = get_params(line);
+                                                       transitions.push_back(event.append(",").append(param));
+                                               }
+                                               if(get_model(line) == 7)
+                                               {
+                                                       param = ",";
+                                                       transitions.push_back(param.append(event));
+                                               }
+                                               break;
+                       case 99 :       find_return_stmt(stmt, event);
+                                               break;
+                       case 101 :      find_return_stmt(dyn_cast<SwitchStmt>(stmt)->getBody(), event); // switch
+                                               break;
+                       case 102 :      find_return_stmt(dyn_cast<WhileStmt>(stmt)->getBody(), event); // while
+                                               break;
+               }
+       }
+
+  virtual void HandleTranslationUnit(clang::ASTContext &Context) {
+    IO_operations io(dest, getStateMachine(), getNameOfFirstState(), getTransitions(), getStates(), getEvents());
+    io.save_to_file();
+    io.print_stats();
+  }
+};
 
-       HeaderSearchOptions hsopts = compInv.getHeaderSearchOpts();
-       hsopts.ResourceDir = LLVM_PREFIX "/lib/clang/" CLANG_VERSION_STRING;
-       LangOptions lang = compInv.getLangOpts();
-       CompilerInvocation::setLangDefaults(lang, IK_ObjCXX);
-       TargetInfo *ti = TargetInfo::CreateTargetInfo(diag, compInv.getTargetOpts());
-       ApplyHeaderSearchOptions(*headers, hsopts, lang, ti->getTriple());
-       FrontendOptions f = compInv.getFrontendOpts();
-       inputFilename = f.Inputs[0].second;
+class VisualizeStatechartAction : public PluginASTAction {
+protected:
+  ASTConsumer *CreateASTConsumer(CompilerInstance &CI, llvm::StringRef) {
+    size_t dot = getCurrentFile().find_last_of('.');
+    std::string dest = getCurrentFile().substr(0, dot);
+    dest.append(".dot");
+    return new FindStates(dest);
+  }
 
-       cout<<"Input filename: "<<inputFilename<<"\n"; // print Input filename
-       cout<<"Output filename: "<<outputFilename<<"\n"; // print Output filename
+  bool ParseArgs(const CompilerInstance &CI,
+                 const std::vector<std::string>& args) {
+    for (unsigned i = 0, e = args.size(); i != e; ++i) {
+      llvm::errs() << "Visualizer arg = " << args[i] << "\n";
 
+      // Example error handling.
+      if (args[i] == "-an-error") {
+        DiagnosticsEngine &D = CI.getDiagnostics();
+        unsigned DiagID = D.getCustomDiagID(
+          DiagnosticsEngine::Error, "invalid argument '" + args[i] + "'");
+        D.Report(DiagID);
+        return false;
+      }
+    }
+    if (args.size() && args[0] == "help")
+      PrintHelp(llvm::errs());
 
-       Preprocessor pp(diag, lang, *ti, sm, *headers);
-       pp.getBuiltinInfo().InitializeBuiltins(pp.getIdentifierTable(), lang);
-               
-       InitializePreprocessor(pp, compInv.getPreprocessorOpts(),hsopts,f);
-       
-       const FileEntry *file = fm.getFile(inputFilename);
-       sm.createMainFileID(file);
-       IdentifierTable tab(lang);
-       Builtin::Context builtins(*ti);
-       FindStates c;
-       ASTContext ctx(lang, sm, *ti, tab, * new SelectorTable(), builtins,0);
-       mdc->BeginSourceFile(lang, &pp);//start using diagnostic
-       ParseAST(pp, &c, ctx, false, false);
-       mdc->EndSourceFile(); //end using diagnostic
-       if(c.states.size()>0) c.save_to_file(outputFilename);
-       else cout<<"No state machine was found\n";
-       c.print_stats();
-       return 0;
-}
+    return true;
+  }
+  void PrintHelp(llvm::raw_ostream& ros) {
+    ros << "Help for Visualize Statechart plugin goes here\n";
+  }
+
+};
+
+static FrontendPluginRegistry::Add<VisualizeStatechartAction> X("visualize-statechart", "visualize statechart");