1 //===- AnnotateFunctions.cpp ----------------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // Example clang plugin which adds an annotation to every function in
10 // translation units that start with #pragma enable_annotate.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "clang/Frontend/FrontendPluginRegistry.h"
15 #include "clang/AST/AST.h"
16 #include "clang/AST/ASTConsumer.h"
17 #include "clang/AST/Attr.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:
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(
36             FD->getASTContext(), "example_annotation", nullptr, 0));
37     return true;
38   }
39 };
40 
41 class AnnotateFunctionsAction : public PluginASTAction {
42 public:
43   std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
44                                                  llvm::StringRef) override {
45     return std::make_unique<AnnotateFunctionsConsumer>();
46   }
47 
48   bool ParseArgs(const CompilerInstance &CI,
49                  const std::vector<std::string> &args) override {
50     return true;
51   }
52 
53   PluginASTAction::ActionType getActionType() override {
54     return AddBeforeMainAction;
55   }
56 };
57 
58 class PragmaAnnotateHandler : public PragmaHandler {
59 public:
60   PragmaAnnotateHandler() : PragmaHandler("enable_annotate") { }
61 
62   void HandlePragma(Preprocessor &PP, PragmaIntroducer 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