1351290Sdim#include "clang/AST/AST.h"
2351290Sdim#include "clang/AST/ASTConsumer.h"
3351290Sdim#include "clang/AST/RecursiveASTVisitor.h"
4351290Sdim#include "clang/CodeGen/ObjectFilePCHContainerOperations.h"
5351290Sdim#include "clang/Frontend/ASTConsumers.h"
6351290Sdim#include "clang/Frontend/CompilerInstance.h"
7351290Sdim#include "clang/Frontend/FrontendActions.h"
8351290Sdim#include "clang/Rewrite/Core/Rewriter.h"
9351290Sdim#include "clang/Tooling/CommonOptionsParser.h"
10351290Sdim#include "clang/Tooling/Tooling.h"
11351290Sdim
12351290Sdim#include "llvm/ADT/StringExtras.h"
13351290Sdim#include "llvm/ADT/StringRef.h"
14351290Sdim#include "llvm/Support/raw_ostream.h"
15351290Sdim
16351290Sdim#include <sstream>
17351290Sdim#include <string>
18351290Sdim
19351290Sdimusing namespace clang;
20351290Sdimusing namespace clang::driver;
21351290Sdimusing namespace clang::tooling;
22351290Sdim
23351290Sdimstatic llvm::cl::OptionCategory InstrCategory("LLDB Instrumentation Generator");
24351290Sdim
25351290Sdim/// Get the macro name for recording method calls.
26351290Sdim///
27351290Sdim/// LLDB_RECORD_METHOD
28351290Sdim/// LLDB_RECORD_METHOD_CONST
29351290Sdim/// LLDB_RECORD_METHOD_NO_ARGS
30351290Sdim/// LLDB_RECORD_METHOD_CONST_NO_ARGS
31351290Sdim/// LLDB_RECORD_STATIC_METHOD
32351290Sdim/// LLDB_RECORD_STATIC_METHOD_NO_ARGS
33351290Sdimstatic std::string GetRecordMethodMacroName(bool Static, bool Const,
34351290Sdim                                            bool NoArgs) {
35351290Sdim  std::string Macro;
36351290Sdim  llvm::raw_string_ostream OS(Macro);
37351290Sdim
38351290Sdim  OS << "LLDB_RECORD";
39351290Sdim  if (Static)
40351290Sdim    OS << "_STATIC";
41351290Sdim  OS << "_METHOD";
42351290Sdim  if (Const)
43351290Sdim    OS << "_CONST";
44351290Sdim  if (NoArgs)
45351290Sdim    OS << "_NO_ARGS";
46351290Sdim
47351290Sdim  return OS.str();
48351290Sdim}
49351290Sdim
50351290Sdim/// Get the macro name for register methods.
51351290Sdim///
52351290Sdim/// LLDB_REGISTER_CONSTRUCTOR
53351290Sdim/// LLDB_REGISTER_METHOD
54351290Sdim/// LLDB_REGISTER_METHOD_CONST
55351290Sdim/// LLDB_REGISTER_STATIC_METHOD
56351290Sdimstatic std::string GetRegisterMethodMacroName(bool Static, bool Const) {
57351290Sdim  std::string Macro;
58351290Sdim  llvm::raw_string_ostream OS(Macro);
59351290Sdim
60351290Sdim  OS << "LLDB_REGISTER";
61351290Sdim  if (Static)
62351290Sdim    OS << "_STATIC";
63351290Sdim  OS << "_METHOD";
64351290Sdim  if (Const)
65351290Sdim    OS << "_CONST";
66351290Sdim
67351290Sdim  return OS.str();
68351290Sdim}
69351290Sdim
70351290Sdimstatic std::string GetRecordMethodMacro(StringRef Result, StringRef Class,
71351290Sdim                                        StringRef Method, StringRef Signature,
72351290Sdim                                        StringRef Values, bool Static,
73351290Sdim                                        bool Const) {
74351290Sdim  std::string Macro;
75351290Sdim  llvm::raw_string_ostream OS(Macro);
76351290Sdim
77351290Sdim  OS << GetRecordMethodMacroName(Static, Const, Values.empty());
78351290Sdim  OS << "(" << Result << ", " << Class << ", " << Method;
79351290Sdim
80351290Sdim  if (!Values.empty()) {
81351290Sdim    OS << ", (" << Signature << "), " << Values << ");\n\n";
82351290Sdim  } else {
83351290Sdim    OS << ");\n\n";
84351290Sdim  }
85351290Sdim
86351290Sdim  return OS.str();
87351290Sdim}
88351290Sdim
89351290Sdimstatic std::string GetRecordConstructorMacro(StringRef Class,
90351290Sdim                                             StringRef Signature,
91351290Sdim                                             StringRef Values) {
92351290Sdim  std::string Macro;
93351290Sdim  llvm::raw_string_ostream OS(Macro);
94351290Sdim  if (!Values.empty()) {
95351290Sdim    OS << "LLDB_RECORD_CONSTRUCTOR(" << Class << ", (" << Signature << "), "
96351290Sdim       << Values << ");\n\n";
97351290Sdim  } else {
98351290Sdim    OS << "LLDB_RECORD_CONSTRUCTOR_NO_ARGS(" << Class << ");\n\n";
99351290Sdim  }
100351290Sdim  return OS.str();
101351290Sdim}
102351290Sdim
103351290Sdimstatic std::string GetRecordDummyMacro(StringRef Result, StringRef Class,
104351290Sdim                                       StringRef Method, StringRef Signature,
105351290Sdim                                       StringRef Values) {
106351290Sdim  assert(!Values.empty());
107351290Sdim  std::string Macro;
108351290Sdim  llvm::raw_string_ostream OS(Macro);
109351290Sdim
110351290Sdim  OS << "LLDB_RECORD_DUMMY(" << Result << ", " << Class << ", " << Method;
111351290Sdim  OS << ", (" << Signature << "), " << Values << ");\n\n";
112351290Sdim
113351290Sdim  return OS.str();
114351290Sdim}
115351290Sdim
116351290Sdimstatic std::string GetRegisterConstructorMacro(StringRef Class,
117351290Sdim                                               StringRef Signature) {
118351290Sdim  std::string Macro;
119351290Sdim  llvm::raw_string_ostream OS(Macro);
120351290Sdim  OS << "LLDB_REGISTER_CONSTRUCTOR(" << Class << ", (" << Signature << "));\n";
121351290Sdim  return OS.str();
122351290Sdim}
123351290Sdim
124351290Sdimstatic std::string GetRegisterMethodMacro(StringRef Result, StringRef Class,
125351290Sdim                                          StringRef Method, StringRef Signature,
126351290Sdim                                          bool Static, bool Const) {
127351290Sdim  std::string Macro;
128351290Sdim  llvm::raw_string_ostream OS(Macro);
129351290Sdim  OS << GetRegisterMethodMacroName(Static, Const);
130351290Sdim  OS << "(" << Result << ", " << Class << ", " << Method << ", (" << Signature
131351290Sdim     << "));\n";
132351290Sdim  return OS.str();
133351290Sdim}
134351290Sdim
135351290Sdimclass SBReturnVisitor : public RecursiveASTVisitor<SBReturnVisitor> {
136351290Sdimpublic:
137351290Sdim  SBReturnVisitor(Rewriter &R) : MyRewriter(R) {}
138351290Sdim
139351290Sdim  bool VisitReturnStmt(ReturnStmt *Stmt) {
140351290Sdim    Expr *E = Stmt->getRetValue();
141351290Sdim
142351290Sdim    if (E->getBeginLoc().isMacroID())
143351290Sdim      return false;
144351290Sdim
145351290Sdim    SourceRange R(E->getBeginLoc(), E->getEndLoc());
146351290Sdim
147351290Sdim    StringRef WrittenExpr = Lexer::getSourceText(
148351290Sdim        CharSourceRange::getTokenRange(R), MyRewriter.getSourceMgr(),
149351290Sdim        MyRewriter.getLangOpts());
150351290Sdim
151351290Sdim    std::string ReplacementText =
152351290Sdim        "LLDB_RECORD_RESULT(" + WrittenExpr.str() + ")";
153351290Sdim    MyRewriter.ReplaceText(R, ReplacementText);
154351290Sdim
155351290Sdim    return true;
156351290Sdim  }
157351290Sdim
158351290Sdimprivate:
159351290Sdim  Rewriter &MyRewriter;
160351290Sdim};
161351290Sdim
162351290Sdimclass SBVisitor : public RecursiveASTVisitor<SBVisitor> {
163351290Sdimpublic:
164351290Sdim  SBVisitor(Rewriter &R, ASTContext &Context)
165351290Sdim      : MyRewriter(R), Context(Context) {}
166351290Sdim
167351290Sdim  bool VisitCXXMethodDecl(CXXMethodDecl *Decl) {
168351290Sdim    // Not all decls should be registered. Please refer to that method's
169351290Sdim    // comment for details.
170351290Sdim    if (ShouldSkip(Decl))
171351290Sdim      return false;
172351290Sdim
173351290Sdim    // Skip CXXMethodDecls that already starts with a macro. This should make
174351290Sdim    // it easier to rerun the tool to find missing macros.
175351290Sdim    Stmt *Body = Decl->getBody();
176351290Sdim    for (auto &C : Body->children()) {
177351290Sdim      if (C->getBeginLoc().isMacroID())
178351290Sdim        return false;
179351290Sdim      break;
180351290Sdim    }
181351290Sdim
182351290Sdim    // Print 'bool' instead of '_Bool'.
183351290Sdim    PrintingPolicy Policy(Context.getLangOpts());
184351290Sdim    Policy.Bool = true;
185351290Sdim
186351290Sdim    // Unsupported signatures get a dummy macro.
187351290Sdim    bool ShouldInsertDummy = false;
188351290Sdim
189351290Sdim    // Collect the functions parameter types and names.
190351290Sdim    std::vector<std::string> ParamTypes;
191351290Sdim    std::vector<std::string> ParamNames;
192351290Sdim    for (auto *P : Decl->parameters()) {
193351290Sdim      QualType T = P->getType();
194351290Sdim      ParamTypes.push_back(T.getAsString(Policy));
195351290Sdim      ParamNames.push_back(P->getNameAsString());
196351290Sdim
197351290Sdim      // Currently we don't support functions that have void pointers or
198351290Sdim      // function pointers as an argument, in which case we insert a dummy
199351290Sdim      // macro.
200351290Sdim      ShouldInsertDummy |= T->isFunctionPointerType() || T->isVoidPointerType();
201351290Sdim    }
202351290Sdim
203351290Sdim    // Convert the two lists to string for the macros.
204351290Sdim    std::string ParamTypesStr = llvm::join(ParamTypes, ", ");
205351290Sdim    std::string ParamNamesStr = llvm::join(ParamNames, ", ");
206351290Sdim
207351290Sdim    CXXRecordDecl *Record = Decl->getParent();
208351290Sdim    QualType ReturnType = Decl->getReturnType();
209351290Sdim
210351290Sdim    // Construct the macros.
211351290Sdim    std::string Macro;
212351290Sdim    if (ShouldInsertDummy) {
213351290Sdim      // Don't insert a register call for dummy macros.
214351290Sdim      Macro = GetRecordDummyMacro(
215351290Sdim          ReturnType.getAsString(Policy), Record->getNameAsString(),
216351290Sdim          Decl->getNameAsString(), ParamTypesStr, ParamNamesStr);
217351290Sdim
218351290Sdim    } else if (isa<CXXConstructorDecl>(Decl)) {
219351290Sdim      llvm::outs() << GetRegisterConstructorMacro(Record->getNameAsString(),
220351290Sdim                                                  ParamTypesStr);
221351290Sdim
222351290Sdim      Macro = GetRecordConstructorMacro(Record->getNameAsString(),
223351290Sdim                                        ParamTypesStr, ParamNamesStr);
224351290Sdim    } else {
225351290Sdim      llvm::outs() << GetRegisterMethodMacro(
226351290Sdim          ReturnType.getAsString(Policy), Record->getNameAsString(),
227351290Sdim          Decl->getNameAsString(), ParamTypesStr, Decl->isStatic(),
228351290Sdim          Decl->isConst());
229351290Sdim
230351290Sdim      Macro = GetRecordMethodMacro(
231351290Sdim          ReturnType.getAsString(Policy), Record->getNameAsString(),
232351290Sdim          Decl->getNameAsString(), ParamTypesStr, ParamNamesStr,
233351290Sdim          Decl->isStatic(), Decl->isConst());
234351290Sdim    }
235351290Sdim
236351290Sdim    // Insert the macro at the beginning of the function. We don't attempt to
237351290Sdim    // fix the formatting and instead rely on clang-format to fix it after the
238351290Sdim    // tool has run. This is also the reason that the macros end with two
239351290Sdim    // newlines, counting on clang-format to normalize this in case the macro
240351290Sdim    // got inserted before an existing newline.
241351290Sdim    SourceLocation InsertLoc = Lexer::getLocForEndOfToken(
242351290Sdim        Body->getBeginLoc(), 0, MyRewriter.getSourceMgr(),
243351290Sdim        MyRewriter.getLangOpts());
244351290Sdim    MyRewriter.InsertTextAfter(InsertLoc, Macro);
245351290Sdim
246351290Sdim    // If the function returns a class or struct, we need to wrap its return
247351290Sdim    // statement(s).
248351290Sdim    bool ShouldRecordResult = ReturnType->isStructureOrClassType() ||
249351290Sdim                              ReturnType->getPointeeCXXRecordDecl();
250351290Sdim    if (!ShouldInsertDummy && ShouldRecordResult) {
251351290Sdim      SBReturnVisitor Visitor(MyRewriter);
252351290Sdim      Visitor.TraverseDecl(Decl);
253351290Sdim    }
254351290Sdim
255351290Sdim    return true;
256351290Sdim  }
257351290Sdim
258351290Sdimprivate:
259351290Sdim  /// Determine whether we need to consider the given CXXMethodDecl.
260351290Sdim  ///
261351290Sdim  /// Currently we skip the following cases:
262351290Sdim  ///  1. Decls outside the main source file,
263351290Sdim  ///  2. Decls that are only present in the source file,
264351290Sdim  ///  3. Decls that are not definitions,
265351290Sdim  ///  4. Non-public methods,
266351290Sdim  ///  5. Variadic methods.
267351290Sdim  ///  6. Destructors.
268351290Sdim  bool ShouldSkip(CXXMethodDecl *Decl) {
269351290Sdim    // Skip anything outside the main file.
270351290Sdim    if (!MyRewriter.getSourceMgr().isInMainFile(Decl->getBeginLoc()))
271351290Sdim      return true;
272351290Sdim
273351290Sdim    // Skip if the canonical decl in the current decl. It means that the method
274351290Sdim    // is declared in the implementation and is therefore not exposed as part
275351290Sdim    // of the API.
276351290Sdim    if (Decl == Decl->getCanonicalDecl())
277351290Sdim      return true;
278351290Sdim
279351290Sdim    // Skip decls that have no body, i.e. are just declarations.
280351290Sdim    Stmt *Body = Decl->getBody();
281351290Sdim    if (!Body)
282351290Sdim      return true;
283351290Sdim
284351290Sdim    // Skip non-public methods.
285351290Sdim    AccessSpecifier AS = Decl->getAccess();
286351290Sdim    if (AS != AccessSpecifier::AS_public)
287351290Sdim      return true;
288351290Sdim
289351290Sdim    // Skip variadic methods.
290351290Sdim    if (Decl->isVariadic())
291351290Sdim      return true;
292351290Sdim
293351290Sdim    // Skip destructors.
294351290Sdim    if (isa<CXXDestructorDecl>(Decl))
295351290Sdim      return true;
296351290Sdim
297351290Sdim    return false;
298351290Sdim  }
299351290Sdim
300351290Sdim  Rewriter &MyRewriter;
301351290Sdim  ASTContext &Context;
302351290Sdim};
303351290Sdim
304351290Sdimclass SBConsumer : public ASTConsumer {
305351290Sdimpublic:
306351290Sdim  SBConsumer(Rewriter &R, ASTContext &Context) : Visitor(R, Context) {}
307351290Sdim
308351290Sdim  // Override the method that gets called for each parsed top-level
309351290Sdim  // declaration.
310351290Sdim  bool HandleTopLevelDecl(DeclGroupRef DR) override {
311351290Sdim    for (DeclGroupRef::iterator b = DR.begin(), e = DR.end(); b != e; ++b) {
312351290Sdim      Visitor.TraverseDecl(*b);
313351290Sdim    }
314351290Sdim    return true;
315351290Sdim  }
316351290Sdim
317351290Sdimprivate:
318351290Sdim  SBVisitor Visitor;
319351290Sdim};
320351290Sdim
321351290Sdimclass SBAction : public ASTFrontendAction {
322351290Sdimpublic:
323351290Sdim  SBAction() = default;
324351290Sdim
325351290Sdim  bool BeginSourceFileAction(CompilerInstance &CI) override {
326351290Sdim    llvm::outs() << "{\n";
327351290Sdim    return true;
328351290Sdim  }
329351290Sdim
330351290Sdim  void EndSourceFileAction() override {
331351290Sdim    llvm::outs() << "}\n";
332351290Sdim    MyRewriter.overwriteChangedFiles();
333351290Sdim  }
334351290Sdim
335351290Sdim  std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
336351290Sdim                                                 StringRef File) override {
337351290Sdim    MyRewriter.setSourceMgr(CI.getSourceManager(), CI.getLangOpts());
338360784Sdim    return std::make_unique<SBConsumer>(MyRewriter, CI.getASTContext());
339351290Sdim  }
340351290Sdim
341351290Sdimprivate:
342351290Sdim  Rewriter MyRewriter;
343351290Sdim};
344351290Sdim
345351290Sdimint main(int argc, const char **argv) {
346351290Sdim  CommonOptionsParser OP(argc, argv, InstrCategory,
347351290Sdim                         "Utility for generating the macros for LLDB's "
348351290Sdim                         "instrumentation framework.");
349351290Sdim
350351290Sdim  auto PCHOpts = std::make_shared<PCHContainerOperations>();
351360784Sdim  PCHOpts->registerWriter(std::make_unique<ObjectFilePCHContainerWriter>());
352360784Sdim  PCHOpts->registerReader(std::make_unique<ObjectFilePCHContainerReader>());
353351290Sdim
354351290Sdim  ClangTool T(OP.getCompilations(), OP.getSourcePathList(), PCHOpts);
355351290Sdim  return T.run(newFrontendActionFactory<SBAction>().get());
356351290Sdim}
357