]> rtime.felk.cvut.cz Git - boost-statechart-viewer.git/blobdiff - src/visualizer.cpp
Support forward-declared states used as context
[boost-statechart-viewer.git] / src / visualizer.cpp
index b24fccac2fa360fef83b35ff88f2af846fe25d1c..0d8f64f5e200dda7636f1ce368163fcae86b1e72 100644 (file)
@@ -1,6 +1,6 @@
-
-////////////////////////////////////////////////////////////////////////////////////////  
-//    
+/** @file */
+////////////////////////////////////////////////////////////////////////////////////////
+//
 //    This file is part of Boost Statechart Viewer.
 //
 //    Boost Statechart Viewer is free software: you can redistribute it and/or modify
 ////////////////////////////////////////////////////////////////////////////////////////
 
 //standard header files
-#include <iostream>
+#include <iomanip>
+#include <fstream>
+#include <map>
 
 //LLVM Header files
 #include "llvm/Support/raw_ostream.h"
-#include "llvm/Support/Host.h"
-#include "llvm/Config/config.h"
+#include "llvm/Support/raw_os_ostream.h"
 
 //clang header files
-#include "clang/Frontend/TextDiagnosticPrinter.h"
-#include "clang/Lex/HeaderSearch.h"
-#include "clang/Basic/FileManager.h"
-#include "clang/Frontend/Utils.h"
-#include "clang/Basic/TargetInfo.h"
-#include "clang/Lex/Preprocessor.h"
-#include "clang/Frontend/CompilerInstance.h"
 #include "clang/AST/ASTConsumer.h"
-#include "clang/Sema/Lookup.h"
-#include "clang/Parse/ParseAST.h"
-#include "clang/Basic/Version.h"
-#include "clang/Driver/Driver.h"
-#include "clang/Driver/Compilation.h"
-
-//my own header files
-#include "iooper.h"
+#include "clang/AST/ASTContext.h"
+#include "clang/AST/CXXInheritance.h"
+#include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Frontend/FrontendPluginRegistry.h"
 
 using namespace clang;
-using namespace clang::driver;
 using namespace std;
 
-class MyDiagnosticClient : public TextDiagnosticPrinter // My diagnostic Client
+namespace Model
 {
-       int nwarnings;
-       int nnotes;
-       int nignored;
-       int nerrors;
-       public:
-       MyDiagnosticClient(llvm::raw_ostream &os, const DiagnosticOptions &diags, bool OwnsOutputStream = false):TextDiagnosticPrinter(os, diags, OwnsOutputStream = false)
-       {
-               nwarnings=0;
-               nnotes=0;
-               nignored=0;
-               nerrors = 0;
-       }
-       virtual void HandleDiagnostic(Diagnostic::Level DiagLevel, const DiagnosticInfo &Info)
-       {
-               TextDiagnosticPrinter::HandleDiagnostic(DiagLevel, Info); // print diagnostic information
-               switch (DiagLevel)
-               {
-                       case 0 : nignored+=1; break;
-                       case 1 : nnotes+=1; break;
-                       case 2 : nwarnings+=1; break;
-                       default : nerrors+=1; 
-                                                print_stats(); 
-                                                exit(1);
-               }
-       }
-       
-       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";
+
+    inline int getIndentLevelIdx() {
+       static int i = ios_base::xalloc();
+       return i;
+    }
+
+    ostream& indent(ostream& os) { os << setw(2*os.iword(getIndentLevelIdx())) << ""; return os; }
+    ostream& indent_inc(ostream& os) { os.iword(getIndentLevelIdx())++; return os; }
+    ostream& indent_dec(ostream& os) { os.iword(getIndentLevelIdx())--; return os; }
+
+    class State;
+
+    class Context : public map<string, State*> {
+    public:
+       iterator add(State *state);
+       Context *findContext(const string &name);
+    };
+
+    class State : public Context
+    {
+       string initialInnerState;
+    public:
+       const string name;
+       explicit State(string name) : name(name) {}
+       void setInitialInnerState(string name) { initialInnerState = name; }
+       friend ostream& operator<<(ostream& os, const State& s);
+    };
+
+
+    Context::iterator Context::add(State *state)
+    {
+       pair<iterator, bool> ret =  insert(value_type(state->name, state));
+       return ret.first;
+    }
+
+    Context *Context::findContext(const string &name)
+    {
+       iterator i = find(name), e;
+       if (i != end())
+           return i->second;
+       for (i = begin(), e = end(); i != e; ++i) {
+           Context *c = i->second->findContext(name);
+           if (c)
+               return c;
        }
-       
-       int getNbrOfWarnings()
-       {
-               return nwarnings;               
+       return 0;
+    }
+
+
+    ostream& operator<<(ostream& os, const Context& c);
+
+    ostream& operator<<(ostream& os, const State& s)
+    {
+       if (s.size()) {
+           os << indent << "subgraph cluster_" << s.name << " {\n" << indent_inc;
+           os << indent << "label = \"" << s.name << "\"\n";
        }
-       
-       int getNbrOfNotes()
-       {
-               return nnotes;          
+       os << indent << "" << s.name << "\n";
+       if (s.size()) {
+           os << indent << s.initialInnerState << " [peripheries=2]\n";
+           os << static_cast<Context>(s);
+           os << indent_dec << indent << "}\n";
        }
+       return os;
+    }
 
-       int getNbrOfIgnored()
-       {
-               return nignored;                
+
+    ostream& operator<<(ostream& os, const Context& c)
+    {
+       for (Context::const_iterator i = c.begin(), e = c.end(); i != e; i++) {
+           os << *i->second;
        }
-};
+       return os;
+    }
 
-class FindStates : public ASTConsumer
-{
-       list<string> transitions;
-       list<string> cReactions;
-       list<string> events;
-       list<string> states;
-       string name_of_machine;
-       string name_of_start;
-       FullSourceLoc *fsloc;
-       
-       public:
-
-       list<string> getStates()
+
+    class Transition
+    {
+    public:
+       const string src, dst, event;
+       Transition(string src, string dst, string event) : src(src), dst(dst), event(event) {}
+    };
+
+    ostream& operator<<(ostream& os, const Transition& t)
+    {
+       os << indent << t.src << " -> " << t.dst << " [label = \"" << t.event << "\"]\n";
+       return os;
+    }
+
+
+    class Machine : public Context
+    {
+    protected:
+       string initial_state;
+    public:
+       const string name;
+       explicit Machine(string name) : name(name) {}
+
+       void setInitialState(string name) { initial_state = name; }
+
+       friend ostream& operator<<(ostream& os, const Machine& m);
+    };
+
+    ostream& operator<<(ostream& os, const Machine& m)
+    {
+       os << indent << "subgraph " << m.name << " {\n" << indent_inc;
+       os << indent << m.initial_state << " [peripheries=2]\n";
+       os << static_cast<Context>(m);
+       os << indent_dec << indent << "}\n";
+       return os;
+    }
+
+
+    class Model : public map<string, Machine>
+    {
+       Context unknown;        // For forward-declared state classes
+    public:
+       list< Transition*> transitions;
+
+       iterator add(const Machine &m)
        {
-               return states;
+           pair<iterator, bool> ret =  insert(value_type(m.name, m));
+           return ret.first;
        }
-       
-       list<string> getTransitions()
+
+       void addUnknownState(State *m)
        {
-               return transitions;
+           unknown[m->name] = m;
        }
-               
-       list<string> getEvents()
+
+
+       Context *findContext(const string &name)
        {
-               return events;
+           Context::iterator ci = unknown.find(name);
+           if (ci != unknown.end())
+               return ci->second;
+           iterator i = find(name), e;
+           if (i != end())
+               return &i->second;
+           for (i = begin(), e = end(); i != e; ++i) {
+               Context *c = i->second.findContext(name);
+               if (c)
+                   return c;
+           }
+           return 0;
        }
 
-       string getStateMachine()
+       State *removeFromUnknownContexts(const string &name)
        {
-               return name_of_machine;
+           Context::iterator ci = unknown.find(name);
+           if (ci == unknown.end())
+               return 0;
+           unknown.erase(ci);
+           return ci->second;
        }
 
-       string getNameOfFirstState()
+       void write_as_dot_file(string fn)
        {
-               return name_of_start;
+           ofstream f(fn.c_str());
+           f << "digraph statecharts {\n" << indent_inc;
+           for (iterator i = begin(), e = end(); i != e; i++)
+               f << i->second;
+           for (list<Transition*>::iterator t = transitions.begin(), e = transitions.end(); t != e; ++t)
+               f << **t;
+           f << indent_dec << "}\n";
        }
-       
-       virtual void Initialize(ASTContext &ctx)//run after the AST is constructed before the consumer starts to work
-       {       
-               fsloc = new FullSourceLoc(* new SourceLocation(), ctx.getSourceManager());
-               name_of_start = "";
-               name_of_machine = "";
+    };
+};
+
+
+class MyCXXRecordDecl : public CXXRecordDecl
+{
+    static bool FindBaseClassString(const CXXBaseSpecifier *Specifier,
+                                   CXXBasePath &Path,
+                                   void *qualName)
+    {
+        string qn(static_cast<const char*>(qualName));
+        const RecordType *rt = Specifier->getType()->getAs<RecordType>();
+        assert(rt);
+        TagDecl *canon = rt->getDecl()->getCanonicalDecl();
+        return canon->getQualifiedNameAsString() == qn;
+    }
+
+public:
+    bool isDerivedFrom(const char *baseStr, CXXBaseSpecifier const **Base = 0) const {
+       CXXBasePaths Paths(/*FindAmbiguities=*/false, /*RecordPaths=*/!!Base, /*DetectVirtual=*/false);
+       Paths.setOrigin(const_cast<MyCXXRecordDecl*>(this));
+       if (!lookupInBases(&FindBaseClassString, const_cast<char*>(baseStr), Paths))
+           return false;
+       if (Base)
+           *Base = Paths.front().back().Base;
+       return true;
+    }
+};
+
+
+class Visitor : public RecursiveASTVisitor<Visitor>
+{
+    ASTContext *ASTCtx;
+    Model::Model &model;
+    DiagnosticsEngine &Diags;
+    unsigned diag_unhandled_reaction_type, diag_unhandled_reaction_decl,
+       diag_found_state, diag_found_statemachine;
+
+public:
+    bool shouldVisitTemplateInstantiations() const { return true; }
+
+    explicit Visitor(ASTContext *Context, Model::Model &model, DiagnosticsEngine &Diags)
+       : ASTCtx(Context), model(model), Diags(Diags)
+    {
+       diag_found_statemachine =
+           Diags.getCustomDiagID(DiagnosticsEngine::Note, "Found statemachine '%0'");
+       diag_found_state =
+           Diags.getCustomDiagID(DiagnosticsEngine::Note, "Found state '%0'");
+       diag_unhandled_reaction_type =
+           Diags.getCustomDiagID(DiagnosticsEngine::Error, "Unhandled reaction type '%0'");
+       diag_unhandled_reaction_decl =
+           Diags.getCustomDiagID(DiagnosticsEngine::Error, "Unhandled reaction decl '%0'");
+    }
+
+    DiagnosticBuilder Diag(SourceLocation Loc, unsigned DiagID) { return Diags.Report(Loc, DiagID); }
+
+    void HandleReaction(const Type *T, const SourceLocation Loc, CXXRecordDecl *SrcState)
+    {
+       // TODO: Improve Loc tracking
+       if (const ElaboratedType *ET = dyn_cast<ElaboratedType>(T))
+           HandleReaction(ET->getNamedType().getTypePtr(), Loc, SrcState);
+       else if (const TemplateSpecializationType *TST = dyn_cast<TemplateSpecializationType>(T)) {
+           string name = TST->getTemplateName().getAsTemplateDecl()->getQualifiedNameAsString();
+           if (name == "boost::statechart::transition") {
+               const Type *EventType = TST->getArg(0).getAsType().getTypePtr();
+               const Type *DstStateType = TST->getArg(1).getAsType().getTypePtr();
+               CXXRecordDecl *Event = EventType->getAsCXXRecordDecl();
+               CXXRecordDecl *DstState = DstStateType->getAsCXXRecordDecl();
+
+               Model::Transition *T = new Model::Transition(SrcState->getName(), DstState->getName(), Event->getName());
+               model.transitions.push_back(T);
+           } else if (name == "boost::mpl::list") {
+               for (TemplateSpecializationType::iterator Arg = TST->begin(), End = TST->end(); Arg != End; ++Arg)
+                   HandleReaction(Arg->getAsType().getTypePtr(), Loc, SrcState);
+           } else
+               Diag(Loc, diag_unhandled_reaction_type) << name;
+       } else
+           Diag(Loc, diag_unhandled_reaction_type) << T->getTypeClassName();
+    }
+
+    void HandleReaction(const NamedDecl *Decl, CXXRecordDecl *SrcState)
+    {
+       if (const TypedefDecl *r = dyn_cast<TypedefDecl>(Decl))
+           HandleReaction(r->getCanonicalDecl()->getUnderlyingType().getTypePtr(),
+                          r->getLocStart(), SrcState);
+       else
+           Diag(Decl->getLocation(), diag_unhandled_reaction_decl) << Decl->getDeclKindName();
+    }
+
+    CXXRecordDecl *getTemplateArgDecl(const Type *T, unsigned ArgNum)
+    {
+       if (const ElaboratedType *ET = dyn_cast<ElaboratedType>(T))
+           return getTemplateArgDecl(ET->getNamedType().getTypePtr(), ArgNum);
+       else if (const TemplateSpecializationType *TST = dyn_cast<TemplateSpecializationType>(T)) {
+           if (TST->getNumArgs() >= ArgNum+1)
+               return TST->getArg(ArgNum).getAsType()->getAsCXXRecordDecl();
        }
+       return 0;
+    }
+
 
-       virtual void HandleTopLevelDecl(DeclGroupRef DGR)// traverse all top level declarations
+    bool VisitCXXRecordDecl(CXXRecordDecl *Declaration)
+    {
+       if (!Declaration->isCompleteDefinition())
+           return true;
+
+       MyCXXRecordDecl *RecordDecl = static_cast<MyCXXRecordDecl*>(Declaration);
+       const CXXBaseSpecifier *Base;
+
+       if (RecordDecl->isDerivedFrom("boost::statechart::simple_state", &Base))
        {
-               SourceLocation loc;
-      string line, output, event;
-               llvm::raw_string_ostream x(output);
-               for (DeclGroupRef::iterator i = DGR.begin(), e = DGR.end(); i != e; ++i) 
-               {
-                       const Decl *decl = *i;
-                       loc = decl->getLocation();
-                       if(loc.isValid())
-                       {
-                               if(decl->getKind()==35)
-                               {                                       
-                                       method_decl(decl);
-                               }
-                               if (const TagDecl *tagDecl = dyn_cast<TagDecl>(decl))
-                               {
-                                       if(tagDecl->isStruct() || tagDecl->isClass()) //is it a struct or class 
-                                       {
-                                               struct_class(decl);
-                                       }
-                               }       
-                               if(const NamespaceDecl *namespaceDecl = dyn_cast<NamespaceDecl>(decl))
-                               {
-                                       
-                                       DeclContext *declCont = namespaceDecl->castToDeclContext(namespaceDecl);
-                                       //cout<<namedDecl->getNameAsString()<<"   sss\n";
-                                       recursive_visit(declCont);
-                               
-                               }
-                       }
-                       output = "";
-               }
+           string name(RecordDecl->getName()); //getQualifiedNameAsString());
+           Diag(RecordDecl->getLocStart(), diag_found_state) << name;
+
+           Model::State *state;
+           // Either we saw a reference to forward declared state
+           // before, or we create a new state.
+           if (!(state = model.removeFromUnknownContexts(name)))
+               state = new Model::State(name);
+
+           CXXRecordDecl *Context = getTemplateArgDecl(Base->getType().getTypePtr(), 1);
+           Model::Context *c = model.findContext(Context->getName());
+           if (!c) {
+               Model::State *s = new Model::State(Context->getName());
+               model.addUnknownState(s);
+               c = s;
+           }
+           c->add(state);
+
+           if (CXXRecordDecl *InnerInitialState = getTemplateArgDecl(Base->getType().getTypePtr(), 2))
+               state->setInitialInnerState(InnerInitialState->getName());
+
+           IdentifierInfo& II = ASTCtx->Idents.get("reactions");
+           // TODO: Lookup for reactions even in base classes - probably by using Sema::LookupQualifiedName()
+           for (DeclContext::lookup_result Reactions = RecordDecl->lookup(DeclarationName(&II));
+                Reactions.first != Reactions.second; ++Reactions.first)
+               HandleReaction(*Reactions.first, RecordDecl);
        }
-       void recursive_visit(const DeclContext *declCont) //recursively visit all decls hidden inside namespaces
+       else if (RecordDecl->isDerivedFrom("boost::statechart::state_machine", &Base))
        {
-      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(decl->getKind()==35)
-                               {
-                                       method_decl(decl);
-                               }
-                               else if (const TagDecl *tagDecl = dyn_cast<TagDecl>(decl))
-                               {
-                                       if(tagDecl->isStruct() || tagDecl->isClass()) //is it a structure or class      
-                                       {
-                                               struct_class(decl);
-                                       }       
-                               }
-                               else if(const NamespaceDecl *namespaceDecl = dyn_cast<NamespaceDecl>(decl))
-                               {
-                                       DeclContext *declCont = namespaceDecl->castToDeclContext(namespaceDecl);
-                                       //cout<<namedDecl->getNameAsString()<<"  sss\n";                        
-                                       recursive_visit(declCont);
-                               }
-                       }
-                       output = "";
-               } 
+           Model::Machine m(RecordDecl->getName());
+           Diag(RecordDecl->getLocStart(), diag_found_statemachine) << m.name;
+
+           if (CXXRecordDecl *InitialState = getTemplateArgDecl(Base->getType().getTypePtr(), 1))
+               m.setInitialState(InitialState->getName());
+           model.add(m);
        }
-               
-       void struct_class(const Decl *decl) // works with struct or class decl
+       else if (RecordDecl->isDerivedFrom("boost::statechart::event"))
        {
-               string output, line, ret, trans, event; 
-               llvm::raw_string_ostream x(output);
-               decl->print(x);
-               line = get_line_of_code(x.str());
-               output = "";
-               int pos;
-               const NamedDecl *namedDecl = dyn_cast<NamedDecl>(decl);
-               if(is_derived(line))
-               {
-                       const CXXRecordDecl *cRecDecl = dyn_cast<CXXRecordDecl>(decl);
-                                       
-                       if(find_events(cRecDecl, line))
-                       {
-                               events.push_back(namedDecl->getNameAsString());
-                               cout<<"New event: "<<namedDecl->getNameAsString()<<"\n";
-                       }
-                       else if(name_of_machine == "")
-                       {
-                               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 = find_states(cRecDecl, line);      
-                               if(!ret.empty())
-                               {       
-                                       cout << "New state: " << namedDecl->getNameAsString() << "\n";
-                                       states.push_back(ret);                  
-                                       methods_in_class(decl,namedDecl->getNameAsString());
-                               }
-                       }
-               }
+           //sc.events.push_back(RecordDecl->getNameAsString());
        }
+       return true;
+    }
+};
 
-       void methods_in_class(const Decl *decl, const string state)
-       {
-               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);              
-               //states.push_back(namedDecl->getNameAsString());
-               
-               output="";
-               for (DeclContext::decl_iterator i = declCont->decls_begin(), e = declCont->decls_end(); i != e; ++i) 
-               {
-                       if (i->getKind()==26) 
-                       {
-                               i->print(x);
-                               output = x.str();
-                               line = clean_spaces(cut_type(output));          
-                               ret = find_transitions(state,line);
-                               if(!ret.empty()) 
-                               {
-                                       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));
-                                               }
-                                               //cout<<ret<<"\n";
-                                               if(i!=num-1) ret = ret.substr(pos+1);
-                                       }
-                                       output="";
-                               }
-                       }
-                       if(i->getKind()==35) method_decl(decl);
-               }
-       }
 
-       void method_decl(const Decl *decl)
-       {
-               string output, line, event;     
-               llvm::raw_string_ostream x(output);
-               if(decl->hasBody())
-               {
-                       decl->print(x);
-                       line = get_return(x.str());
-                       if(test_model(line,"result"))
-                       {
-                               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++)
-                               {
-                                       event = *i;
-                                       if(line.compare(event)==0) 
-                                       {
-                                               cReactions.erase(i);
-                                               break;
-                                       }
-                               }
-                       }
-               }
-       }
+class VisualizeStatechartConsumer : public clang::ASTConsumer
+{
+    Model::Model model;
+    Visitor visitor;
+    string destFileName;
+public:
+    explicit VisualizeStatechartConsumer(ASTContext *Context, std::string destFileName,
+                                        DiagnosticsEngine &D)
+       : visitor(Context, model, D), destFileName(destFileName) {}
 
-       void find_return_stmt(Stmt *statemt,string event)
-       {
-               if(statemt->getStmtClass() == 99) test_stmt(dyn_cast<CaseStmt>(statemt)->getSubStmt(), event);
-               else
-               {
-                       for (Stmt::child_range range = statemt->children(); range; ++range)    
-                       {
-                               test_stmt(*range, event);
-                       }
-               }
-       }
-       
-       void test_stmt(Stmt *stmt, string event)
-       {
-               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(test_model(line,"transit"))
-                                                       {
-                                                               param = get_params(line);
-                                                               transitions.push_back(event.append(",").append(param));
-                                                       }
-                                                       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) {
+       visitor.TraverseDecl(Context.getTranslationUnitDecl());
+       model.write_as_dot_file(destFileName);
+    }
+};
+
+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 VisualizeStatechartConsumer(&CI.getASTContext(), dest, CI.getDiagnostics());
+  }
+
+  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());
+
+    return true;
+  }
+  void PrintHelp(llvm::raw_ostream& ros) {
+    ros << "Help for Visualize Statechart plugin goes here\n";
+  }
 
 };
 
-int main(int argc, char **argv)
-{ 
-       string inputFilename = "";
-       string outputFilename = "graph.dot"; // initialize output Filename
-       DiagnosticOptions dopts;
-       dopts.ShowColors=1;
-       MyDiagnosticClient *mdc = new MyDiagnosticClient(llvm::errs(), dopts);
-       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_BINDIR, llvm::sys::getHostTriple(), "", false, false, diag);
-       TheDriver.setCheckInputsExist(true);
-       TheDriver.CCCIsCXX = 1; 
-       TheDriver.ResourceDir = LLVM_PREFIX "/lib/clang/" CLANG_VERSION_STRING;
-
-       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) 
-               {
-                       if(strlen(Args[i])>2)
-                       {
-                               string str = Args[i];
-                               outputFilename = str.substr(2);
-                       }
-                       else outputFilename = Args[i+1];
-                       break;
-               }
-       }
-               
-       CompilerInvocation::CreateFromArgs(compInv,
-                                         const_cast<const char **>(CCArgs.data()),
-                                         const_cast<const char **>(CCArgs.data())+CCArgs.size(),
-                                         diag);
-
-       HeaderSearchOptions hsopts = compInv.getHeaderSearchOpts();
-       LangOptions lang = compInv.getLangOpts();
-       CompilerInvocation::setLangDefaults(lang, IK_ObjCXX);
-       TargetInfo *ti = TargetInfo::CreateTargetInfo(diag, compInv.getTargetOpts());
-       FrontendOptions f = compInv.getFrontendOpts();
-       inputFilename = f.Inputs[0].second;
-
-       cout<<"Input filename: "<<inputFilename<<"\n"; // print Input filename
-       cout<<"Output filename: "<<outputFilename<<"\n"; // print Output filename
-
-
-       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
-       IO_operations *io = new IO_operations(outputFilename, c.getStateMachine(), c.getNameOfFirstState(), c.getTransitions(), c.getStates(), c.getEvents());
-       io->save_to_file();
-       io->print_stats();
-       mdc->print_stats();
-       return 0;
-}
+static FrontendPluginRegistry::Add<VisualizeStatechartAction> X("visualize-statechart", "visualize statechart");
+
+// Local Variables:
+// c-basic-offset: 4
+// End: