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 }
174