Make it possible for AST plugins to enable themselves by default

Currently when an AST plugin is loaded it must then be enabled by passing
-plugin pluginname or -add-plugin pluginname to the -cc1 command line. This
patch adds a method to PluginASTAction which allows it to declare that the
action happens before, instead of, or after the main AST action, plus the
relevant changes to make the plugin action happen at that time automatically.

Differential Revision: http://reviews.llvm.org/D17959


git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@263546 91177308-0d34-0410-b5e6-96231b3b80d8

docs/ClangPlugins.rst

@@ -54,6 +54,10 @@ the `latest version of PrintFunctionNames.cpp
 Running the plugin
 ==================
 
+
+Using the cc1 command line
+--------------------------
+
 To run a plugin, the dynamic library containing the plugin registry must be
 loaded via the :option:`-load` command line option. This will load all plugins
 that are registered, and you can select the plugins to run by specifying the
@@ -88,3 +92,19 @@ source tree:
 Also see the print-function-name plugin example's
 `README <http://llvm.org/viewvc/llvm-project/cfe/trunk/examples/PrintFunctionNames/README.txt?view=markup>`_
 
+
+Using the clang command line
+----------------------------
+
+Using :option:`-fplugin=plugin` on the clang command line passes the plugin
+through as an argument to :option:`-load` on the cc1 command line. If the plugin
+class implements the ``getActionType`` method then the plugin is run
+automatically. For example, to run the plugin automatically after the main AST
+action (i.e. the same as using :option:`-add-plugin`):
+
+.. code-block:: c++
+
+  // Automatically run the plugin after the main AST action
+  PluginASTAction::ActionType getActionType() override {
+    return AddAfterMainAction;
+  }

examples/AnnotateFunctions/AnnotateFunctions.cpp

@@ -0,0 +1,52 @@
+//===- AnnotateFunctions.cpp ----------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Example clang plugin which adds an annotation to every function.
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/Frontend/FrontendPluginRegistry.h"
+#include "clang/AST/AST.h"
+#include "clang/AST/ASTConsumer.h"
+using namespace clang;
+
+namespace {
+
+class AnnotateFunctionsConsumer : public ASTConsumer {
+public:
+  bool HandleTopLevelDecl(DeclGroupRef DG) override {
+    for (auto D : DG)
+      if (FunctionDecl *FD = dyn_cast<FunctionDecl>(D))
+        FD->addAttr(AnnotateAttr::CreateImplicit(FD->getASTContext(),
+                                                 "example_annotation"));
+    return true;
+  }
+};
+
+class AnnotateFunctionsAction : public PluginASTAction {
+public:
+  std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
+                                                 llvm::StringRef) override {
+    return llvm::make_unique<AnnotateFunctionsConsumer>();
+  }
+
+  bool ParseArgs(const CompilerInstance &CI,
+                 const std::vector<std::string> &args) override {
+    return true;
+  }
+
+  PluginASTAction::ActionType getActionType() override {
+    return AddBeforeMainAction;
+  }
+};
+
+}
+
+static FrontendPluginRegistry::Add<AnnotateFunctionsAction>
+X("annotate-fns", "annotate functions");

examples/AnnotateFunctions/CMakeLists.txt

@@ -0,0 +1,9 @@
+add_llvm_loadable_module(AnnotateFunctions AnnotateFunctions.cpp)
+
+if(LLVM_ENABLE_PLUGINS AND (WIN32 OR CYGWIN))
+  target_link_libraries(AnnotateFunctions ${cmake_2_8_12_PRIVATE}
+    clangAST
+    clangFrontend
+    LLVMSupport
+    )
+endif()

examples/CMakeLists.txt

@@ -8,3 +8,4 @@ add_subdirectory(analyzer-plugin)
 endif()
 add_subdirectory(clang-interpreter)
 add_subdirectory(PrintFunctionNames)
+add_subdirectory(AnnotateFunctions)

include/clang/Frontend/FrontendAction.h

@@ -249,6 +249,19 @@ class PluginASTAction : public ASTFrontendAction {
   /// CompilerInstance's Diagnostic object to report errors.
   virtual bool ParseArgs(const CompilerInstance &CI,
                          const std::vector<std::string> &arg) = 0;
+
+  enum ActionType {
+    Cmdline,             //< Action is determined by the cc1 command-line
+    ReplaceAction,       //< Replace the main action
+    AddBeforeMainAction, //< Execute the action before the main action
+    AddAfterMainAction   //< Execute the action after the main action
+  };
+  /// \brief Get the action type for this plugin
+  ///
+  /// \return The action type. If the type is Cmdline then by default the
+  /// plugin does nothing and what it does is determined by the cc1
+  /// command-line.
+  virtual ActionType getActionType() { return Cmdline; }
 };
 
 /// \brief Abstract base class to use for preprocessor-based frontend actions.

include/clang/Frontend/FrontendOptions.h

@@ -16,6 +16,7 @@
 #include "llvm/ADT/StringRef.h"
 #include <string>
 #include <vector>
+#include <unordered_map>
 
 namespace llvm {
 class MemoryBuffer;
@@ -227,15 +228,12 @@ class FrontendOptions {
   /// The name of the action to run when using a plugin action.
   std::string ActionName;
 
-  /// Args to pass to the plugin
-  std::vector<std::string> PluginArgs;
+  /// Args to pass to the plugins
+  std::unordered_map<std::string,std::vector<std::string>> PluginArgs;
 
   /// The list of plugin actions to run in addition to the normal action.
   std::vector<std::string> AddPluginActions;
 
-  /// Args to pass to the additional plugins
-  std::vector<std::vector<std::string> > AddPluginArgs;
-
   /// The list of plugins to load.
   std::vector<std::string> Plugins;
 

lib/Frontend/CompilerInvocation.cpp

@@ -1051,18 +1051,10 @@ static InputKind ParseFrontendArgs(FrontendOptions &Opts, ArgList &Args,
     Opts.Plugins.emplace_back(A->getValue(0));
     Opts.ProgramAction = frontend::PluginAction;
     Opts.ActionName = A->getValue();
-
-    for (const Arg *AA : Args.filtered(OPT_plugin_arg))
-      if (AA->getValue(0) == Opts.ActionName)
-        Opts.PluginArgs.emplace_back(AA->getValue(1));
   }
-
   Opts.AddPluginActions = Args.getAllArgValues(OPT_add_plugin);
-  Opts.AddPluginArgs.resize(Opts.AddPluginActions.size());
-  for (int i = 0, e = Opts.AddPluginActions.size(); i != e; ++i)
-    for (const Arg *A : Args.filtered(OPT_plugin_arg))
-      if (A->getValue(0) == Opts.AddPluginActions[i])
-        Opts.AddPluginArgs[i].emplace_back(A->getValue(1));
+  for (const Arg *AA : Args.filtered(OPT_plugin_arg))
+    Opts.PluginArgs[AA->getValue(0)].emplace_back(AA->getValue(1));
 
   for (const std::string &Arg :
          Args.getAllArgValues(OPT_ftest_module_file_extension_EQ)) {

lib/Frontend/FrontendAction.cpp

@@ -141,28 +141,46 @@ FrontendAction::CreateWrappedASTConsumer(CompilerInstance &CI,
   if (!Consumer)
     return nullptr;
 
-  if (CI.getFrontendOpts().AddPluginActions.size() == 0)
+  // If there are no registered plugins we don't need to wrap the consumer
+  if (FrontendPluginRegistry::begin() == FrontendPluginRegistry::end())
     return Consumer;
 
-  // Make sure the non-plugin consumer is first, so that plugins can't
-  // modifiy the AST.
+  // Collect the list of plugins that go before the main action (in Consumers)
+  // or after it (in AfterConsumers)
   std::vector<std::unique_ptr<ASTConsumer>> Consumers;
-  Consumers.push_back(std::move(Consumer));
-
-  for (size_t i = 0, e = CI.getFrontendOpts().AddPluginActions.size();
-       i != e; ++i) { 
-    // This is O(|plugins| * |add_plugins|), but since both numbers are
-    // way below 50 in practice, that's ok.
-    for (FrontendPluginRegistry::iterator
-        it = FrontendPluginRegistry::begin(),
-        ie = FrontendPluginRegistry::end();
-        it != ie; ++it) {
-      if (it->getName() != CI.getFrontendOpts().AddPluginActions[i])
-        continue;
-      std::unique_ptr<PluginASTAction> P = it->instantiate();
-      if (P->ParseArgs(CI, CI.getFrontendOpts().AddPluginArgs[i]))
-        Consumers.push_back(P->CreateASTConsumer(CI, InFile));
+  std::vector<std::unique_ptr<ASTConsumer>> AfterConsumers;
+  for (FrontendPluginRegistry::iterator it = FrontendPluginRegistry::begin(),
+                                        ie = FrontendPluginRegistry::end();
+       it != ie; ++it) {
+    std::unique_ptr<PluginASTAction> P = it->instantiate();
+    PluginASTAction::ActionType ActionType = P->getActionType();
+    if (ActionType == PluginASTAction::Cmdline) {
+      // This is O(|plugins| * |add_plugins|), but since both numbers are
+      // way below 50 in practice, that's ok.
+      for (size_t i = 0, e = CI.getFrontendOpts().AddPluginActions.size();
+           i != e; ++i) {
+        if (it->getName() == CI.getFrontendOpts().AddPluginActions[i]) {
+          ActionType = PluginASTAction::AddAfterMainAction;
+          break;
+        }
+      }
     }
+    if ((ActionType == PluginASTAction::AddBeforeMainAction ||
+         ActionType == PluginASTAction::AddAfterMainAction) &&
+        P->ParseArgs(CI, CI.getFrontendOpts().PluginArgs[it->getName()])) {
+      std::unique_ptr<ASTConsumer> PluginConsumer = P->CreateASTConsumer(CI, InFile);
+      if (ActionType == PluginASTAction::AddBeforeMainAction) {
+        Consumers.push_back(std::move(PluginConsumer));
+      } else {
+        AfterConsumers.push_back(std::move(PluginConsumer));
+      }
+    }
+  }
+
+  // Add to Consumers the main consumer, then all the plugins that go after it
+  Consumers.push_back(std::move(Consumer));
+  for (auto &C : AfterConsumers) {
+    Consumers.push_back(std::move(C));
   }
 
   return llvm::make_unique<MultiplexConsumer>(std::move(Consumers));

lib/FrontendTool/ExecuteCompilerInvocation.cpp

@@ -66,7 +66,9 @@ CreateFrontendBaseAction(CompilerInstance &CI) {
          it != ie; ++it) {
       if (it->getName() == CI.getFrontendOpts().ActionName) {
         std::unique_ptr<PluginASTAction> P(it->instantiate());
-        if (!P->ParseArgs(CI, CI.getFrontendOpts().PluginArgs))
+        if ((P->getActionType() != PluginASTAction::ReplaceAction &&
+             P->getActionType() != PluginASTAction::Cmdline) ||
+            !P->ParseArgs(CI, CI.getFrontendOpts().PluginArgs[it->getName()]))
           return nullptr;
         return std::move(P);
       }
@@ -194,6 +196,18 @@ bool clang::ExecuteCompilerInvocation(CompilerInstance *Clang) {
         << Path << Error;
   }
 
+  // Check if any of the loaded plugins replaces the main AST action
+  for (FrontendPluginRegistry::iterator it = FrontendPluginRegistry::begin(),
+                                        ie = FrontendPluginRegistry::end();
+       it != ie; ++it) {
+    std::unique_ptr<PluginASTAction> P(it->instantiate());
+    if (P->getActionType() == PluginASTAction::ReplaceAction) {
+      Clang->getFrontendOpts().ProgramAction = clang::frontend::PluginAction;
+      Clang->getFrontendOpts().ActionName = it->getName();
+      break;
+    }
+  }
+
   // Honor -mllvm.
   //
   // FIXME: Remove this, one day.

test/Frontend/plugin-annotate-functions.c

@@ -0,0 +1,7 @@
+// RUN: %clang -fplugin=%llvmshlibdir/AnnotateFunctions%pluginext -emit-llvm -S %s -o - | FileCheck %s
+// REQUIRES: plugins, examples
+
+// CHECK: [[STR_VAR:@.+]] = private unnamed_addr constant [19 x i8] c"example_annotation\00"
+// CHECK: @llvm.global.annotations = {{.*}}@fn1{{.*}}[[STR_VAR]]{{.*}}@fn2{{.*}}[[STR_VAR]]
+void fn1() { }
+void fn2() { }