1 //===- AnnotateFunctions.cpp ----------------------------------------------===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // Example clang plugin which adds an annotation to every function in
11 // translation units that start with #pragma enable_annotate.
12 //
13 //===----------------------------------------------------------------------===//
14 
15 #include "clang/Frontend/FrontendPluginRegistry.h"
16 #include "clang/AST/AST.h"
17 #include "clang/AST/ASTConsumer.h"
18 #include "clang/Lex/Preprocessor.h"
19 #include "clang/Lex/LexDiagnostic.h"
20 using namespace clang;
21 
22 namespace {
23 
24 static bool EnableAnnotate = false;
25 static bool HandledDecl = false;
26 
27 class AnnotateFunctionsConsumer : public ASTConsumer {
28 public:
HandleTopLevelDecl(DeclGroupRef DG)29   bool HandleTopLevelDecl(DeclGroupRef DG) override {
30     HandledDecl = true;
31     if (!EnableAnnotate)
32       return true;
33     for (auto D : DG)
34       if (FunctionDecl *FD = dyn_cast<FunctionDecl>(D))
35         FD->addAttr(AnnotateAttr::CreateImplicit(FD->getASTContext(),
36                                                  "example_annotation"));
37     return true;
38   }
39 };
40 
41 class AnnotateFunctionsAction : public PluginASTAction {
42 public:
CreateASTConsumer(CompilerInstance & CI,llvm::StringRef)43   std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
44                                                  llvm::StringRef) override {
45     return llvm::make_unique<AnnotateFunctionsConsumer>();
46   }
47 
ParseArgs(const CompilerInstance & CI,const std::vector<std::string> & args)48   bool ParseArgs(const CompilerInstance &CI,
49                  const std::vector<std::string> &args) override {
50     return true;
51   }
52 
getActionType()53   PluginASTAction::ActionType getActionType() override {
54     return AddBeforeMainAction;
55   }
56 };
57 
58 class PragmaAnnotateHandler : public PragmaHandler {
59 public:
PragmaAnnotateHandler()60   PragmaAnnotateHandler() : PragmaHandler("enable_annotate") { }
61 
HandlePragma(Preprocessor & PP,PragmaIntroducerKind Introducer,Token & PragmaTok)62   void HandlePragma(Preprocessor &PP, PragmaIntroducerKind Introducer,
63                     Token &PragmaTok) override {
64 
65     Token Tok;
66     PP.LexUnexpandedToken(Tok);
67     if (Tok.isNot(tok::eod))
68       PP.Diag(Tok, diag::ext_pp_extra_tokens_at_eol) << "pragma";
69 
70     if (HandledDecl) {
71       DiagnosticsEngine &D = PP.getDiagnostics();
72       unsigned ID = D.getCustomDiagID(
73         DiagnosticsEngine::Error,
74         "#pragma enable_annotate not allowed after declarations");
75       D.Report(PragmaTok.getLocation(), ID);
76     }
77 
78     EnableAnnotate = true;
79   }
80 };
81 
82 }
83 
84 static FrontendPluginRegistry::Add<AnnotateFunctionsAction>
85 X("annotate-fns", "annotate functions");
86 
87 static PragmaHandlerRegistry::Add<PragmaAnnotateHandler>
88 Y("enable_annotate","enable annotation");
89