106f32e7eSjoerg //===- AnnotateFunctions.cpp ----------------------------------------------===//
206f32e7eSjoerg //
306f32e7eSjoerg // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
406f32e7eSjoerg // See https://llvm.org/LICENSE.txt for license information.
506f32e7eSjoerg // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
606f32e7eSjoerg //
706f32e7eSjoerg //===----------------------------------------------------------------------===//
806f32e7eSjoerg //
906f32e7eSjoerg // Example clang plugin which adds an annotation to every function in
1006f32e7eSjoerg // translation units that start with #pragma enable_annotate.
1106f32e7eSjoerg //
1206f32e7eSjoerg //===----------------------------------------------------------------------===//
1306f32e7eSjoerg 
1406f32e7eSjoerg #include "clang/Frontend/FrontendPluginRegistry.h"
1506f32e7eSjoerg #include "clang/AST/AST.h"
1606f32e7eSjoerg #include "clang/AST/ASTConsumer.h"
17*13fbcb42Sjoerg #include "clang/AST/Attr.h"
1806f32e7eSjoerg #include "clang/Lex/Preprocessor.h"
1906f32e7eSjoerg #include "clang/Lex/LexDiagnostic.h"
2006f32e7eSjoerg using namespace clang;
2106f32e7eSjoerg 
2206f32e7eSjoerg namespace {
2306f32e7eSjoerg 
2406f32e7eSjoerg static bool EnableAnnotate = false;
2506f32e7eSjoerg static bool HandledDecl = false;
2606f32e7eSjoerg 
2706f32e7eSjoerg class AnnotateFunctionsConsumer : public ASTConsumer {
2806f32e7eSjoerg public:
HandleTopLevelDecl(DeclGroupRef DG)2906f32e7eSjoerg   bool HandleTopLevelDecl(DeclGroupRef DG) override {
3006f32e7eSjoerg     HandledDecl = true;
3106f32e7eSjoerg     if (!EnableAnnotate)
3206f32e7eSjoerg       return true;
3306f32e7eSjoerg     for (auto D : DG)
3406f32e7eSjoerg       if (FunctionDecl *FD = dyn_cast<FunctionDecl>(D))
35*13fbcb42Sjoerg         FD->addAttr(AnnotateAttr::CreateImplicit(
36*13fbcb42Sjoerg             FD->getASTContext(), "example_annotation", nullptr, 0));
3706f32e7eSjoerg     return true;
3806f32e7eSjoerg   }
3906f32e7eSjoerg };
4006f32e7eSjoerg 
4106f32e7eSjoerg class AnnotateFunctionsAction : public PluginASTAction {
4206f32e7eSjoerg public:
CreateASTConsumer(CompilerInstance & CI,llvm::StringRef)4306f32e7eSjoerg   std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
4406f32e7eSjoerg                                                  llvm::StringRef) override {
4506f32e7eSjoerg     return std::make_unique<AnnotateFunctionsConsumer>();
4606f32e7eSjoerg   }
4706f32e7eSjoerg 
ParseArgs(const CompilerInstance & CI,const std::vector<std::string> & args)4806f32e7eSjoerg   bool ParseArgs(const CompilerInstance &CI,
4906f32e7eSjoerg                  const std::vector<std::string> &args) override {
5006f32e7eSjoerg     return true;
5106f32e7eSjoerg   }
5206f32e7eSjoerg 
getActionType()5306f32e7eSjoerg   PluginASTAction::ActionType getActionType() override {
5406f32e7eSjoerg     return AddBeforeMainAction;
5506f32e7eSjoerg   }
5606f32e7eSjoerg };
5706f32e7eSjoerg 
5806f32e7eSjoerg class PragmaAnnotateHandler : public PragmaHandler {
5906f32e7eSjoerg public:
PragmaAnnotateHandler()6006f32e7eSjoerg   PragmaAnnotateHandler() : PragmaHandler("enable_annotate") { }
6106f32e7eSjoerg 
HandlePragma(Preprocessor & PP,PragmaIntroducer Introducer,Token & PragmaTok)6206f32e7eSjoerg   void HandlePragma(Preprocessor &PP, PragmaIntroducer Introducer,
6306f32e7eSjoerg                     Token &PragmaTok) override {
6406f32e7eSjoerg 
6506f32e7eSjoerg     Token Tok;
6606f32e7eSjoerg     PP.LexUnexpandedToken(Tok);
6706f32e7eSjoerg     if (Tok.isNot(tok::eod))
6806f32e7eSjoerg       PP.Diag(Tok, diag::ext_pp_extra_tokens_at_eol) << "pragma";
6906f32e7eSjoerg 
7006f32e7eSjoerg     if (HandledDecl) {
7106f32e7eSjoerg       DiagnosticsEngine &D = PP.getDiagnostics();
7206f32e7eSjoerg       unsigned ID = D.getCustomDiagID(
7306f32e7eSjoerg         DiagnosticsEngine::Error,
7406f32e7eSjoerg         "#pragma enable_annotate not allowed after declarations");
7506f32e7eSjoerg       D.Report(PragmaTok.getLocation(), ID);
7606f32e7eSjoerg     }
7706f32e7eSjoerg 
7806f32e7eSjoerg     EnableAnnotate = true;
7906f32e7eSjoerg   }
8006f32e7eSjoerg };
8106f32e7eSjoerg 
8206f32e7eSjoerg }
8306f32e7eSjoerg 
8406f32e7eSjoerg static FrontendPluginRegistry::Add<AnnotateFunctionsAction>
8506f32e7eSjoerg X("annotate-fns", "annotate functions");
8606f32e7eSjoerg 
8706f32e7eSjoerg static PragmaHandlerRegistry::Add<PragmaAnnotateHandler>
8806f32e7eSjoerg Y("enable_annotate","enable annotation");
89