175 lines · 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 20using namespace clang;21using namespace clang::driver;22using namespace clang::tooling;23 24static llvm::cl::OptionCategory InstrCategory("LLDB Instrumentation Generator");25 26class SBVisitor : public RecursiveASTVisitor<SBVisitor> {27public: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's33 // 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 76private: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 method92 // is declared in the implementation and is therefore not exposed as part93 // 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 122class SBConsumer : public ASTConsumer {123public:124 SBConsumer(Rewriter &R, ASTContext &Context) : Visitor(R, Context) {}125 126 // Override the method that gets called for each parsed top-level127 // 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 135private:136 SBVisitor Visitor;137};138 139class SBAction : public ASTFrontendAction {140public: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 153private:154 Rewriter MyRewriter;155};156 157int 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}175