Clang rewriter comparison

[rvantonder]

// File: BugRefactor.cpp
// Auth: John Kotheimer (jkotheim@andrew.cmu.edu)
// Desc: Refactoring tool for source transformation for fuzzing resarch project

#include <fstream>
#include <regex>

#include "clang/AST/AST.h"
#include "clang/AST/ASTConsumer.h"
#include "clang/AST/ASTContext.h"
#include "clang/AST/RecursiveASTVisitor.h"
#include "clang/ASTMatchers/ASTMatchers.h"
#include "clang/ASTMatchers/ASTMatchFinder.h"
#include "clang/Basic/DiagnosticOptions.h"
#include "clang/Basic/Diagnostic.h"
#include "clang/Basic/FileManager.h"
#include "clang/Basic/SourceManager.h"
#include "clang/Frontend/ASTConsumers.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Frontend/FrontendActions.h"
#include "clang/Frontend/TextDiagnosticPrinter.h"
#include "clang/Rewrite/Core/Rewriter.h"
#include "clang/Tooling/CommonOptionsParser.h"
#include "clang/Tooling/Refactoring.h"
#include "clang/Tooling/Tooling.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/YAMLParser.h"

using namespace clang;
using namespace clang::ast_matchers;
using namespace clang::driver;
using namespace clang::tooling;
using namespace llvm;
// using namespace llvm::yaml;

// Set up help text and command line options
static llvm::cl::OptionCategory MyToolCategory("bug-refactor options");
static cl::extrahelp CommonHelp(CommonOptionsParser::HelpMessage);
// TODO add more help text if needed
// static cl::extrahelp MoreHelp("\nMore help text...");
// Type of bug we are repairing as a string (used to deduce the template file)
static cl::opt<std::string> BugType("bug", cl::cat(MyToolCategory), cl::ValueExpected(cl::ValueRequired),
                                    cl::desc("Type of bug to correct. Should have the same filename as the patch template."));
// Program variable we are correcting
static cl::opt<std::string> PVar("pvar", cl::cat(MyToolCategory), cl::ValueExpected(cl::ValueRequired),
                                 cl::desc("Program variable to repair."));
// Line number we are trying to patch
static cl::opt<unsigned> LineNum("l", cl::cat(MyToolCategory), cl::ValueExpected(cl::ValueRequired),
                            cl::desc("Line containing the bug to be repaired."));
// Path to the YAML config file (not yet implemented)
static cl::opt<std::string> YAMLPath("c", cl::cat(MyToolCategory),
                                     cl::desc("YAML configuration file. (Not yet implemented.)"));
// Debug mode flag to output extra information
static cl::opt<bool> DebugMode("d", cl::cat(MyToolCategory),
                                cl::desc("Debug mode. If true, outputs extra debug info."));

// Class to track the state of refactoring operations
class RefactorSession {
    public:
        RefactorSession(std::string bug, std::string pvar, unsigned line) :
            BugType(bug), PVar(pvar), RefactorLine(line) {}

        void genPatchCode(std::string path) {
            std::ifstream infile(path, std::ios::in | std::ios::binary);
            std::string contents;

            // read the entire patch file
            if(infile) {
                infile.seekg(0, std::ios::end);
                contents.resize(infile.tellg());
                infile.seekg(0, std::ios::beg);
                infile.read(&contents[0], contents.size());
                infile.close();
            } else {
                llvm::errs() << "Fatal: could not read patch file.\n";
                exit(1);
            }

            // perform a regex substitution for the correct program variable
            // TODO: we are using the tag @@PVAR@@ to denote this in a template
            // but may add the ability to customize this further in future versions
            std::regex pvar_regex("@@PVAR@@", std::regex_constants::ECMAScript);
            // replace the program var tag with the correct variable name and store it
            PatchCode = std::regex_replace(contents, pvar_regex, PVar);
        }

        // Accessor methods for the instance variables
        unsigned getLine() {
            return RefactorLine;
        }

        std::string getBugType() {
            return BugType;
        }

        std::string getPVar() {
            return PVar;
        }

        std::string getPatchCode() {
            return PatchCode;
        }

    private:
        std::string BugType;
        std::string PVar;
        std::string PatchCode;
        unsigned RefactorLine;
};

// The visitor visits all Stmts in the code and applies the patch
class FindLineVisitor : public RecursiveASTVisitor<FindLineVisitor> {
    public:
        explicit FindLineVisitor(ASTContext *context,
            std::map<std::string, Replacements> *replace,
            RefactorSession &session)
            : Context(context), Replace(replace), Session(session) {}

        bool VisitStmt(Stmt *st) {
            SourceManager &mgr = Context->getSourceManager();
            SourceLocation loc = st->getLocStart();
            // Check if this statement is on the line we're looking for
            if(mgr.getExpansionLineNumber(loc) == Session.getLine()) {
                Replacement Rep(mgr, loc, 0, Session.getPatchCode());
                Replacements Reps(Rep);
                // The Replacement interface takes a std::map of <filename, Replacements>
                std::string Filename = mgr.getFilename(loc);
                // Insert the patch if this is the correct location
                Replace->insert(std::pair<std::string,Replacements>(Filename, Reps));
            }
            return true;
        }

    private:
        ASTContext *Context;
        std::map<std::string, Replacements> *Replace;
        RefactorSession &Session;
};

// The ASTConsumer that defines where our visitor goes
class FindLineConsumer : public ASTConsumer {
    public:
        explicit FindLineConsumer(ASTContext *context,
            std::map<std::string, Replacements> *rep,
            RefactorSession &session) 
            : Visitor(context, rep, session) {}

        // HandleTranslationUnit is called when the ASTs
        // for the entire translation unit have been parsed (per LLVM docs)
        virtual void HandleTranslationUnit(ASTContext &context) {
            Visitor.TraverseDecl(context.getTranslationUnitDecl());
        }

    private:
        FindLineVisitor Visitor;
};

// The ASTFrontendAction that sets up the consumer
class FindLineAction : public clang::ASTFrontendAction {
    public:
        explicit FindLineAction(std::map<std::string, Replacements> *rep,
            RefactorSession &session)
            : Replace(rep), Session(session) {}

        virtual std::unique_ptr<clang::ASTConsumer> CreateASTConsumer(
            clang::CompilerInstance &Compiler, llvm::StringRef InFile) {
                return std::unique_ptr<clang::ASTConsumer>(
                    new FindLineConsumer(&Compiler.getASTContext(),
                        Replace, Session));
            }

    private:
        std::map<std::string, Replacements> *Replace;
        RefactorSession &Session;     
};

// We have to subclass FrontendActionFactory in order to invoke
// the proper constructor for FindLineAction
class FindLineActionFactory : public clang::tooling::FrontendActionFactory {
    public:
        FindLineActionFactory(std::map<std::string, Replacements> *rep,
            RefactorSession &session)
            : Replace(rep), Session(session) {} 

        // All we have to do is invoke the constructor within create()
        FrontendAction *create() {
            return new FindLineAction(Replace, Session);
        }

    private:
        std::map<std::string, Replacements> *Replace;
        RefactorSession &Session;
};

int main(int argc, const char **argv) {
    // Set up the tool and parse options
    CommonOptionsParser OptionsParser(argc, argv, MyToolCategory);
    RefactoringTool Tool(OptionsParser.getCompilations(),
        OptionsParser.getSourcePathList());

    // RefactorSession tracks internal state variables used during the refactoring session
    RefactorSession Session(BugType.getValue(), PVar.getValue(), LineNum.getValue());

    // Patch templates will have file names identical to the bug they fix
    // For example, the file "null_deref" could patch a null dereference bug
    Session.genPatchCode(BugType.getValue());
    // Print the code we are to insert, if debug mode is enabled
    if(DebugMode.getValue()) {
        llvm::outs() << "Generated patch code:\n";
        llvm::outs() << Session.getPatchCode() << "\n";
    }

    if(int res = Tool.run(new FindLineActionFactory(&Tool.getReplacements(), Session)))
        return res;

    // Apply the replacments generated by the tool
    // We need a SourceManager to set up the Rewriter
    IntrusiveRefCntPtr<DiagnosticOptions> DiagOpts = new DiagnosticOptions();
    DiagnosticsEngine Diagnostics(IntrusiveRefCntPtr<DiagnosticIDs>(new DiagnosticIDs()), &*DiagOpts,
                                  new TextDiagnosticPrinter(llvm::errs(), &*DiagOpts), true);
    SourceManager Sources(Diagnostics, Tool.getFiles());

    // Apply all replacements to a rewriter
    Rewriter Rewrite(Sources, LangOptions());
    Tool.applyAllReplacements(Rewrite);

    // Query the rewriter for all the files it has rewritten,
    // dumping their new contents to stdout
    for (Rewriter::buffer_iterator I = Rewrite.buffer_begin(),
                                   E = Rewrite.buffer_end();
                                   I != E; ++I) {
        const FileEntry *Entry = Sources.getFileEntryForID(I->first);
        if(DebugMode.getValue())
            llvm::outs() << "Rewrite buffer for file: " << Entry->getName() << "\n";
        I->second.write(llvm::outs());
    }

    return 0;
}

[rvantonder]

Moving to projects