clang/examples/AnnotateFunctions/AnnotateFunctions.cpp

   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");