Clang教程之实现源源变化(3)

一个高级的源源变换，主要是为了展示如何进行AST Declaration的遍历和代码的规整化处理，可以将其作为一个代码自动规整化工具进行使用

实现代码如下：

//------------------------------------------------------------------------------
// Simple Rewrtie Sample 
//
//
// jourluohua (jourluohua@gmail.com)
// This code is in the public domain
//------------------------------------------------------------------------------
#include <sstream>
#include <string>
#include <fstream>

#include "clang/AST/AST.h"
#include "clang/AST/ASTConsumer.h"
#include "clang/AST/RecursiveASTVisitor.h"
#include "clang/Frontend/ASTConsumers.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Frontend/FrontendActions.h"
#include "clang/Rewrite/Core/Rewriter.h"
#include "clang/Tooling/CommonOptionsParser.h"
#include "clang/Tooling/Tooling.h"
#include "llvm/Support/raw_ostream.h"

using namespace clang;
using namespace clang::driver;
using namespace clang::tooling;

static llvm::cl::OptionCategory ToolingSampleCategory("Simple Rewriter");

// Implementation of the ASTConsumer interface for reading an AST produced
// by the Clang parser.
class MyASTConsumer : public ASTConsumer {
private:
    std::string InFile;
    SourceManager *SM;
    Rewriter TheRewriter;
    
  public:
    MyASTConsumer(SourceManager* sm, StringRef _InFile, Rewriter& theRew)
        : SM(sm), InFile(_InFile), TheRewriter(theRew){}

    void HandleTranslationUnit(ASTContext &Context) override;

    StringRef getInFile() { return InFile; }
};
void MyASTConsumer::HandleTranslationUnit(ASTContext &Context){
  DeclContext::decl_iterator D = Context.getTranslationUnitDecl()->decls_begin();
  DeclContext::decl_iterator DEnd = Context.getTranslationUnitDecl()->decls_end();
  PrintingPolicy Policy(Context.getLangOpts());
  while (D != DEnd)
  {
    SourceLocation Loc = SM->getExpansionLoc(D->getLocation());
    SourceLocation LocStart = D->getBeginLoc();
    SourceLocation LocEnd = D->getEndLoc();
    //(*D)->dump();
    if (Loc.isInvalid() || !SM->isWrittenInMainFile(Loc) || D->isImplicit() ||
        LocStart.isInvalid()) {
      ++D;
      continue;
    }

    std::string prettyBufS;
    llvm::raw_string_ostream prettyBuf(prettyBufS);
    D->print(prettyBuf, Policy);

    if (LocEnd.isInvalid()) {
      TheRewriter.InsertTextAfter(LocStart, "\n" + prettyBuf.str());
      ++D;
      continue;
    }

    const char *startBuf = SM->getCharacterData(LocStart);
    const char *endBuf = SM->getCharacterData(LocEnd);

    char separator;
    const char *separatorBuf;

    switch (D->getKind()) {
      case Decl::Function: {
        FunctionDecl *FD = cast<FunctionDecl>(*D);

        if (!FD->isThisDeclarationADefinition()) {
          separator = ';';
          separatorBuf = strchr(endBuf, separator);
          TheRewriter.ReplaceText(LocStart, separatorBuf - startBuf, prettyBuf.str());
        }
        else {
          separator = '}';
          separatorBuf = strchr(endBuf, separator);
          TheRewriter.ReplaceText(LocStart, separatorBuf - startBuf + 1, prettyBuf.str());
        }
        ++D;
        break;
      }
      case Decl::Var: {
        ++D;
        bool flag = true;
        while (D != DEnd && flag == true) {
          flag = false;
          if (VarDecl *VD = dyn_cast<VarDecl>(*D)) {
            if (VD->getBeginLoc() == LocStart) {
              prettyBuf << "; ";
              VD->print(prettyBuf, Policy);
              ++D;
              flag = true;
            }
          }
        }

        separator = ';';
        separatorBuf = strchr(endBuf, separator);
        TheRewriter.ReplaceText(LocStart, separatorBuf - startBuf, prettyBuf.str());
        break;
      }
      case Decl::TypeAlias:
      case Decl::Typedef: {
        separator = ';';
        separatorBuf = strchr(endBuf, separator);
        TheRewriter.ReplaceText(LocStart, separatorBuf - startBuf, prettyBuf.str());
        ++D;
        break;
      }
      case Decl::Record: {
        ++D;
        if (D != DEnd) {
          SourceLocation NextLocStart = D->getBeginLoc();
          SourceLocation NextLocEnd = D->getEndLoc();
          // Handle the following case:
          // typedef struct abc {
          //   ...;
          // } abc;
          if (NextLocStart.isValid() && NextLocStart < LocStart) {
            prettyBuf << ";\n";
            D->print(prettyBuf, Policy);
            startBuf = SM->getCharacterData(NextLocStart);
            endBuf = SM->getCharacterData(NextLocEnd);
            LocStart = NextLocStart;
            ++D;
          }
        }

        separator = ';';
        separatorBuf = strchr(endBuf, separator);
        TheRewriter.ReplaceText(LocStart, separatorBuf - startBuf, prettyBuf.str());
        break;
      }
      
      default:
        llvm_unreachable("should handle the decl kind");
    }
  }

  const RewriteBuffer *RewriteBuf = TheRewriter.getRewriteBufferFor(SM->getMainFileID());
  
  //*OutFile << std::string(RewriteBuf->begin(), RewriteBuf->end()); 
  std::ofstream outEdF("dpu_test.cpp", std::ofstream::out);

  if(outEdF.good()){
      std::string tmp(RewriteBuf->begin(), RewriteBuf->end()); 
      outEdF << tmp;
      //std::cout << tmp;
  }


  outEdF.close();


}
// For each source file provided to the tool, a new FrontendAction is created.
class MyFrontendAction : public ASTFrontendAction {
public:
  MyFrontendAction() {}


  std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                 StringRef inFile) override {
    llvm::errs() << "** Creating AST consumer for: " << inFile << "\n";
    TheRewriter.setSourceMgr(CI.getSourceManager(), CI.getLangOpts());
    return std::make_unique<MyASTConsumer>(&(CI.getSourceManager()), inFile, TheRewriter);
  }
  Rewriter TheRewriter;
};

int main(int argc, const char **argv) {
  llvm::Expected<CommonOptionsParser> op=CommonOptionsParser::create(argc, argv, ToolingSampleCategory);
  
  ClangTool Tool(op.get().getCompilations(), op.get().getSourcePathList());

  // ClangTool::run accepts a FrontendActionFactory, which is then used to
  // create new objects implementing the FrontendAction interface. Here we use
  // the helper newFrontendActionFactory to create a default factory that will
  // return a new MyFrontendAction object every time.
  // To further customize this, we could create our own factory class.
  return Tool.run(newFrontendActionFactory<MyFrontendAction>().get());
}