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