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

在这一节中，对loop-convert的例子进行了修改，以展示对ForStmt的处理。

这里主要是通过AST树结构的判断，增加对类似for(int i=0;i<5;i++)这种循环的循环界判断，在实际程序优化中，意义不是特别大，但是作为示例和练习还是不错的。

主要使用的test case是：

int foo(int a, int b, int *c) {
    int ret = 0;
    if (a > b) {
        ret = a;
    } else {
        ret = b;
    }
    for (int temp = 0; temp < 100; ++temp) {
        *c = (*c + temp);
    }
    return ret;
}


int main() {
    int a = 1, b = 2;
    int c = 0;
    int d = foo(a, b, &c);
    return 0;
}

对ForStmt的处理在Visitor的VisitForStmt中进行（前排提示，不要尝试在RecursiveASTVisitor中对节点进行修改，大的节点修改会破坏整个AST树的Context，这里的Visitor仅建议作为遍历读取使用）

//------------------------------------------------------------------------------
// Tooling sample. Demonstrates:
//
// ForStmt Demo to show how to use ForStmt
//
// jourluohua (jourluohua@sina.com)
// This code is in the public domain
//------------------------------------------------------------------------------
#include <sstream>
#include <string>
#include <map>

#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("Tooling Sample");

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

  bool VisitForStmt(ForStmt *forStmt) {
    // Only care about ForStmt statements.
    // now we just judge for(int i=0;i<100;i++) case
    VarDecl* initVar = nullptr;
    Expr* initExpr = nullptr;
    LoopBound* lpBound = new LoopBound();
        DeclStmt* DS = nullptr;

    if ((DS = dyn_cast_or_null<DeclStmt>(dyn_cast<ForStmt>(forStmt)->getInit())) != 0 && DS->isSingleDecl() && 
               (initVar = dyn_cast<VarDecl>(DS->getSingleDecl())) != 0){
            initExpr = initVar->getInit();
    }
    else {
      return true;
    }
    IntegerLiteral* initInteg = dyn_cast_or_null<IntegerLiteral>(initExpr);
    lpBound->initVal = initInteg->getValue().getSExtValue();

    Expr *expCond = dyn_cast<ForStmt>(forStmt)->getCond();
    if(isa<BinaryOperator>(expCond))
    {
      BinaryOperator *binaryOpCond = dyn_cast<BinaryOperator>(expCond);
      if(binaryOpCond->getOpcode() ==BO_LT)
      {
        Expr * condExpr  = binaryOpCond ->getRHS();
        assert(condExpr && "condExpr  ");
        if(isa<IntegerLiteral>(condExpr))
        {
          lpBound->maxVal=(dyn_cast<IntegerLiteral>(condExpr)->getValue()).getSExtValue();
        }
        else
        {
          return true;
        }
      }
      else
      {
        return true;
      }
    }
    else
    {
      return true;
    }

    Expr *incExpr = dyn_cast<ForStmt>(forStmt)->getInc();
    assert(incExpr && "forStmt->getInc() cannot be NULL");
    if(isa<UnaryOperator>(incExpr)){    
      UnaryOperator *unaryInc = dyn_cast<UnaryOperator>(incExpr);
      if(unaryInc->isIncrementOp () )        //just like '++i' or 'i++', we think that they work same  
      {
        lpBound->incVal = 1;
      }
      else
      {
        return true;
      }
      
    }
    conditionMap[initVar] = lpBound;
    return true;
  }

  std::map<VarDecl*, LoopBound*>& getCondMap(){
    return conditionMap;
  }


private:
  Rewriter &TheRewriter;
  std::map<VarDecl*, LoopBound*> conditionMap;
};

// 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.
  bool HandleTopLevelDecl(DeclGroupRef DR) override {
    for (DeclGroupRef::iterator b = DR.begin(), e = DR.end(); b != e; ++b) {
      // Traverse the declaration using our AST visitor.
      Visitor.TraverseDecl(*b);
      (*b)->dump();
    }
    std::map<VarDecl*, LoopBound*> conditionMap = Visitor.getCondMap();
    if(conditionMap.size()>0){
      auto iter = conditionMap.begin();
      LoopBound *lpFirst = iter->second;
      VarDecl* varFirst = iter->first;
      llvm::errs()<<varFirst->getDeclName().getAsString()<<" ["<<lpFirst->initVal<<" "
                  <<lpFirst->maxVal <<" "<<lpFirst->incVal<<"]\n";
    }
    else{
      llvm::errs()<<"conditionMap.size()==0\n";
    }
    return true;
  }

private:
  MyASTVisitor Visitor;
};

// For each source file provided to the tool, a new FrontendAction is created.
class MyFrontendAction : public ASTFrontendAction {
public:
  MyFrontendAction() {}
  void EndSourceFileAction() override {
    SourceManager &SM = TheRewriter.getSourceMgr();
    llvm::errs() << "** EndSourceFileAction for: "
                 << SM.getFileEntryForID(SM.getMainFileID())->getName() << "\n";

    // Now emit the rewritten buffer.
    TheRewriter.getEditBuffer(SM.getMainFileID()).write(llvm::outs());
  }

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

private:
  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());
}

结果展示如下：