lldb/tools/lldb-instr/Instrument.cpp

   1 #include "clang/AST/AST.h"
   2 #include "clang/AST/ASTConsumer.h"
   3 #include "clang/AST/RecursiveASTVisitor.h"
   4 #include "clang/CodeGen/ObjectFilePCHContainerWriter.h"
   5 #include "clang/Frontend/ASTConsumers.h"
   6 #include "clang/Frontend/CompilerInstance.h"
   7 #include "clang/Frontend/FrontendActions.h"
   8 #include "clang/Rewrite/Core/Rewriter.h"
   9 #include "clang/Serialization/ObjectFilePCHContainerReader.h"
  10 #include "clang/Tooling/CommonOptionsParser.h"
  11 #include "clang/Tooling/Tooling.h"
  12
  13 #include "llvm/ADT/StringExtras.h"
  14 #include "llvm/ADT/StringRef.h"
  15 #include "llvm/Support/raw_ostream.h"
  16
  17 #include <sstream>
  18 #include <string>
  19
  20 using namespace clang;
  21 using namespace clang::driver;
  22 using namespace clang::tooling;
  23
  24 static llvm::cl::OptionCategory InstrCategory("LLDB Instrumentation Generator");
  25
  26 class SBVisitor : public RecursiveASTVisitor<SBVisitor> {
  27 public:
  28   SBVisitor(Rewriter &R, ASTContext &Context)
  29       : MyRewriter(R), Context(Context) {}
  30
  31   bool VisitCXXMethodDecl(CXXMethodDecl *Decl) {
  32     // Not all decls should be registered. Please refer to that method's
  33     // comment for details.
  34     if (ShouldSkip(Decl))
  35       return false;
  36
  37     // Print 'bool' instead of '_Bool'.
  38     PrintingPolicy Policy(Context.getLangOpts());
  39     Policy.Bool = true;
  40
  41     // Collect the functions parameter types and names.
  42     std::vector<std::string> ParamNames;
  43     if (!Decl->isStatic())
  44       ParamNames.push_back("this");
  45     for (auto *P : Decl->parameters())
  46       ParamNames.push_back(P->getNameAsString());
  47
  48     // Construct the macros.
  49     std::string Buffer;
  50     llvm::raw_string_ostream Macro(Buffer);
  51     if (ParamNames.empty()) {
  52       Macro << "LLDB_INSTRUMENT()";
  53     } else {
  54       Macro << "LLDB_INSTRUMENT_VA(" << llvm::join(ParamNames, ", ") << ")";
  55     }
  56
  57     Stmt *Body = Decl->getBody();
  58     for (auto &C : Body->children()) {
  59       if (C->getBeginLoc().isMacroID()) {
  60         CharSourceRange Range =
  61             MyRewriter.getSourceMgr().getExpansionRange(C->getSourceRange());
  62         MyRewriter.ReplaceText(Range, Buffer);
  63       } else {
  64         Macro << ";";
  65         SourceLocation InsertLoc = Lexer::getLocForEndOfToken(
  66             Body->getBeginLoc(), 0, MyRewriter.getSourceMgr(),
  67             MyRewriter.getLangOpts());
  68         MyRewriter.InsertTextAfter(InsertLoc, Buffer);
  69       }
  70       break;
  71     }
  72
  73     return true;
  74   }
  75
  76 private:
  77   /// Determine whether we need to consider the given CXXMethodDecl.
  78   ///
  79   /// Currently we skip the following cases:
  80   ///  1. Decls outside the main source file,
  81   ///  2. Decls that are only present in the source file,
  82   ///  3. Decls that are not definitions,
  83   ///  4. Non-public methods,
  84   ///  5. Variadic methods.
  85   ///  6. Destructors.
  86   bool ShouldSkip(CXXMethodDecl *Decl) {
  87     // Skip anything outside the main file.
  88     if (!MyRewriter.getSourceMgr().isInMainFile(Decl->getBeginLoc()))
  89       return true;
  90
  91     // Skip if the canonical decl in the current decl. It means that the method
  92     // is declared in the implementation and is therefore not exposed as part
  93     // of the API.
  94     if (Decl == Decl->getCanonicalDecl())
  95       return true;
  96
  97     // Skip decls that have no body, i.e. are just declarations.
  98     Stmt *Body = Decl->getBody();
  99     if (!Body)
 100       return true;
 101
 102     // Skip non-public methods.
 103     AccessSpecifier AS = Decl->getAccess();
 104     if (AS != AccessSpecifier::AS_public)
 105       return true;
 106
 107     // Skip variadic methods.
 108     if (Decl->isVariadic())
 109       return true;
 110
 111     // Skip destructors.
 112     if (isa<CXXDestructorDecl>(Decl))
 113       return true;
 114
 115     return false;
 116   }
 117
 118   Rewriter &MyRewriter;
 119   ASTContext &Context;
 120 };
 121
 122 class SBConsumer : public ASTConsumer {
 123 public:
 124   SBConsumer(Rewriter &R, ASTContext &Context) : Visitor(R, Context) {}
 125
 126   // Override the method that gets called for each parsed top-level
 127   // declaration.
 128   bool HandleTopLevelDecl(DeclGroupRef DR) override {
 129     for (DeclGroupRef::iterator b = DR.begin(), e = DR.end(); b != e; ++b) {
 130       Visitor.TraverseDecl(*b);
 131     }
 132     return true;
 133   }
 134
 135 private:
 136   SBVisitor Visitor;
 137 };
 138
 139 class SBAction : public ASTFrontendAction {
 140 public:
 141   SBAction() = default;
 142
 143   bool BeginSourceFileAction(CompilerInstance &CI) override { return true; }
 144
 145   void EndSourceFileAction() override { MyRewriter.overwriteChangedFiles(); }
 146
 147   std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
 148                                                  StringRef File) override {
 149     MyRewriter.setSourceMgr(CI.getSourceManager(), CI.getLangOpts());
 150     return std::make_unique<SBConsumer>(MyRewriter, CI.getASTContext());
 151   }
 152
 153 private:
 154   Rewriter MyRewriter;
 155 };
 156
 157 int main(int argc, const char **argv) {
 158   auto ExpectedParser = CommonOptionsParser::create(
 159       argc, argv, InstrCategory, llvm::cl::OneOrMore,
 160       "Utility for generating the macros for LLDB's "
 161       "instrumentation framework.");
 162   if (!ExpectedParser) {
 163     llvm::errs() << ExpectedParser.takeError();
 164     return 1;
 165   }
 166   CommonOptionsParser &OP = ExpectedParser.get();
 167
 168   auto PCHOpts = std::make_shared<PCHContainerOperations>();
 169   PCHOpts->registerWriter(std::make_unique<ObjectFilePCHContainerWriter>());
 170   PCHOpts->registerReader(std::make_unique<ObjectFilePCHContainerReader>());
 171
 172   ClangTool T(OP.getCompilations(), OP.getSourcePathList(), PCHOpts);
 173   return T.run(newFrontendActionFactory<SBAction>().get());
 174 }