lldb-tidy.cpp

#include "clang/AST/AST.h"
#include "clang/AST/ASTConsumer.h"
#include "clang/AST/RecursiveASTVisitor.h"
#include "clang/CodeGen/ObjectFilePCHContainerOperations.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/ADT/StringExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/raw_ostream.h"

#include <sstream>
#include <string>

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

static llvm::cl::OptionCategory InstrCategory("LLDB Instrumentation Generator");

class LLDBVisitor : public RecursiveASTVisitor<LLDBVisitor> {
public:
  LLDBVisitor(Rewriter &R, ASTContext &Context)
      : MyRewriter(R), Context(Context) {}

  bool VisitCXXMemberCallExpr(CXXMemberCallExpr *Expr) {
    auto *Decl = Expr->getMethodDecl();
    llvm::StringRef Name = Decl->getNameAsString();
    if ((Name.endswith("WithFormat") || Name.endswith("WithFormatv")) &&
        Expr->getNumArgs() == 1) {

      const auto *Callee = llvm::cast<MemberExpr>(Expr->getCallee());
      SourceRange R(Callee->getMemberLoc(), Callee->getMemberLoc());

      StringRef WrittenExpr = Lexer::getSourceText(
          CharSourceRange::getTokenRange(R), MyRewriter.getSourceMgr(),
          MyRewriter.getLangOpts());
      auto pos = WrittenExpr.rfind("WithFormat");
      StringRef ReplacementText = WrittenExpr.substr(0, pos);
      llvm::outs() << WrittenExpr << " -> " << ReplacementText << '\n';
      MyRewriter.ReplaceText(R, ReplacementText);
    }
    return true;
  }

private:
  Rewriter &MyRewriter;
  ASTContext &Context;
};

class LLDBConsumer : public ASTConsumer {
public:
  LLDBConsumer(Rewriter &R, ASTContext &Context) : Visitor(R, Context) {}

  virtual void HandleTranslationUnit(clang::ASTContext &Context) {
    Visitor.TraverseDecl(Context.getTranslationUnitDecl());
  }

private:
  LLDBVisitor Visitor;
};

class LLDBAction : public ASTFrontendAction {
public:
  LLDBAction() = default;

  std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                 StringRef File) override {
    MyRewriter.setSourceMgr(CI.getSourceManager(), CI.getLangOpts());
    return std::make_unique<LLDBConsumer>(MyRewriter, CI.getASTContext());
  }

  void EndSourceFileAction() override { MyRewriter.overwriteChangedFiles(); }

private:
  Rewriter MyRewriter;
};

int main(int argc, const char **argv) {
  CommonOptionsParser OP(argc, argv, InstrCategory,
                         "Utility for generating the macros for LLDB's "
                         "instrumentation framework.");

  auto PCHOpts = std::make_shared<PCHContainerOperations>();
  PCHOpts->registerWriter(std::make_unique<ObjectFilePCHContainerWriter>());
  PCHOpts->registerReader(std::make_unique<ObjectFilePCHContainerReader>());

  ClangTool T(OP.getCompilations(), OP.getSourcePathList(), PCHOpts);
  return T.run(newFrontendActionFactory<LLDBAction>().get());
}