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"
20using namespace clang;
21
22namespace {
23
24static bool EnableAnnotate = false;
25static bool HandledDecl = false;
26
27class AnnotateFunctionsConsumer : public ASTConsumer {
28public:
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
41class AnnotateFunctionsAction : public PluginASTAction {
42public:
43 std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
44 llvm::StringRef) override {
45 return llvm::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
58class PragmaAnnotateHandler : public PragmaHandler {
59public:
60 PragmaAnnotateHandler() : PragmaHandler("enable_annotate") { }
61
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
84static FrontendPluginRegistry::Add<AnnotateFunctionsAction>
85X("annotate-fns", "annotate functions");
86
87static PragmaHandlerRegistry::Add<PragmaAnnotateHandler>
88Y("enable_annotate","enable annotation");
89