rewritersample.cpp

//-------------------------------------------------------------------------
//
// rewritersample.cpp: Source-to-source transformation sample with Clang,
// using Rewriter - the code rewriting interface.
//
// This code is in the public domain
//
#include <cstdio>
#include <string>
#include <sstream>
#include <iostream>
#include <fstream>

#include "clang/AST/ASTConsumer.h"
#include "clang/AST/RecursiveASTVisitor.h"
#include "clang/Basic/Diagnostic.h"
#include "clang/Basic/FileManager.h"
#include "clang/Basic/SourceManager.h"
#include "clang/Basic/TargetOptions.h"
#include "clang/Basic/TargetInfo.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Lex/Preprocessor.h"
#include "clang/Parse/ParseAST.h"
#include "clang/Rewrite/Core/Rewriter.h"
#include "clang/Rewrite/Frontend/Rewriters.h"
#include "llvm/Support/Host.h"
#include "llvm/Support/raw_ostream.h"

using namespace clang;
using namespace std;

CompilerInstance* pTheCompInst_g;
std::set<int> unexecutedlines_g;

void printheader (stringstream& mystr) {
return;

mystr << "\
<html>\
<head>\
\
<script language=\"javascript\"> \
    function toggle(xyz) {\
        var displayText = \"displayText\"+xyz;\
        var toggleText = \"toggleText\"+xyz;\
\
    var ele = document.getElementById(toggleText);\
    var text = document.getElementById(displayText);\
    if(ele.style.display == \"block\") {\
            ele.style.display = \"none\";\
        text.innerHTML = \"show\";\
    }\
    else {\
        ele.style.display = \"block\";\
        text.innerHTML = \"hide\";\
    }\
} \
</script>\
</head>\
\
<body>\
    \
<pre>\
";
    return;

}

// By implementing RecursiveASTVisitor, we can specify which AST nodes
// we're interested in by overriding relevant methods.
class MyASTVisitor : public RecursiveASTVisitor<MyASTVisitor>
{
public:
    MyASTVisitor(Rewriter &R)
        : TheRewriter(R)
    {}

    void printbefore(stringstream& str,int id) {
    return;

        str << "<a id=\"displayText"<<id<<"\" href=\"javascript:toggle("<<id<<");\">show</a> <== click Here\n"
            << "<div id=\"toggleText"<<id<<"\" style=\"display: none\">\n";

         return;
    }
    void printafter (stringstream& str) {
    return;
        str << "</div>\n";

        return;
    }

    bool VisitStmt(Stmt *s) {
        // Only care about If statements.
        static int divid = 0;
        if (isa<IfStmt>(s)) {
            IfStmt *IfStatement = cast<IfStmt>(s);
            Stmt *Then = IfStatement->getThen();

            PresumedLoc PLoc = pTheCompInst_g->getSourceManager().getPresumedLoc(IfStatement->getLocStart(), 1);//DiagOpts->ShowPresumedLoc);
            cout << "Line number is " << PLoc.getLine() << endl;
            cout << "Line number is " << pTheCompInst_g->getSourceManager().getSpellingLineNumber(IfStatement->getLocStart(), 0) << endl;
            cout << "Line number is " << pTheCompInst_g->getSourceManager().getExpansionLineNumber(IfStatement->getLocStart(), 0) << endl;
            cout << "Line number is " << pTheCompInst_g->getSourceManager().getPresumedLineNumber(IfStatement->getLocStart(), 0) << endl;
            cout << "-----------------\n";


            stringstream str_before;
            //printbefore(str_before, divid);
            divid++;
            TheRewriter.InsertText(Then->getLocStart(), str_before.str(), true,true);

            SourceLocation END = Then->getLocEnd();
            int offset = Lexer::MeasureTokenLength(END,
                                           pTheCompInst_g->getSourceManager(),
                                           pTheCompInst_g->getLangOpts()) + 1;

            SourceLocation END1 = END.getLocWithOffset(offset);

            stringstream str_after;
            //printafter(str_after);
            TheRewriter.InsertText(END1,
                                   str_after.str(),
                                   true, true);


            Stmt *Else = IfStatement->getElse();
            if (Else) {
                SourceLocation END = Else->getLocEnd();
                int offset = Lexer::MeasureTokenLength(END,
                                           pTheCompInst_g->getSourceManager(),
                                           pTheCompInst_g->getLangOpts()) + 1;

                SourceLocation END1 = END.getLocWithOffset(offset);

                TheRewriter.InsertText(Else->getLocStart(),
                                       "// the 'else' part\n",
                                       true, true);
            }
        }

        return true;
    }

    bool VisitFunctionDecl(FunctionDecl *f) {
        // Only function definitions (with bodies), not declarations.
        if (f->hasBody()) {
            Stmt *FuncBody = f->getBody();

            // Type name as string
            QualType QT = f->getResultType();
            string TypeStr = QT.getAsString();

            // Function name
            DeclarationName DeclName = f->getNameInfo().getName();
            string FuncName = DeclName.getAsString();

            // Add comment before
            stringstream SSBefore;
            SSBefore << "// Begin function " << FuncName << " returning "
                     << TypeStr << "\n";
            SourceLocation ST = f->getSourceRange().getBegin();
            TheRewriter.InsertText(ST, SSBefore.str(), true, true);

            // And after
            stringstream SSAfter;
            SSAfter << "\n// End function " << FuncName << "\n";
            ST = FuncBody->getLocEnd().getLocWithOffset(1);
            TheRewriter.InsertText(ST, SSAfter.str(), true, true);
        }

        return true;
    }

private:
    void AddBraces(Stmt *s);

    Rewriter &TheRewriter;
};


// Implementation of the ASTConsumer interface for reading an AST produced
// by the Clang parser.
class MyASTConsumer : public ASTConsumer
{
public:
    MyASTConsumer(Rewriter &R)
        : Visitor(R)
    {}

    // Override the method that gets called for each parsed top-level
    // declaration.
    virtual bool HandleTopLevelDecl(DeclGroupRef DR) {
        for (DeclGroupRef::iterator b = DR.begin(), e = DR.end();
             b != e; ++b)
            // Traverse the declaration using our AST visitor.
            Visitor.TraverseDecl(*b);
        return true;
    }

private:
    MyASTVisitor Visitor;
};


int main(int argc, char *argv[])
{
    if (argc != 2) {
        llvm::errs() << "Usage: rewritersample <filename>\n";
        //llvm::errs() << "Usage: rewritersample <filename> <gcov-file>\n";
        return 1;
    }

    // CompilerInstance will hold the instance of the Clang compiler for us,
    // managing the various objects needed to run the compiler.
    CompilerInstance TheCompInst;
    pTheCompInst_g = &TheCompInst;

    /*
    LangOptions& Opts = TheCompInst.getLangOpts();
    TheCompInst.getInvocation().setLangDefaults (Opts, IK_CXX);
    */

    //TheCompInst.createDiagnostics(0, 0);
    TheCompInst.createDiagnostics(NULL, false);

    // Initialize target info with the default triple for our platform.
    TargetOptions TO;
    TO.Triple = llvm::sys::getDefaultTargetTriple();
    TargetInfo *TI = TargetInfo::CreateTargetInfo(
        TheCompInst.getDiagnostics(), &TO);
    TheCompInst.setTarget(TI);

    TheCompInst.createFileManager();
    FileManager &FileMgr = TheCompInst.getFileManager();
    TheCompInst.createSourceManager(FileMgr);
    SourceManager &SourceMgr = TheCompInst.getSourceManager();
    TheCompInst.createPreprocessor();
    TheCompInst.createASTContext();

    // A Rewriter helps us manage the code rewriting task.
    Rewriter TheRewriter;
    TheRewriter.setSourceMgr(SourceMgr, TheCompInst.getLangOpts());

#if 0
    //Parse the gcov file.
    std::ifstream infile (argv[2]);
    string line;
    int linenum = 0;
    while (std::getline(infile,line)) {
        //cout << line << endl;
        linenum++;
        const char* str = line.c_str();
        if (!strncmp (str,"    #####:", strlen ("    #####:"))) {
            unexecutedlines_g.insert (linenum);
        }
    }
#endif

    // Set the main file handled by the source manager to the input file.
    const FileEntry *FileIn = FileMgr.getFile(argv[1]);
    SourceMgr.createMainFileID(FileIn);
    TheCompInst.getDiagnosticClient().BeginSourceFile(
        TheCompInst.getLangOpts(),
        &TheCompInst.getPreprocessor());

    // Create an AST consumer instance which is going to get called by
    // ParseAST.
    MyASTConsumer TheConsumer(TheRewriter);

    // Parse the file to AST, registering our consumer as the AST consumer.
    ParseAST(TheCompInst.getPreprocessor(), &TheConsumer,
             TheCompInst.getASTContext());


    stringstream str;
    printheader (str);
    cout << str.str() <<endl;

    // At this point the rewriter's buffer should be full with the rewritten
    // file contents.
    const RewriteBuffer *RewriteBuf =
        TheRewriter.getRewriteBufferFor(SourceMgr.getMainFileID());
    llvm::outs() << string(RewriteBuf->begin(), RewriteBuf->end());

    llvm::outs() << string("</pre></body></html>\n");
    exit (-1);
    return 0;
}